diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 441b3b836a1..7e7439a1030 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,6 +1,6 @@ import enum import json -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union import httpx @@ -1406,10 +1406,25 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase): user_role: Literal["user", "admin"] = "user" +def _assert_initial_budget_reset_at_valid(value: Optional[datetime]) -> None: + """Assert that initial_budget_reset_at is not in the past. Raises ValueError if invalid.""" + if value is not None: + reset_at = value + if reset_at.tzinfo is None: + reset_at = reset_at.replace(tzinfo=timezone.utc) + current_time = datetime.now(timezone.utc) + if reset_at < current_time: + raise ValueError( + f"initial_budget_reset_at cannot be in the past. " + f"Provided: {reset_at.isoformat()}, Current time: {current_time.isoformat()}" + ) + + class NewUserRequest(GenerateRequestBase): max_budget: Optional[float] = None user_email: Optional[str] = None user_alias: Optional[str] = None + initial_budget_reset_at: Optional[datetime] = None user_role: Optional[ Literal[ LitellmUserRoles.PROXY_ADMIN, @@ -1426,6 +1441,11 @@ class NewUserRequest(GenerateRequestBase): sso_user_id: Optional[str] = None organizations: Optional[List[str]] = None + @model_validator(mode="after") + def validate_initial_budget_reset_at(self) -> "NewUserRequest": + _assert_initial_budget_reset_at_valid(self.initial_budget_reset_at) + return self + class NewUserResponse(GenerateKeyResponse): max_budget: Optional[float] = None @@ -1443,6 +1463,7 @@ class NewUserResponse(GenerateKeyResponse): model_max_budget: Optional[dict] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None + initial_budget_reset_at: Optional[datetime] = None class UpdateUserRequestNoUserIDorEmail( @@ -1461,6 +1482,12 @@ class UpdateUserRequestNoUserIDorEmail( ] ] = None max_budget: Optional[float] = None + initial_budget_reset_at: Optional[datetime] = None + + @model_validator(mode="after") + def validate_initial_budget_reset_at(self) -> "UpdateUserRequestNoUserIDorEmail": + _assert_initial_budget_reset_at_valid(self.initial_budget_reset_at) + return self class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail): diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 646779e6f81..d3c8ff10e20 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -119,6 +119,12 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d data_json["budget_duration"] = litellm.internal_user_budget_duration data_json.pop("teams", None) # handled separately + + if "initial_budget_reset_at" in data_json and data_json.get("initial_budget_reset_at") is not None: + data_json["budget_reset_at"] = data_json.pop("initial_budget_reset_at") + else: + data_json.pop("initial_budget_reset_at", None) + return data_json @@ -369,6 +375,7 @@ async def new_user( - user_role: Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20` - max_budget: Optional[float] - Specify max budget for a given user. - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo"). + - initial_budget_reset_at: Optional[datetime] - Specify the exact datetime when the budget should first reset (ISO 8601 format, e.g., "2025-11-15T00:00:00Z"). If provided along with budget_duration, this value takes precedence for the initial reset; subsequent resets will follow the budget_duration interval. - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models). Set to ['no-default-models'] to block all model access. Restricting user to only team-based model access. - tpm_limit: Optional[int] - Specify tpm limit for a given user (Tokens per minute) - rpm_limit: Optional[int] - Specify rpm limit for a given user (Requests per minute) @@ -503,6 +510,9 @@ async def new_user( response_dict["key"] = response.get("token", "") + if data.initial_budget_reset_at is not None and response.get("budget_reset_at") is not None: + response_dict["initial_budget_reset_at"] = response.get("budget_reset_at") + new_user_response = NewUserResponse(**response_dict) ######################################################### @@ -1086,12 +1096,29 @@ def _update_internal_user_params( if data.user_role == LitellmUserRoles.INTERNAL_USER: is_internal_user = True + if "initial_budget_reset_at" in non_default_values and non_default_values["initial_budget_reset_at"] is not None: + _reset_at = non_default_values.pop("initial_budget_reset_at") + try: + # convert to datetime if not already + if not isinstance(_reset_at, datetime): + _reset_at = datetime.fromisoformat(str(_reset_at).replace("Z", "+00:00")) + # ensure timezone aware + if _reset_at.tzinfo is None: + _reset_at = _reset_at.replace(tzinfo=timezone.utc) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to parse initial_budget_reset_at value: {_reset_at}. Error: {e}. Falling back to None." + ) + _reset_at = None + non_default_values["budget_reset_at"] = _reset_at + if "budget_duration" in non_default_values: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - - non_default_values["budget_reset_at"] = get_budget_reset_time( - budget_duration=non_default_values["budget_duration"] - ) + # only auto-compute budget_reset_at if not explicitly provided + if "budget_reset_at" not in non_default_values or non_default_values["budget_reset_at"] is None: + non_default_values["budget_reset_at"] = get_budget_reset_time( + budget_duration=non_default_values["budget_duration"] + ) if "max_budget" not in non_default_values: if ( @@ -1108,9 +1135,11 @@ def _update_internal_user_params( ] = litellm.internal_user_budget_duration from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - non_default_values["budget_reset_at"] = get_budget_reset_time( - budget_duration=non_default_values["budget_duration"] - ) + # only auto-compute budget_reset_at if not explicitly provided + if "budget_reset_at" not in non_default_values or non_default_values["budget_reset_at"] is None: + non_default_values["budget_reset_at"] = get_budget_reset_time( + budget_duration=non_default_values["budget_duration"] + ) return non_default_values @@ -1312,6 +1341,7 @@ async def user_update( - user_role: Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20` - max_budget: Optional[float] - Specify max budget for a given user. - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo"). + - initial_budget_reset_at: Optional[datetime] - Specify the exact datetime when the budget should first reset (ISO 8601 format, e.g., "2025-11-15T00:00:00Z"). If provided along with budget_duration, this value takes precedence for the initial reset; subsequent resets will follow the budget_duration interval. - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models) - tpm_limit: Optional[int] - Specify tpm limit for a given user (Tokens per minute) - rpm_limit: Optional[int] - Specify rpm limit for a given user (Requests per minute) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 323ff7fd531..aca26892f49 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2801,6 +2801,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 max_budget: Optional[float] = None, # max_budget is used to Budget Per user blocked: Optional[bool] = None, budget_duration: Optional[str] = None, # max_budget is used to Budget Per user + budget_reset_at: Optional[datetime] = None, # manual override for initial user/key budget reset time token: Optional[str] = None, key: Optional[ str @@ -2869,10 +2870,30 @@ async def generate_key_helper_fn( # noqa: PLR0915 else: key_reset_at = get_budget_reset_time(budget_duration=key_budget_duration) - if budget_duration is None: # one-time budget - reset_at = None + if budget_reset_at is not None: + try: + # try to parse budget_reset_at + if isinstance(budget_reset_at, datetime): + _reset_candidate = budget_reset_at + else: + _reset_candidate = datetime.fromisoformat(str(budget_reset_at).replace("Z", "+00:00")) + # ensure tz-aware datetime + if _reset_candidate.tzinfo is None: + _reset_candidate = _reset_candidate.replace(tzinfo=timezone.utc) + reset_at = _reset_candidate + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to parse budget_reset_at value: {budget_reset_at}. Error: {e}. Falling back to budget_duration-based computation." + ) + if budget_duration is None: # fallback to one-time budget + reset_at = None + else: + reset_at = get_budget_reset_time(budget_duration=budget_duration) # calculate based on budget duration else: - reset_at = get_budget_reset_time(budget_duration=budget_duration) + if budget_duration is None: # one-time budget + reset_at = None + else: + reset_at = get_budget_reset_time(budget_duration=budget_duration) # calculate based on budget duration aliases_json = json.dumps(aliases) config_json = json.dumps(config) diff --git a/tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py b/tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py new file mode 100644 index 00000000000..e1becc7d0f0 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py @@ -0,0 +1,188 @@ +""" +Test cases for initial_budget_reset_at feature in user management endpoints. +""" + +import os +import sys +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +pytest_plugins = ('pytest_asyncio',) +sys.path.insert( + 0, os.path.abspath("../../../..") +) + +from litellm.proxy._types import NewUserRequest, UpdateUserRequest, UpdateUserRequestNoUserIDorEmail +from litellm.proxy.management_endpoints.internal_user_endpoints import new_user, user_update + + +class TestInitialBudgetResetAtValidation: + """Test Pydantic validation for initial_budget_reset_at field""" + + def test_past_date_rejection_user_request(self): + """Test that NewUserRequest rejects past dates for initial_budget_reset_at""" + past_date = datetime.now(timezone.utc) - timedelta(days=1) + + with pytest.raises(ValueError) as exc_info: + NewUserRequest( + user_id="test_user", + budget_duration="10m", + initial_budget_reset_at=past_date, + ) + + assert "initial_budget_reset_at cannot be in the past" in str(exc_info.value) + + def test_future_date_accepted(self): + """Test that future dates are accepted""" + future_date = datetime.now(timezone.utc) + timedelta(days=10) + + user_request = NewUserRequest( + user_id="test_user", + budget_duration="10m", + initial_budget_reset_at=future_date, + ) + + assert user_request.initial_budget_reset_at == future_date + + def test_past_date_rejection_update_user_request(self): + """Test that UpdateUserRequest rejects past dates for initial_budget_reset_at""" + past_date = datetime.now(timezone.utc) - timedelta(days=1) + + with pytest.raises(ValueError) as exc_info: + UpdateUserRequest( + user_id="test_user", + budget_duration="10m", + initial_budget_reset_at=past_date, + ) + + assert "initial_budget_reset_at cannot be in the past" in str(exc_info.value) + + def test_past_date_rejection_update_user_request_no_id_or_email(self): + """Test that UpdateUserRequestNoUserIDorEmail rejects past dates for initial_budget_reset_at""" + past_date = datetime.now(timezone.utc) - timedelta(days=1) + + with pytest.raises(ValueError) as exc_info: + UpdateUserRequestNoUserIDorEmail( + budget_duration="10m", + initial_budget_reset_at=past_date, + ) + + assert "initial_budget_reset_at cannot be in the past" in str(exc_info.value) + + def test_future_date_accepted_update_user_request(self): + """Test that future dates are accepted on UpdateUserRequest""" + future_date = datetime.now(timezone.utc) + timedelta(days=10) + + update_request = UpdateUserRequest( + user_id="test_user", + budget_duration="10m", + initial_budget_reset_at=future_date, + ) + + assert update_request.initial_budget_reset_at == future_date + + +class TestUserCreationWithInitialBudgetResetAt: + """Integration test for creating users with initial_budget_reset_at""" + + @pytest.mark.asyncio + async def test_create_user_explicit_initial_budget_reset_at_takes_precedence(self): + """Test that explicit initial_budget_reset_at is honored and takes precedence over duration-based computation""" + mock_prisma_client = MagicMock() + mock_prisma_client.insert_data = AsyncMock( + return_value=MagicMock(user_id="new_user", spend=0.0, models=[]) + ) + mock_prisma_client.db.litellm_usertable.count = AsyncMock(return_value=0) + + explicit_date = datetime.now(timezone.utc) + timedelta(days=30) + + user_request = NewUserRequest( + user_id="new_user", + budget_duration="10m", + initial_budget_reset_at=explicit_date, + max_budget=100.0, + ) + + mock_license_check = MagicMock() + mock_license_check.is_over_limit.return_value = False + + with patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ), patch( + "litellm.proxy.proxy_server._license_check", + mock_license_check, + ): + result = await new_user( + data=user_request, + user_api_key_dict=MagicMock(user_id="admin", user_role="proxy_admin"), + ) + + assert result is not None + assert mock_prisma_client.insert_data.called + + # find the user insert call and verify budget_reset_at + for call in mock_prisma_client.insert_data.call_args_list: + if call.kwargs.get("table_name") == "user": + user_data = call.kwargs["data"] + assert "budget_reset_at" in user_data + stored_date = user_data["budget_reset_at"] + assert isinstance(stored_date, datetime) + assert abs((stored_date - explicit_date).total_seconds()) < 1 # verify dates are effectively equal + break + else: + raise AssertionError("User insert call not found") + + +class TestUserUpdateWithInitialBudgetResetAt: + """Integration test for updating users with initial_budget_reset_at""" + + @pytest.mark.asyncio + async def test_update_user_explicit_initial_budget_reset_at_translated(self): + """Test that initial_budget_reset_at is translated to budget_reset_at on the update path""" + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.update_data = AsyncMock( + return_value={"user_id": "existing_user"} + ) + + explicit_date = datetime.now(timezone.utc) + timedelta(days=30) + + update_request = UpdateUserRequest( + user_id="existing_user", + budget_duration="10m", + initial_budget_reset_at=explicit_date, + max_budget=100.0, + ) + + with patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ), patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", + "admin", + ): + result = await user_update( + data=update_request, + user_api_key_dict=MagicMock( + user_id="existing_user", + user_role="proxy_admin", + ), + ) + + assert result is not None + assert mock_prisma_client.update_data.called + + call_kwargs = mock_prisma_client.update_data.call_args.kwargs + user_data = call_kwargs["data"] + assert "budget_reset_at" in user_data, "budget_reset_at should be set from initial_budget_reset_at" + assert "initial_budget_reset_at" not in user_data, "initial_budget_reset_at should be removed" + stored_date = user_data["budget_reset_at"] + assert isinstance(stored_date, datetime) + assert abs((stored_date - explicit_date).total_seconds()) < 1 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])