From ee165498220a7b3c1c7f75aeadf8303cceb991ed Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 23:17:48 +0000 Subject: [PATCH] feat(proxy): enforce recurring access schedules on virtual keys --- litellm/proxy/_types.py | 5 + litellm/proxy/auth/access_schedule.py | 157 ++++++++++++++++ litellm/proxy/auth/user_api_key_auth.py | 13 ++ .../key_management_endpoints.py | 18 +- .../proxy/auth/test_access_schedule.py | 171 ++++++++++++++++++ .../proxy/auth/test_user_api_key_auth.py | 123 +++++++++++++ .../test_key_management_endpoints.py | 45 +++++ 7 files changed, 530 insertions(+), 2 deletions(-) create mode 100644 litellm/proxy/auth/access_schedule.py create mode 100644 tests/test_litellm/proxy/auth/test_access_schedule.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7575091be54..602761279aa 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3539,6 +3539,11 @@ class ProxyErrorTypes(str, enum.Enum): Key has expired """ + key_access_schedule_denied = "key_access_schedule_denied" + """ + Request was made outside the key's allowed recurring access schedule + """ + auth_error = "auth_error" """ General authentication error diff --git a/litellm/proxy/auth/access_schedule.py b/litellm/proxy/auth/access_schedule.py new file mode 100644 index 00000000000..a351fec07cf --- /dev/null +++ b/litellm/proxy/auth/access_schedule.py @@ -0,0 +1,157 @@ +"""Recurring access-schedule evaluation for virtual keys. + +A key may carry ``permissions.access_schedule`` describing recurring time +windows (per-weekday, in a named IANA timezone) during which the key is +allowed to make requests. Requests outside every window are denied. Invalid +persisted schedules fail closed (deny). +""" + +from __future__ import annotations + +from datetime import datetime, time +from typing import Literal, Mapping, Union +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError + +from pydantic import BaseModel, ConfigDict, ValidationError, field_validator, model_validator + +ACCESS_SCHEDULE_PERMISSION_KEY = "access_schedule" + +Weekday = Literal["mon", "tue", "wed", "thu", "fri", "sat", "sun"] + +_WEEKDAY_INDEX: Mapping[Weekday, int] = { + "mon": 0, + "tue": 1, + "wed": 2, + "thu": 3, + "fri": 4, + "sat": 5, + "sun": 6, +} + + +class AccessWindow(BaseModel): + model_config = ConfigDict(extra="forbid") + + days: tuple[Weekday, ...] + start: time + end: time + + @field_validator("days") + @classmethod + def _days_non_empty(cls, value: tuple[Weekday, ...]) -> tuple[Weekday, ...]: + if len(value) == 0: + raise ValueError("'days' must contain at least one weekday") + return value + + @field_validator("start", "end") + @classmethod + def _time_is_naive(cls, value: time) -> time: + if value.tzinfo is not None: + raise ValueError("'start'/'end' must be a local wall-clock time without an offset") + return value + + @model_validator(mode="after") + def _start_differs_from_end(self) -> AccessWindow: + if self.start == self.end: + raise ValueError("'start' and 'end' must differ (use two windows for a full day)") + return self + + +class AccessSchedule(BaseModel): + model_config = ConfigDict(extra="forbid") + + timezone: str + windows: tuple[AccessWindow, ...] + + @field_validator("timezone") + @classmethod + def _timezone_is_known(cls, value: str) -> str: + try: + ZoneInfo(value) + except (ZoneInfoNotFoundError, ValueError) as exc: + raise ValueError(f"unknown timezone '{value}'") from exc + return value + + @field_validator("windows") + @classmethod + def _windows_non_empty(cls, value: tuple[AccessWindow, ...]) -> tuple[AccessWindow, ...]: + if len(value) == 0: + raise ValueError("'windows' must contain at least one window") + return value + + +class ScheduleAbsent(BaseModel): + tag: Literal["absent"] = "absent" + + +class ScheduleValid(BaseModel): + tag: Literal["valid"] = "valid" + schedule: AccessSchedule + + +class ScheduleInvalid(BaseModel): + tag: Literal["invalid"] = "invalid" + error: str + + +ScheduleParseResult = Union[ScheduleAbsent, ScheduleValid, ScheduleInvalid] + + +def _format_validation_error(exc: ValidationError) -> str: + first = exc.errors()[0] + location = ".".join(str(part) for part in first["loc"]) + prefix = f"{location}: " if location else "" + return f"{prefix}{first['msg']}" + + +def parse_access_schedule(permissions: Mapping[str, object] | None) -> ScheduleParseResult: + if not permissions or ACCESS_SCHEDULE_PERMISSION_KEY not in permissions: + return ScheduleAbsent() + try: + schedule = AccessSchedule.model_validate(permissions[ACCESS_SCHEDULE_PERMISSION_KEY]) + except ValidationError as exc: + return ScheduleInvalid(error=_format_validation_error(exc)) + return ScheduleValid(schedule=schedule) + + +def _window_is_open(window: AccessWindow, weekday_index: int, wall_clock: time) -> bool: + day_indexes = frozenset(_WEEKDAY_INDEX[day] for day in window.days) + if window.start < window.end: + return weekday_index in day_indexes and window.start <= wall_clock < window.end + starts_today = weekday_index in day_indexes and wall_clock >= window.start + spills_from_yesterday = ((weekday_index - 1) % 7) in day_indexes and wall_clock < window.end + return starts_today or spills_from_yesterday + + +def is_within_schedule(schedule: AccessSchedule, now: datetime) -> bool: + local_now = now.astimezone(ZoneInfo(schedule.timezone)) + weekday_index = local_now.weekday() + wall_clock = local_now.time() + return any(_window_is_open(window, weekday_index, wall_clock) for window in schedule.windows) + + +class AccessAllowed(BaseModel): + tag: Literal["allowed"] = "allowed" + + +class AccessDenied(BaseModel): + tag: Literal["denied"] = "denied" + reason: str + + +AccessDecision = Union[AccessAllowed, AccessDenied] + + +def evaluate_access_schedule(permissions: Mapping[str, object] | None, now: datetime) -> AccessDecision: + parsed = parse_access_schedule(permissions) + match parsed: + case ScheduleAbsent(): + return AccessAllowed() + case ScheduleInvalid(error=error): + return AccessDenied(reason=f"key has an invalid access_schedule and is denied (fail closed): {error}") + case ScheduleValid(schedule=schedule): + if is_within_schedule(schedule, now): + return AccessAllowed() + return AccessDenied( + reason=f"request is outside the key's allowed access_schedule (timezone {schedule.timezone})" + ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e7c12448435..2494b45db83 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -55,6 +55,7 @@ from litellm.proxy.auth.auth_checks import ( is_valid_fallback_model, resolve_and_validate_end_user_id, ) +from litellm.proxy.auth.access_schedule import AccessDenied, evaluate_access_schedule from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, @@ -1842,6 +1843,18 @@ async def _user_api_key_auth_builder( param=abbreviate_api_key(api_key=api_key), ) + access_decision = evaluate_access_schedule( + permissions=valid_token.permissions, + now=datetime.now(timezone.utc), + ) + if isinstance(access_decision, AccessDenied): + raise ProxyException( + message=f"Authentication Error - {access_decision.reason}", + type=ProxyErrorTypes.key_access_schedule_denied, + code=status.HTTP_403_FORBIDDEN, + param=abbreviate_api_key(api_key=api_key), + ) + if not skip_budget_checks: with tracer.trace("litellm.proxy.auth.budget_checks"): # Check 4. Max Budget Alert Check (runs before budget enforcement diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ac6a2a4a7db..52d06ffdb64 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -54,6 +54,7 @@ from litellm.proxy.auth.auth_checks import ( get_project_object, get_team_object, ) +from litellm.proxy.auth.access_schedule import ScheduleInvalid, parse_access_schedule from litellm.proxy.auth.auth_utils import abbreviate_api_key from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import ( @@ -599,6 +600,17 @@ def _check_permissions_caller_permission( ) +def _validate_access_schedule(data: GenerateRequestBase) -> None: + if "permissions" not in data.model_fields_set and not data.permissions: + return + parsed = parse_access_schedule(data.permissions) + if isinstance(parsed, ScheduleInvalid): + raise HTTPException( + status_code=400, + detail={"error": f"Invalid permissions.access_schedule: {parsed.error}"}, + ) + + def _check_budget_limits_delegation_ceiling( budget_limits: Optional[List[BudgetLimitEntry]], delegation_ceiling: Optional[float], @@ -1506,7 +1518,7 @@ async def generate_key_fn( - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. - - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} + - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} It may also carry an `access_schedule` object ({"timezone": "Europe/Berlin", "windows": [{"days": ["mon","tue"], "start": "09:00", "end": "18:00"}]}) restricting the key to recurring time windows; requests outside every window are rejected with 403 - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. @@ -1604,6 +1616,7 @@ async def generate_key_fn( data=data, user_api_key_dict=user_api_key_dict, ) + _validate_access_schedule(data) # For non-admin internal users: auto-assign caller's user_id if not provided # This prevents creating unbound keys with no user association (LIT-1884) @@ -1714,7 +1727,7 @@ async def generate_service_account_key_fn( - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - guardrails: Optional[List[str]] - List of active guardrails for the key - - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} + - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} It may also carry an `access_schedule` object ({"timezone": "Europe/Berlin", "windows": [{"days": ["mon","tue"], "start": "09:00", "end": "18:00"}]}) restricting the key to recurring time windows; requests outside every window are rejected with 403 - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. @@ -2281,6 +2294,7 @@ async def _validate_update_key_data( data=data, user_api_key_dict=user_api_key_dict, ) + _validate_access_schedule(data) _validate_caller_can_change_key_ownership( data=data, diff --git a/tests/test_litellm/proxy/auth/test_access_schedule.py b/tests/test_litellm/proxy/auth/test_access_schedule.py new file mode 100644 index 00000000000..5d713a7b460 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_access_schedule.py @@ -0,0 +1,171 @@ +"""Unit tests for recurring access-schedule evaluation on virtual keys.""" + +from datetime import datetime, timezone + +import pytest + +from litellm.proxy.auth.access_schedule import ( + AccessAllowed, + AccessDenied, + ScheduleAbsent, + ScheduleInvalid, + ScheduleValid, + evaluate_access_schedule, + is_within_schedule, + parse_access_schedule, +) + +WORKDAY_SCHEDULE = { + "access_schedule": { + "timezone": "Europe/Berlin", + "windows": [ + {"days": ["mon", "tue", "wed", "thu", "fri"], "start": "09:00", "end": "18:00"} + ], + } +} + + +def _utc(year: int, month: int, day: int, hour: int, minute: int = 0) -> datetime: + return datetime(year, month, day, hour, minute, tzinfo=timezone.utc) + + +def test_absent_when_no_permissions(): + assert isinstance(parse_access_schedule(None), ScheduleAbsent) + assert isinstance(parse_access_schedule({}), ScheduleAbsent) + assert isinstance(parse_access_schedule({"pii": False}), ScheduleAbsent) + + +def test_valid_schedule_parses(): + parsed = parse_access_schedule(WORKDAY_SCHEDULE) + assert isinstance(parsed, ScheduleValid) + assert parsed.schedule.timezone == "Europe/Berlin" + assert len(parsed.schedule.windows) == 1 + + +@pytest.mark.parametrize( + "raw", + [ + {"timezone": "Europe/Berlin"}, + {"timezone": "Europe/Berlin", "windows": []}, + {"timezone": "Not/AZone", "windows": [{"days": ["mon"], "start": "09:00", "end": "18:00"}]}, + {"timezone": "UTC", "windows": [{"days": [], "start": "09:00", "end": "18:00"}]}, + {"timezone": "UTC", "windows": [{"days": ["funday"], "start": "09:00", "end": "18:00"}]}, + {"timezone": "UTC", "windows": [{"days": ["mon"], "start": "9am", "end": "18:00"}]}, + {"timezone": "UTC", "windows": [{"days": ["mon"], "start": "09:00", "end": "18:00", "x": 1}]}, + {"timezone": "UTC", "windows": [{"days": ["mon"], "start": "09:00", "end": "09:00"}]}, + [{"days": ["mon"], "start": "09:00", "end": "18:00"}], + ], +) +def test_invalid_schedules_are_rejected(raw): + assert isinstance(parse_access_schedule({"access_schedule": raw}), ScheduleInvalid) + + +def test_within_window_allows(): + now = _utc(2026, 7, 24, 10) # Fri 12:00 Berlin + assert evaluate_access_schedule(WORKDAY_SCHEDULE, now) == AccessAllowed() + + +def test_start_is_inclusive(): + now = _utc(2026, 7, 24, 7) # Fri 09:00 Berlin exactly + assert isinstance(evaluate_access_schedule(WORKDAY_SCHEDULE, now), AccessAllowed) + + +def test_end_is_exclusive(): + now = _utc(2026, 7, 24, 16) # Fri 18:00 Berlin exactly + assert isinstance(evaluate_access_schedule(WORKDAY_SCHEDULE, now), AccessDenied) + + +def test_after_hours_denies(): + now = _utc(2026, 7, 24, 18) # Fri 20:00 Berlin + decision = evaluate_access_schedule(WORKDAY_SCHEDULE, now) + assert isinstance(decision, AccessDenied) + assert "access_schedule" in decision.reason + + +def test_before_hours_denies(): + now = _utc(2026, 7, 24, 6) # Fri 08:00 Berlin + assert isinstance(evaluate_access_schedule(WORKDAY_SCHEDULE, now), AccessDenied) + + +def test_weekend_denies(): + now = _utc(2026, 7, 26, 10) # Sun 12:00 Berlin + assert isinstance(evaluate_access_schedule(WORKDAY_SCHEDULE, now), AccessDenied) + + +def test_timezone_is_respected(): + schedule = { + "access_schedule": { + "timezone": "America/New_York", + "windows": [{"days": ["fri"], "start": "09:00", "end": "18:00"}], + } + } + # Fri 20:00 UTC == Fri 16:00 New York (inside) + assert isinstance(evaluate_access_schedule(schedule, _utc(2026, 7, 24, 20)), AccessAllowed) + # Fri 08:00 UTC == Fri 04:00 New York (before window) + assert isinstance(evaluate_access_schedule(schedule, _utc(2026, 7, 24, 8)), AccessDenied) + + +OVERNIGHT_SCHEDULE = { + "access_schedule": { + "timezone": "UTC", + "windows": [{"days": ["fri"], "start": "22:00", "end": "06:00"}], + } +} + + +def test_overnight_open_on_start_day_evening(): + assert isinstance(evaluate_access_schedule(OVERNIGHT_SCHEDULE, _utc(2026, 7, 24, 23)), AccessAllowed) + + +def test_overnight_spills_into_next_morning(): + assert isinstance(evaluate_access_schedule(OVERNIGHT_SCHEDULE, _utc(2026, 7, 25, 5)), AccessAllowed) + + +def test_overnight_closed_after_end_next_morning(): + assert isinstance(evaluate_access_schedule(OVERNIGHT_SCHEDULE, _utc(2026, 7, 25, 7)), AccessDenied) + + +def test_overnight_closed_before_start_on_start_day(): + assert isinstance(evaluate_access_schedule(OVERNIGHT_SCHEDULE, _utc(2026, 7, 24, 21)), AccessDenied) + + +def test_overnight_does_not_open_on_non_start_evening(): + # Saturday 23:00 is not covered: only Friday evenings start the window + assert isinstance(evaluate_access_schedule(OVERNIGHT_SCHEDULE, _utc(2026, 7, 25, 23)), AccessDenied) + + +def test_multiple_windows_any_match_allows(): + schedule = { + "access_schedule": { + "timezone": "UTC", + "windows": [ + {"days": ["mon"], "start": "09:00", "end": "12:00"}, + {"days": ["mon"], "start": "13:00", "end": "17:00"}, + ], + } + } + # Monday 2026-07-20 + assert isinstance(evaluate_access_schedule(schedule, _utc(2026, 7, 20, 10)), AccessAllowed) + assert isinstance(evaluate_access_schedule(schedule, _utc(2026, 7, 20, 14)), AccessAllowed) + # lunch gap 12:00-13:00 is denied + assert isinstance(evaluate_access_schedule(schedule, _utc(2026, 7, 20, 12, 30)), AccessDenied) + + +def test_absent_schedule_allows(): + assert isinstance(evaluate_access_schedule({}, _utc(2026, 7, 26, 3)), AccessAllowed) + + +def test_invalid_persisted_schedule_fails_closed(): + decision = evaluate_access_schedule( + {"access_schedule": {"timezone": "Not/AZone", "windows": []}}, + _utc(2026, 7, 24, 10), + ) + assert isinstance(decision, AccessDenied) + assert "fail closed" in decision.reason + + +def test_is_within_schedule_direct(): + parsed = parse_access_schedule(OVERNIGHT_SCHEDULE) + assert isinstance(parsed, ScheduleValid) + assert is_within_schedule(parsed.schedule, _utc(2026, 7, 24, 23)) is True + assert is_within_schedule(parsed.schedule, _utc(2026, 7, 24, 21)) is False diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index ca2e282a119..d0040f949d6 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1329,6 +1329,129 @@ async def test_scim_deactivated_user_key_is_rejected(): setattr(_proxy_server_mod, attr, val) +@pytest.mark.parametrize( + "access_schedule, expect_denied", + [ + ({"timezone": "UTC", "windows": [{"days": ["__OTHER_DAY__"], "start": "00:00", "end": "23:59"}]}, True), + ({"timezone": "Not/AZone", "windows": []}, True), + ( + { + "timezone": "UTC", + "windows": [ + {"days": ["mon", "tue", "wed", "thu", "fri", "sat", "sun"], "start": "00:00", "end": "12:00"}, + {"days": ["mon", "tue", "wed", "thu", "fri", "sat", "sun"], "start": "12:00", "end": "23:59:59.999999"}, + ], + }, + False, + ), + ], +) +@pytest.mark.asyncio +async def test_virtual_key_access_schedule_enforced_in_auth(access_schedule, expect_denied): + """The auth flow must deny a virtual key whose recurring access_schedule + excludes the current time (or is invalid), and allow one that includes it. + """ + from datetime import datetime, timezone + + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + from litellm.proxy.proxy_server import hash_token + + _day_names = ["mon", "tue", "wed", "thu", "fri", "sat", "sun"] + other_day = _day_names[(datetime.now(timezone.utc).weekday() + 3) % 7] + windows = [ + {**w, "days": [other_day if d == "__OTHER_DAY__" else d for d in w["days"]]} + for w in access_schedule["windows"] + ] + permissions = {"access_schedule": {**access_schedule, "windows": windows}} + + api_key = "sk-access-schedule-key" + hashed_key = hash_token(api_key) + valid_token = UserAPIKeyAuth( + api_key=api_key, + token=hashed_key, + user_id="schedule-user", + permissions=permissions, + ) + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.delete_cache = MagicMock() + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs_to_set = { + "prisma_client": MagicMock(), + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} + try: + for attr, val in _attrs_to_set.items(): + setattr(_proxy_server_mod, attr, val) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + with ( + patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=valid_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + if expect_denied: + with pytest.raises(ProxyException) as exc_info: + await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + assert exc_info.value.type == ProxyErrorTypes.key_access_schedule_denied + assert int(exc_info.value.code) == status.HTTP_403_FORBIDDEN + assert exc_info.value.param != api_key + else: + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + assert result.token == hashed_key + finally: + for attr, val in _original_values.items(): + setattr(_proxy_server_mod, attr, val) + + @pytest.mark.asyncio async def test_cached_proxy_admin_key_sets_via_virtual_key_marker(): """Cached PROXY_ADMIN auth objects early-return before the marked DB and diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 51f72f91dc3..ea2e8bf5b44 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15067,3 +15067,48 @@ async def test_rotate_master_key_rotates_sso_identity_assertions( prisma_client=mock_prisma_client, new_master_key="sk-new-master-key", ) + + +def test_validate_access_schedule_accepts_valid_schedule(): + from litellm.proxy._types import GenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_access_schedule, + ) + + data = GenerateKeyRequest( + permissions={ + "access_schedule": { + "timezone": "Europe/Berlin", + "windows": [ + {"days": ["mon", "tue", "wed", "thu", "fri"], "start": "09:00", "end": "18:00"} + ], + } + } + ) + _validate_access_schedule(data) + + +def test_validate_access_schedule_rejects_invalid_schedule(): + from litellm.proxy._types import UpdateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_access_schedule, + ) + + data = UpdateKeyRequest( + key="sk-1", + permissions={"access_schedule": {"timezone": "Not/AZone", "windows": []}}, + ) + with pytest.raises(HTTPException) as exc_info: + _validate_access_schedule(data) + assert exc_info.value.status_code == 400 + assert "access_schedule" in str(exc_info.value.detail) + + +def test_validate_access_schedule_ignores_other_permissions(): + from litellm.proxy._types import GenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_access_schedule, + ) + + _validate_access_schedule(GenerateKeyRequest(permissions={"pii": False})) + _validate_access_schedule(GenerateKeyRequest())