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 1/4] 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() From d85205b41b77016e7d2c5d81046e13c1466a231b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 27 Oct 2026 00:10:22 +0000 Subject: [PATCH 2/4] test: align startup fixtures with owned database lifespan --- tests/unit/proxy/test_proxy_server.py | 4 ++-- .../proxy/test_proxy_server_endpoints_and_startup.py | 11 +++++++++-- 2 files changed, 11 insertions(+), 4 deletions(-) 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 a859ce96c70..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 @@ -15587,4 +15595,3 @@ async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is call(api_provider="openai", capture_rate=0.97), call(api_provider="openai", capture_rate=None), ] - From ee2cb0b6e02ceeb55d954bdae86c7c03f35bc194 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 26 Oct 2026 17:21:14 -0700 Subject: [PATCH 3/4] refactor(db): keep the database handle feature-agnostic Database now holds the backend handle and its close operation only, and each feature provider builds its own repository from it, so _db never grows a field per feature. get_database checks the concrete type instead of validating a runtime Protocol on every request, and the banner audit runs as a FastAPI background task so the in-flight drain waits for it before the database closes The proxy-dependencies skill moves to the repo root .agents/skills and gains bridge rules against passing SQL text or identifiers from Python Co-Authored-By: Claude Opus 5.5 --- .../skills/proxy-dependencies/SKILL.md | 7 +++-- .../references/route-migration.md | 5 +-- .../references/runtime-ownership.md | 7 +++-- AGENTS.md | 2 +- litellm-rust/AGENTS.md | 2 -- litellm/_db.py | 27 +++------------- litellm/proxy/deps.py | 14 ++++++--- litellm/proxy/proxy_server.py | 2 +- .../user_banner_dependencies.py | 4 +-- .../user_banner_endpoints.py | 6 ++-- .../unit/proxy/proxy_server/test_lifecycle.py | 8 ++--- tests/unit/proxy/test_deps.py | 31 ++++--------------- .../test_user_banner_endpoints.py | 2 +- tests/unit/test__db.py | 2 +- 14 files changed, 44 insertions(+), 75 deletions(-) rename {litellm-rust/.agents => .agents}/skills/proxy-dependencies/SKILL.md (82%) rename {litellm-rust/.agents => .agents}/skills/proxy-dependencies/references/route-migration.md (85%) rename {litellm-rust/.agents => .agents}/skills/proxy-dependencies/references/runtime-ownership.md (82%) diff --git a/litellm-rust/.agents/skills/proxy-dependencies/SKILL.md b/.agents/skills/proxy-dependencies/SKILL.md similarity index 82% rename from litellm-rust/.agents/skills/proxy-dependencies/SKILL.md rename to .agents/skills/proxy-dependencies/SKILL.md index 4d714adc48c..7ff6c54aa86 100644 --- a/litellm-rust/.agents/skills/proxy-dependencies/SKILL.md +++ b/.agents/skills/proxy-dependencies/SKILL.md @@ -5,7 +5,7 @@ description: Inject typed repositories and services into Python proxy handlers, # Proxy dependencies -- Follow the root [database-access skill](../../../../.agents/skills/database-access/SKILL.md) for queries, repository extraction and backend migration +- 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 @@ -19,8 +19,9 @@ description: Inject typed repositories and services into Python proxy handlers, ## 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 +- `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 diff --git a/litellm-rust/.agents/skills/proxy-dependencies/references/route-migration.md b/.agents/skills/proxy-dependencies/references/route-migration.md similarity index 85% rename from litellm-rust/.agents/skills/proxy-dependencies/references/route-migration.md rename to .agents/skills/proxy-dependencies/references/route-migration.md index d3c2aad4204..c5e2c5da878 100644 --- a/litellm-rust/.agents/skills/proxy-dependencies/references/route-migration.md +++ b/.agents/skills/proxy-dependencies/references/route-migration.md @@ -3,7 +3,7 @@ ## 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 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 @@ -17,7 +17,8 @@ - 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 +- 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 diff --git a/litellm-rust/.agents/skills/proxy-dependencies/references/runtime-ownership.md b/.agents/skills/proxy-dependencies/references/runtime-ownership.md similarity index 82% rename from litellm-rust/.agents/skills/proxy-dependencies/references/runtime-ownership.md rename to .agents/skills/proxy-dependencies/references/runtime-ownership.md index de867addc36..e2d218b5de5 100644 --- a/litellm-rust/.agents/skills/proxy-dependencies/references/runtime-ownership.md +++ b/.agents/skills/proxy-dependencies/references/runtime-ownership.md @@ -3,8 +3,8 @@ ## 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 +- 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 @@ -13,6 +13,9 @@ ## 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 diff --git a/AGENTS.md b/AGENTS.md index 69e34be00fc..24aeb7b615e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,4 +1,4 @@ -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) +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) diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index c21837858e4..e62a9463cb5 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,7 +1,5 @@ # 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 index 587cbca54a5..f78a18e80f2 100644 --- a/litellm/_db.py +++ b/litellm/_db.py @@ -1,37 +1,18 @@ 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 +from typing import TYPE_CHECKING 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) +class Database: + prisma: "PrismaClient" 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) + await self.prisma.disconnect() @asynccontextmanager diff --git a/litellm/proxy/deps.py b/litellm/proxy/deps.py index 4ee0aa15c90..a10ddfedd69 100644 --- a/litellm/proxy/deps.py +++ b/litellm/proxy/deps.py @@ -1,18 +1,22 @@ 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() +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: - return _DATABASE_ADAPTER.validate_python(getattr(request.state, "database", _MISSING_DATABASE)) + 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 d2667f5b052..c2d341a5bba 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1552,7 +1552,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener ) database: Final = await resources.enter_async_context( - _db.lifespan(None if prisma_client is None else _db.from_prisma(prisma_client)) + _db.lifespan(None if prisma_client is None else _db.Database(prisma=prisma_client)) ) await migrate_if_requested( diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py index b4016d4a226..1ad9e578e5d 100644 --- a/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py +++ b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py @@ -6,7 +6,7 @@ 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 +from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerRepository, UserBannerStore class UserBannerAudit(Protocol): @@ -14,7 +14,7 @@ class UserBannerAudit(Protocol): async def get_user_banner_repository(database: DatabaseDependency) -> UserBannerStore | None: - return None if database is None else database.user_banners + return None if database is None else UserBannerRepository(database.prisma) async def get_user_banner_audit() -> UserBannerAudit: diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py index 1ba6cfed3f0..ac2c3be81e9 100644 --- a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py @@ -1,7 +1,6 @@ -import asyncio from typing import Annotated, Final -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException from pydantic import ConfigDict from litellm._uuid import uuid4 @@ -51,6 +50,7 @@ async def update_user_banner( 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. @@ -72,6 +72,6 @@ async def update_user_banner( await repository.save_banner(banner) - asyncio.create_task(audit(before, banner, user_api_key_dict)) + background_tasks.add_task(audit, before, banner, user_api_key_dict) return UpdateUserBannerResponse(message="User banner updated successfully", banner=banner) diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index b182d3ed33f..707bd8ca237 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -118,7 +118,7 @@ async def test_database_lifespan_disconnects_prisma_after_proxy_shutdown(monkeyp monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - async with _db.lifespan(_db.from_prisma(fake_prisma)): + async with _db.lifespan(_db.Database(prisma=fake_prisma)): await proxy_shutdown_event() fake_prisma.disconnect.assert_not_awaited() @@ -225,7 +225,7 @@ 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) - async with _db.lifespan(_db.from_prisma(fake_prisma)): + async with _db.lifespan(_db.Database(prisma=fake_prisma)): await proxy_shutdown_event() assert calls == ["flush", "disconnect"] @@ -266,7 +266,7 @@ 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())) - async with _db.lifespan(_db.from_prisma(fake_prisma)): + async with _db.lifespan(_db.Database(prisma=fake_prisma)): await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) assert calls == ["flush", "audit", "disconnect"] @@ -338,7 +338,7 @@ async def test_database_lifespan_propagates_disconnect_error_after_shutdown(monk monkeypatch.setattr(litellm, "success_callback", [], raising=False) with pytest.raises(RuntimeError, match="db gone"): - async with _db.lifespan(_db.from_prisma(fake_prisma)): + async with _db.lifespan(_db.Database(prisma=fake_prisma)): await proxy_shutdown_event() diff --git a/tests/unit/proxy/test_deps.py b/tests/unit/proxy/test_deps.py index 712b206c8bb..0591064cfe2 100644 --- a/tests/unit/proxy/test_deps.py +++ b/tests/unit/proxy/test_deps.py @@ -6,15 +6,12 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from pydantic import ValidationError +from starlette.requests import Request 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, -) +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: @@ -36,7 +33,7 @@ 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)) + 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( @@ -63,7 +60,7 @@ def test_dependency_override_changes_next_request_without_changing_lifespan_owne prisma.db.litellm_uisettings.find_unique = AsyncMock( return_value=SimpleNamespace(ui_settings={"message": "stored"}) ) - app: Final = _app(_db.from_prisma(prisma)) + 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) @@ -74,22 +71,6 @@ def test_dependency_override_changes_next_request_without_changing_lifespan_owne @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"): + with pytest.raises(DatabaseStateMissingError): 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 7fdf3222869..dc3e748a5de 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 @@ -58,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.setitem(app.dependency_overrides, get_database, lambda: _db.from_prisma(mock_prisma)) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: _db.Database(prisma=mock_prisma)) return mock_prisma diff --git a/tests/unit/test__db.py b/tests/unit/test__db.py index e6c18defd97..829bff2d014 100644 --- a/tests/unit/test__db.py +++ b/tests/unit/test__db.py @@ -12,7 +12,7 @@ from litellm import _db async def test_lifespan_closes_database_when_owner_fails(failure): prisma: Final = MagicMock() prisma.disconnect = AsyncMock() - database: Final = _db.from_prisma(prisma) + database: Final = _db.Database(prisma=prisma) async def fail(): async with _db.lifespan(database) as current: From cebc7e6dd7ff52e936e1c53cc4021ab72bf1f7aa Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 26 Oct 2026 17:39:06 -0700 Subject: [PATCH 4/4] wip --- litellm-rust/Cargo.lock | 1 + ...ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json | 26 ++++ ...8dc13486a1f56d03400f54edf3f25be95cccc.json | 14 +++ .../queries/user_banner/find_user_banner.sql | 3 + .../queries/user_banner/save_user_banner.sql | 5 + litellm-rust/crates/db/src/lib.rs | 2 + litellm-rust/crates/db/src/user_banner.rs | 114 ++++++++++++++++++ litellm-rust/crates/python-bridge/Cargo.toml | 13 +- .../crates/python-bridge/src/database.rs | 68 +++++++++++ litellm-rust/crates/python-bridge/src/lib.rs | 4 + litellm/_db.py | 8 +- litellm/constants.py | 1 + litellm/proxy/proxy_server.py | 17 ++- .../user_banner_dependencies.py | 13 +- .../repositories/user_banner_repository.py | 14 +++ litellm/rust_bridge/_native.pyi | 9 ++ litellm/rust_bridge/database.py | 33 +++++ .../test_user_banner_endpoints.py | 41 +++++++ tests/unit/test__db.py | 14 +++ 19 files changed, 391 insertions(+), 9 deletions(-) create mode 100644 litellm-rust/crates/db/.sqlx/query-136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json create mode 100644 litellm-rust/crates/db/.sqlx/query-c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc.json create mode 100644 litellm-rust/crates/db/queries/user_banner/find_user_banner.sql create mode 100644 litellm-rust/crates/db/queries/user_banner/save_user_banner.sql create mode 100644 litellm-rust/crates/db/src/user_banner.rs create mode 100644 litellm-rust/crates/python-bridge/src/database.rs create mode 100644 litellm/rust_bridge/database.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 32ae431f863..33a7c7aa30d 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4229,6 +4229,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 6fc05c43146..0a7a072d171 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 index f78a18e80f2..c064cbf2fdb 100644 --- a/litellm/_db.py +++ b/litellm/_db.py @@ -5,14 +5,20 @@ 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: - await self.prisma.disconnect() + try: + if self.native is not None: + await self.native.close() + finally: + await self.prisma.disconnect() @asynccontextmanager diff --git a/litellm/constants.py b/litellm/constants.py index d1daf2ee2d1..dbd9351a4e8 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/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c2d341a5bba..9b9b28709e1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -86,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, ) @@ -214,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: @@ -1420,6 +1422,15 @@ 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: @@ -1552,7 +1563,11 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener ) database: Final = await resources.enter_async_context( - _db.lifespan(None if prisma_client is None else _db.Database(prisma=prisma_client)) + _db.lifespan( + None + if prisma_client is None + else _db.Database(prisma=prisma_client, native=await _connect_native_database()) + ) ) await migrate_if_requested( diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py index 1ad9e578e5d..895ee54cba9 100644 --- a/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py +++ b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py @@ -6,7 +6,12 @@ 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, UserBannerRepository, UserBannerStore +from litellm.repositories.user_banner_repository import ( + USER_BANNER_ROW_ID, + NativeUserBannerRepository, + UserBannerRepository, + UserBannerStore, +) class UserBannerAudit(Protocol): @@ -14,7 +19,11 @@ class UserBannerAudit(Protocol): async def get_user_banner_repository(database: DatabaseDependency) -> UserBannerStore | None: - return None if database is None else UserBannerRepository(database.prisma) + 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: diff --git a/litellm/repositories/user_banner_repository.py b/litellm/repositories/user_banner_repository.py index e8298f4aa0d..ffdaecdf14c 100644 --- a/litellm/repositories/user_banner_repository.py +++ b/litellm/repositories/user_banner_repository.py @@ -1,4 +1,5 @@ import json +from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Protocol from pydantic import ValidationError @@ -9,6 +10,8 @@ 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" @@ -28,6 +31,17 @@ def parse_user_banner(raw_settings: object) -> UserBanner: 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" 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/ui_crud_endpoints/test_user_banner_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py index dc3e748a5de..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 @@ -229,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 index 829bff2d014..35b1ec9623a 100644 --- a/tests/unit/test__db.py +++ b/tests/unit/test__db.py @@ -34,3 +34,17 @@ async def test_lifespan_without_database_preserves_owner_error(): 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()