diff --git a/litellm/models/deployment_warning.py b/litellm/models/deployment_warning.py new file mode 100644 index 00000000000..1fcec9163bb --- /dev/null +++ b/litellm/models/deployment_warning.py @@ -0,0 +1,24 @@ +from collections.abc import Mapping +from typing import Literal, Protocol + +from pydantic import AwareDatetime, ConfigDict, Field + +from litellm.types.llms.base import LiteLLMBaseModel + + +class AdminIgnore(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + source: Literal["admin"] = "admin" + ignored_by: str + ignored_at: AwareDatetime + ignored_until: AwareDatetime | None = Field( + default=None, + description="When the snooze ends. Null means the warning is ignored until an admin restores it.", + ) + + +class DeploymentWarningIgnoreStore(Protocol): + async def get_ignores(self) -> Mapping[str, AdminIgnore]: ... + + async def save_ignores(self, ignores: Mapping[str, AdminIgnore]) -> None: ... diff --git a/litellm/proxy/deps/repositories.py b/litellm/proxy/deps/repositories.py index 3f4f0dce32d..f5308f837ae 100644 --- a/litellm/proxy/deps/repositories.py +++ b/litellm/proxy/deps/repositories.py @@ -2,8 +2,10 @@ from typing import Annotated, TypeAlias from fastapi import Depends +from litellm.models.deployment_warning import DeploymentWarningIgnoreStore from litellm.models.user_banner import UserBannerStore from litellm.proxy.deps.database import DatabaseDep +from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository from litellm.repositories.user_banner_repository import UserBannerRepository @@ -12,3 +14,12 @@ async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | UserBannerRepositoryDep: TypeAlias = Annotated[UserBannerStore | None, Depends(get_user_banner_repository)] + + +async def get_deployment_warning_ignore_repository(database: DatabaseDep) -> DeploymentWarningIgnoreStore | None: + return None if database is None else DeploymentWarningIgnoreRepository(database.prisma, use_writer=True) + + +DeploymentWarningIgnoreRepositoryDep: TypeAlias = Annotated[ + DeploymentWarningIgnoreStore | None, Depends(get_deployment_warning_ignore_repository) +] diff --git a/litellm/proxy/health_endpoints/deployment_warning_endpoints.py b/litellm/proxy/health_endpoints/deployment_warning_endpoints.py index 115068f79f8..3936eaba73f 100644 --- a/litellm/proxy/health_endpoints/deployment_warning_endpoints.py +++ b/litellm/proxy/health_endpoints/deployment_warning_endpoints.py @@ -1,4 +1,3 @@ -import json import logging from collections.abc import Mapping from dataclasses import dataclass @@ -7,12 +6,14 @@ from types import MappingProxyType from typing import Annotated, Final, Literal, Protocol, TypeAlias from fastapi import APIRouter, Depends, HTTPException, status -from pydantic import AwareDatetime, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, model_validator +from pydantic import ConfigDict, Field, JsonValue, TypeAdapter, model_validator -from litellm._logging import verbose_logger, verbose_proxy_logger +from litellm._logging import verbose_logger from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.models.deployment_warning import AdminIgnore, DeploymentWarningIgnoreStore from litellm.proxy._types import AUDIT_ACTIONS, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import DeploymentWarningIgnoreRepositoryDep from litellm.proxy.health_endpoints._health_endpoints import ( DISABLE_NO_REDIS_WARNING_ENV_VAR, env_credential_login_enabled, @@ -22,7 +23,6 @@ from litellm.proxy.health_endpoints._health_endpoints import ( ) from litellm.repositories.deployment_warning_ignore_repository import ( DEPLOYMENT_WARNING_IGNORES_ROW_ID, - DeploymentWarningIgnoreRepository, ) from litellm.types.llms.base import LiteLLMBaseModel @@ -50,18 +50,6 @@ class WarningCondition: env_var_ignore: str | None = None -class AdminIgnore(LiteLLMBaseModel): - model_config = ConfigDict(frozen=True) - - source: Literal["admin"] = "admin" - ignored_by: str - ignored_at: AwareDatetime - ignored_until: AwareDatetime | None = Field( - default=None, - description="When the snooze ends. Null means the warning is ignored until an admin restores it.", - ) - - class EnvVarIgnore(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) @@ -109,12 +97,6 @@ class DeploymentWarningIgnoreResponse(LiteLLMBaseModel): ignore: AdminIgnore | None -class DeploymentWarningIgnoreStore(Protocol): - async def get_raw_ignores(self) -> object: ... - - async def upsert_ignores(self, payload: str) -> None: ... - - class ConfigAuditLogger(Protocol): async def __call__( self, @@ -127,7 +109,6 @@ class ConfigAuditLogger(Protocol): ) -> None: ... -_IGNORES_ADAPTER: Final[TypeAdapter[Mapping[str, AdminIgnore]]] = TypeAdapter(Mapping[str, AdminIgnore]) _JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _AUDIT_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[DeploymentWarningIgnoreResponse, ...]]] = TypeAdapter( tuple[DeploymentWarningIgnoreResponse, ...] @@ -185,18 +166,8 @@ def is_known_warning_id(warning_id: str) -> bool: return False -def _parse_stored_ignores(raw: object) -> Mapping[str, AdminIgnore]: - if raw is None: - return MappingProxyType({}) - try: - return MappingProxyType(_IGNORES_ADAPTER.validate_python(json.loads(raw) if isinstance(raw, str) else raw)) - except (json.JSONDecodeError, ValidationError): - verbose_proxy_logger.warning("Ignoring unreadable %s row", DEPLOYMENT_WARNING_IGNORES_ROW_ID) - return MappingProxyType({}) - - async def _read_unexpired_ignores(store: DeploymentWarningIgnoreStore, now: datetime) -> Mapping[str, AdminIgnore]: - stored: Final = _parse_stored_ignores(await store.get_raw_ignores()) + stored: Final = await store.get_ignores() return MappingProxyType( { warning_id: ignore @@ -225,14 +196,6 @@ def utc_now() -> datetime: return datetime.now(timezone.utc) -def deployment_warning_ignore_store() -> DeploymentWarningIgnoreStore | None: - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - return None - return DeploymentWarningIgnoreRepository(prisma_client, use_writer=True) - - def config_audit_logger() -> ConfigAuditLogger: from litellm.proxy.proxy_server import create_config_audit_log @@ -281,7 +244,7 @@ async def current_warning_conditions(now: Annotated[datetime, Depends(utc_now)]) async def get_deployment_warnings( _: Annotated[UserAPIKeyAuth, Depends(_deployment_warning_reader)], conditions: Annotated[tuple[WarningCondition, ...], Depends(current_warning_conditions)], - store: Annotated[DeploymentWarningIgnoreStore | None, Depends(deployment_warning_ignore_store)], + store: DeploymentWarningIgnoreRepositoryDep, now: Annotated[datetime, Depends(utc_now)], ) -> DeploymentWarningsResponse: """ @@ -309,7 +272,7 @@ async def update_deployment_warning_ignore( warning_id: str, update: DeploymentWarningIgnoreUpdate, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_deployment_warning_admin)], - store: Annotated[DeploymentWarningIgnoreStore | None, Depends(deployment_warning_ignore_store)], + store: DeploymentWarningIgnoreRepositoryDep, audit_logger: Annotated[ConfigAuditLogger, Depends(config_audit_logger)], now: Annotated[datetime, Depends(utc_now)], ) -> DeploymentWarningIgnoreResponse: @@ -338,7 +301,7 @@ async def update_deployment_warning_ignore( others: Final = {other_id: other for other_id, other in before.items() if other_id != warning_id} after: Final = MappingProxyType(others if ignore is None else {**others, warning_id: ignore}) - await store.upsert_ignores(_IGNORES_ADAPTER.dump_json(dict(after)).decode()) + await store.save_ignores(after) await audit_logger( param_name=DEPLOYMENT_WARNING_IGNORES_ROW_ID, action="updated", diff --git a/litellm/repositories/deployment_warning_ignore_repository.py b/litellm/repositories/deployment_warning_ignore_repository.py index da02d424f80..95c9121d4e1 100644 --- a/litellm/repositories/deployment_warning_ignore_repository.py +++ b/litellm/repositories/deployment_warning_ignore_repository.py @@ -1,5 +1,12 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.models.deployment_warning import AdminIgnore from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: @@ -7,15 +14,28 @@ if TYPE_CHECKING: DEPLOYMENT_WARNING_IGNORES_ROW_ID: Final = "deployment_warning_ignores" +_IGNORES_ADAPTER: Final[TypeAdapter[Mapping[str, AdminIgnore]]] = TypeAdapter(Mapping[str, AdminIgnore]) + + +def parse_ignores(raw: object) -> Mapping[str, AdminIgnore]: + if raw is None: + return MappingProxyType({}) + try: + return MappingProxyType(_IGNORES_ADAPTER.validate_python(json.loads(raw) if isinstance(raw, str) else raw)) + except (json.JSONDecodeError, ValidationError): + verbose_proxy_logger.warning("Ignoring unreadable %s row", DEPLOYMENT_WARNING_IGNORES_ROW_ID) + return MappingProxyType({}) + class DeploymentWarningIgnoreRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]): table_name = "litellm_uisettings" - async def get_raw_ignores(self) -> object: + async def get_ignores(self) -> Mapping[str, AdminIgnore]: db_record: Final = await self.table.find_unique(where={"id": DEPLOYMENT_WARNING_IGNORES_ROW_ID}) - return db_record.ui_settings if db_record is not None else None + return parse_ignores(db_record.ui_settings if db_record is not None else None) - async def upsert_ignores(self, payload: str) -> None: + async def save_ignores(self, ignores: Mapping[str, AdminIgnore]) -> None: + payload: Final = _IGNORES_ADAPTER.dump_json(dict(ignores)).decode() row: Final = {"id": DEPLOYMENT_WARNING_IGNORES_ROW_ID, "ui_settings": payload} await self.table.upsert( where={"id": DEPLOYMENT_WARNING_IGNORES_ROW_ID}, diff --git a/tests/unit/proxy/deps/test_repositories.py b/tests/unit/proxy/deps/test_repositories.py index 92ae02fc08b..aa4af5310ff 100644 --- a/tests/unit/proxy/deps/test_repositories.py +++ b/tests/unit/proxy/deps/test_repositories.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from contextlib import asynccontextmanager from types import SimpleNamespace from typing import Final @@ -6,9 +7,10 @@ from unittest.mock import AsyncMock, MagicMock from fastapi import FastAPI from fastapi.testclient import TestClient +from litellm.models.deployment_warning import AdminIgnore from litellm.models.user_banner import UserBanner from litellm.proxy.deps.database import get_database -from litellm.proxy.deps.repositories import UserBannerRepositoryDep +from litellm.proxy.deps.repositories import DeploymentWarningIgnoreRepositoryDep, UserBannerRepositoryDep from litellm.repositories import database as database_runtime from tests._support.native_database import native_database @@ -26,6 +28,10 @@ def _app(database: database_runtime.Database | None) -> FastAPI: async def banner(banners: UserBannerRepositoryDep) -> UserBanner: return UserBanner() if banners is None else await banners.get_banner() + @app.get("/ignores") + async def ignores(ignores: DeploymentWarningIgnoreRepositoryDep) -> Mapping[str, AdminIgnore] | None: + return None if ignores is None else await ignores.get_ignores() + return app @@ -44,6 +50,22 @@ def test_lifespan_database_is_shared_until_shutdown(): ) +def test_both_repositories_use_the_lifespan_database(): + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock( + return_value=SimpleNamespace(ui_settings={"message": "stored banner"}) + ) + stored_ignores: Final = {"detailed_debug": {"ignored_by": "admin", "ignored_at": "2026-10-14T12:00:00Z"}} + prisma.writer_db.litellm_uisettings.find_unique = AsyncMock( + return_value=SimpleNamespace(ui_settings=stored_ignores) + ) + with TestClient(_app(database_runtime.Database(prisma=prisma, native=NATIVE))) as client: + assert client.get("/banner").json() == UserBanner(message="stored banner").model_dump() + assert client.get("/ignores").json() == { + "detailed_debug": {**stored_ignores["detailed_debug"], "source": "admin", "ignored_until": None} + } + + def test_no_database_state_preserves_optional_repository(): with TestClient(_app(None)) as client: response: Final = client.get("/banner") diff --git a/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py b/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py index 4583d8fd978..67aea8dfeb8 100644 --- a/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py @@ -1,10 +1,11 @@ import json import logging -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final +from unittest.mock import MagicMock import pytest from fastapi import FastAPI @@ -13,14 +14,15 @@ from pydantic import JsonValue from litellm._logging import verbose_logger from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.models.deployment_warning import AdminIgnore from litellm.proxy._types import AUDIT_ACTIONS, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import get_deployment_warning_ignore_repository from litellm.proxy.health_endpoints.deployment_warning_endpoints import ( WarningCondition, config_audit_logger, current_warning_conditions, deployment_warning_conditions, - deployment_warning_ignore_store, license_expiry_condition, router, utc_now, @@ -28,7 +30,9 @@ from litellm.proxy.health_endpoints.deployment_warning_endpoints import ( from litellm.proxy.proxy_server import ( _dump_redacted_config, # pyright: ignore[reportPrivateUsage] # the exact redaction every config audit row goes through ) +from litellm.repositories.database import Database from litellm.repositories.deployment_warning_ignore_repository import DEPLOYMENT_WARNING_IGNORES_ROW_ID +from litellm.rust_bridge._native import NativeDatabase T0: Final = datetime(2026, 10, 14, 12, 0, tzinfo=timezone.utc) THIRTY_DAYS: Final = 30 * 24 * 60 * 60 @@ -36,13 +40,13 @@ THIRTY_DAYS: Final = 30 * 24 * 60 * 60 @dataclass class InMemoryIgnoreStore: - raw: object = None + saved: list[Mapping[str, AdminIgnore]] = field(default_factory=list) - async def get_raw_ignores(self) -> object: - return self.raw + async def get_ignores(self) -> Mapping[str, AdminIgnore]: + return self.saved[-1] if self.saved else {} - async def upsert_ignores(self, payload: str) -> None: - self.raw = json.loads(payload) + async def save_ignores(self, ignores: Mapping[str, AdminIgnore]) -> None: + self.saved.append(ignores) @dataclass @@ -87,7 +91,7 @@ class Harness: self.app.dependency_overrides[current_warning_conditions] = lambda: conditions def without_database(self) -> None: - self.app.dependency_overrides[deployment_warning_ignore_store] = lambda: None + self.app.dependency_overrides[get_deployment_warning_ignore_repository] = lambda: None def warnings(self) -> dict[str, dict[str, object]]: response = self.client.get("/health/deployment_warnings") @@ -102,7 +106,7 @@ def harness() -> Iterator[Harness]: store = InMemoryIgnoreStore() audit = RecordingAuditLogger() clock = [T0] - app.dependency_overrides[deployment_warning_ignore_store] = lambda: store + app.dependency_overrides[get_deployment_warning_ignore_repository] = lambda: store app.dependency_overrides[config_audit_logger] = lambda: audit app.dependency_overrides[utc_now] = lambda: clock[0] h = Harness(app=app, client=TestClient(app), store=store, audit=audit, clock=clock) @@ -233,28 +237,29 @@ class TestCurrentWarningConditions: class TestDeploymentWarningIgnoreStore: - def test_no_database_means_no_store(self, monkeypatch): - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) - assert deployment_warning_ignore_store() is None + @pytest.mark.asyncio + async def test_no_database_means_no_store(self): + assert await get_deployment_warning_ignore_repository(None) is None @pytest.mark.asyncio - async def test_reads_ignores_from_the_writer_not_a_lagging_replica(self, monkeypatch): + async def test_reads_ignores_from_the_writer_not_a_lagging_replica(self): def ui_settings_table(row: object) -> SimpleNamespace: async def find_unique(where: dict[str, str]) -> object: return row return SimpleNamespace(litellm_uisettings=SimpleNamespace(find_unique=find_unique)) - fresh = {"detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00Z"}} - monkeypatch.setattr( - "litellm.proxy.proxy_server.prisma_client", - SimpleNamespace( - db=ui_settings_table(None), writer_db=ui_settings_table(SimpleNamespace(ui_settings=fresh)) - ), + fresh: Final = {"detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00Z"}} + prisma: Final = MagicMock( + db=ui_settings_table(None), writer_db=ui_settings_table(SimpleNamespace(ui_settings=fresh)) + ) + store: Final = await get_deployment_warning_ignore_repository( + Database(prisma=prisma, native=await NativeDatabase.connect()) ) - store = deployment_warning_ignore_store() assert store is not None - assert await store.get_raw_ignores() == fresh + assert await store.get_ignores() == { + "detailed_debug": AdminIgnore(ignored_by="a", ignored_at=datetime(2026, 10, 14, 12, tzinfo=timezone.utc)) + } class TestGetDeploymentWarnings: @@ -295,20 +300,6 @@ class TestGetDeploymentWarnings: ) assert harness.warnings()["no_redis"]["ignore"]["source"] == "env_var" - def test_stored_ignore_without_timezone_reads_as_unreadable(self, harness): - harness.store.raw = { - "detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00", "ignored_until": None} - } - warnings = harness.warnings() - assert set(warnings) == {"env_credential_login", "detailed_debug"} - assert warnings["detailed_debug"]["ignore"] is None - - def test_corrupt_stored_ignores_read_as_none(self, harness): - harness.store.raw = "{not json" - warnings = harness.warnings() - assert set(warnings) == {"env_credential_login", "detailed_debug"} - assert all(w["ignore"] is None for w in warnings.values()) - def test_without_database_ignoring_is_reported_unavailable(self, harness): harness.without_database() response = harness.client.get("/health/deployment_warnings") @@ -438,7 +429,7 @@ class TestUpdateDeploymentWarningIgnore: "/health/deployment_warnings/detailed_debug", json={"ignored": True, "snooze_seconds": snooze_seconds} ) assert response.status_code == 422 - assert harness.store.raw is None + assert harness.store.saved == [] def test_thirty_day_snooze_is_allowed(self, harness): response = harness.client.patch( @@ -457,14 +448,14 @@ class TestUpdateDeploymentWarningIgnore: def test_unknown_warning_is_rejected(self, harness, warning_id): response = harness.client.patch(f"/health/deployment_warnings/{warning_id}", json={"ignored": True}) assert response.status_code == 404 - assert harness.store.raw is None + assert harness.store.saved == [] @pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER]) def test_only_proxy_admins_can_ignore(self, harness, role): harness.as_role(role) response = harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": True}) assert response.status_code == 403 - assert harness.store.raw is None + assert harness.store.saved == [] assert harness.audit.calls == [] def test_without_database_ignoring_is_unavailable(self, harness): diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py index 4e410e22b7f..56843640888 100644 --- a/tests/unit/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -3,7 +3,7 @@ Tests for gateway repository layer. """ import json -from datetime import datetime +from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Any, Dict, Final, List, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -23,7 +23,11 @@ from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.credentials_repository import CredentialsRepository -from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository +from litellm.models.deployment_warning import AdminIgnore +from litellm.repositories.deployment_warning_ignore_repository import ( + DeploymentWarningIgnoreRepository, + parse_ignores, +) from litellm.repositories.model_repository import ModelRepository from litellm.repositories.object_permission_repository import ( ObjectPermissionRepository, @@ -1618,15 +1622,27 @@ class TestDeploymentWarningIgnoreRepository: } writer_table._records.update({row_id: dict(row) for row_id, row in other_rows.items()}) repo = DeploymentWarningIgnoreRepository(client, use_writer=True) + ignored_at = datetime(2026, 10, 14, 12, tzinfo=timezone.utc) + snoozed = AdminIgnore(ignored_by="admin", ignored_at=ignored_at, ignored_until=ignored_at + timedelta(hours=1)) - assert await repo.get_raw_ignores() is None - await repo.upsert_ignores('{"no_redis": {}}') - await repo.upsert_ignores('{"detailed_debug": {}}') + assert await repo.get_ignores() == {} + await repo.save_ignores({"no_redis": AdminIgnore(ignored_by="admin", ignored_at=ignored_at)}) + await repo.save_ignores({"detailed_debug": snoozed}) - assert await repo.get_raw_ignores() == '{"detailed_debug": {}}' + assert await repo.get_ignores() == {"detailed_debug": snoozed} assert {row_id: writer_table._records[row_id] for row_id in other_rows} == other_rows assert len(writer_table._records) == len(other_rows) + 1 + @pytest.mark.parametrize( + "stored", + [ + "{not json", + {"detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00", "ignored_until": None}}, + ], + ) + def test_unreadable_rows_read_as_no_ignores(self, stored): + assert parse_ignores(stored) == {} + class TestVerificationTokenRepositoryExtended: @pytest.fixture