diff --git a/litellm/proxy/common_utils/timezone_utils.py b/litellm/proxy/common_utils/timezone_utils.py index a50daf40144..a6c6021efe7 100644 --- a/litellm/proxy/common_utils/timezone_utils.py +++ b/litellm/proxy/common_utils/timezone_utils.py @@ -1,5 +1,7 @@ from datetime import datetime, time, timezone +from typing import Optional +from fastapi import HTTPException from pydantic import BaseModel, ConfigDict import litellm @@ -78,3 +80,34 @@ def get_budget_reset_time(budget_duration: str) -> datetime: `BudgetResetSettings` by injection (creation/update endpoints, startup backfill). """ return compute_budget_reset_at(budget_duration, get_budget_reset_settings()) + + +def validate_budget_duration(budget_duration: Optional[str]) -> None: + """Reject budget durations that can't be parsed, are non-positive, or overflow + date math, so a bad value can't be persisted and later silently reset on the + wrong cadence (or crash the budget reset job). + + Shared by every management endpoint that accepts a `budget_duration` (key, team, + customer, org, budget). `get_next_standardized_reset_time` fails open (warns and + falls back to a next-midnight reset) so a bad row never crashes the reset job; + this is the fail-closed counterpart at the write boundary that stops the bad row + from being written in the first place. + """ + if budget_duration is None: + return + + from litellm.litellm_core_utils.duration_parser import duration_in_seconds + + try: + if duration_in_seconds(budget_duration) <= 0: + raise ValueError("budget_duration must be positive") + get_budget_reset_time(budget_duration=budget_duration) + except (ValueError, OverflowError): + raise HTTPException( + status_code=400, + detail={ + "error": "Invalid budget_duration '{}'. Use a format like '1h', '24h', '7d', or '30d'.".format( + budget_duration + ) + }, + ) diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 6b70a9064df..def1f6eda2e 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -18,7 +18,10 @@ from fastapi import APIRouter, Depends, HTTPException from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.common_utils.timezone_utils import ( + get_budget_reset_time, + validate_budget_duration, +) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.utils import jsonify_object from litellm.repositories.budget_repository import BudgetRepository @@ -71,6 +74,8 @@ async def new_budget( detail={"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"}, ) + validate_budget_duration(budget_obj.budget_duration) + # Validate model_max_budget if present if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0: from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -152,6 +157,8 @@ async def update_budget( detail={"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"}, ) + validate_budget_duration(budget_obj.budget_duration) + # Validate model_max_budget if present in update if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0: from litellm.proxy.management_endpoints.key_management_endpoints import ( diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 84f67bdc3bc..8fafa4be0e5 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -20,6 +20,7 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy.common_utils.timezone_utils import validate_budget_duration from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity @@ -326,6 +327,7 @@ async def new_end_user( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) + validate_budget_duration(data.budget_duration) try: ## VALIDATION ## if data.default_model is not None: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ac6a2a4a7db..2c771b7149d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -61,7 +61,10 @@ from litellm.proxy.common_utils.callback_utils import ( encrypt_callback_vars, ) from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.common_utils.timezone_utils import ( + get_budget_reset_time, + validate_budget_duration, +) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import ( @@ -1585,6 +1588,8 @@ async def generate_key_fn( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) + validate_budget_duration(data.budget_duration) + if user_custom_key_generate is not None: if inspect.iscoroutinefunction(user_custom_key_generate): result = await user_custom_key_generate(data) # type: ignore @@ -2266,6 +2271,8 @@ async def _validate_update_key_data( # Reject NaN/±inf spend before it can reach the DB / spend counter. validate_finite_spend(data.spend) + validate_budget_duration(data.budget_duration) + _is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value _check_allowed_routes_caller_permission( diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 138a55d9227..d693ae23041 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -23,6 +23,7 @@ from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import can_user_call_model, get_user_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.timezone_utils import validate_budget_duration from litellm.proxy.management_endpoints.budget_management_endpoints import ( new_budget, update_budget, @@ -232,6 +233,8 @@ async def new_organization( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) + validate_budget_duration(data.budget_duration) + user_object_correct_type: Optional[LiteLLM_UserTable] = None if user_api_key_dict.user_id is not None: @@ -465,6 +468,8 @@ async def update_organization( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) + validate_budget_duration(data.budget_duration) + if data.updated_by is None: data.updated_by = user_api_key_dict.user_id diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 59b0cbc4ae7..7cb01174b43 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2845,28 +2845,9 @@ def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> Dict[str, Any]: def _validate_budget_duration(budget_duration: Optional[str]) -> None: - """Reject budget durations that can't be parsed, are non-positive, or - overflow date math, so a bad value can't be persisted and later crash the - budget reset job.""" - if budget_duration is None: - return + from litellm.proxy.common_utils.timezone_utils import validate_budget_duration - from litellm.litellm_core_utils.duration_parser import duration_in_seconds - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - - try: - if duration_in_seconds(budget_duration) <= 0: - raise ValueError("budget_duration must be positive") - get_budget_reset_time(budget_duration=budget_duration) - except (ValueError, OverflowError): - raise HTTPException( - status_code=400, - detail={ - "error": "Invalid budget_duration '{}'. Use a format like '1h', '24h', '7d', or '30d'.".format( - budget_duration - ) - }, - ) + validate_budget_duration(budget_duration) @router.post( diff --git a/tests/test_litellm/proxy/common_utils/test_timezone_utils.py b/tests/test_litellm/proxy/common_utils/test_timezone_utils.py index 7f686c53c95..980f1b56caa 100644 --- a/tests/test_litellm/proxy/common_utils/test_timezone_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_timezone_utils.py @@ -4,12 +4,12 @@ from datetime import datetime, time, timezone from zoneinfo import ZoneInfo import pytest +from fastapi import HTTPException -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path import litellm +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy.common_utils.timezone_utils import ( BudgetResetSettings, compute_budget_reset_at, @@ -17,6 +17,7 @@ from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_time, get_budget_reset_timezone, parse_budget_reset_time, + validate_budget_duration, ) @@ -158,9 +159,7 @@ def test_get_budget_reset_settings_reads_globals(): def test_compute_budget_reset_at_applies_offset(): - settings = BudgetResetSettings( - timezone="Asia/Jerusalem", reset_time_of_day=time(12, 0) - ) + settings = BudgetResetSettings(timezone="Asia/Jerusalem", reset_time_of_day=time(12, 0)) reset_at = compute_budget_reset_at("1d", settings) jerusalem = reset_at.astimezone(ZoneInfo("Asia/Jerusalem")) assert jerusalem.hour == 12 @@ -180,3 +179,40 @@ def test_get_budget_reset_time_honors_global_budget_reset_time(): finally: _restore_attr(litellm, "timezone", orig_tz) _restore_attr(litellm, "budget_reset_time", orig_rt) + + +class TestValidateBudgetDuration: + """`validate_budget_duration` is the fail-closed write-boundary guard shared by + the key/team/customer/org/budget endpoints. It must reject any value the budget + reset job can't honor, so a bad duration can never be persisted and then silently + reset on the wrong cadence. + """ + + def test_none_is_a_noop(self): + assert validate_budget_duration(None) is None + + @pytest.mark.parametrize("duration", ["1s", "30m", "1h", "24h", "7d", "30d", "1mo"]) + def test_canonical_durations_pass(self, duration): + assert validate_budget_duration(duration) is None + + @pytest.mark.parametrize("duration", ["hourly", "daily", "weekly", "monthly", " MONTHLY ", "Weekly"]) + def test_word_form_durations_pass(self, duration): + assert validate_budget_duration(duration) is None + + @pytest.mark.parametrize("duration", ["garbage", "5x", "abc", "", "d30", "1 day"]) + def test_unparseable_durations_raise_400(self, duration): + with pytest.raises(HTTPException) as exc_info: + validate_budget_duration(duration) + assert exc_info.value.status_code == 400 + assert "budget_duration" in exc_info.value.detail["error"] + + @pytest.mark.parametrize("duration", ["0s", "0m", "0h", "0d", "0w"]) + def test_non_positive_durations_raise_400(self, duration): + with pytest.raises(HTTPException) as exc_info: + validate_budget_duration(duration) + assert exc_info.value.status_code == 400 + + @pytest.mark.parametrize("duration", ["1s", "30m", "1h", "24h", "7d", "30d", "1mo", "hourly", "weekly", "monthly"]) + def test_accepted_iff_reset_job_can_compute_it(self, duration): + validate_budget_duration(duration) + assert duration_in_seconds(duration) > 0