From 6d1a2761ea056c3da6ce2e748000c9baff3c3c30 Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Mon, 24 Nov 2025 11:35:42 +0530 Subject: [PATCH 01/11] Honor explicit budget reset datetime for user --- litellm/proxy/_types.py | 1 + .../internal_user_endpoints.py | 31 ++++++++++++++----- .../key_management_endpoints.py | 24 ++++++++++++-- 3 files changed, 46 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d6b2d664f9f..b4f64a76fd1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -786,6 +786,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 66085b69b3d..681087451bc 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -306,6 +306,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"). + - 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) @@ -763,12 +764,26 @@ def _update_internal_user_params( if data.user_role == LitellmUserRoles.INTERNAL_USER: is_internal_user = True + if "budget_reset_at" in non_default_values and non_default_values["budget_reset_at"] is not None: + _reset_at = non_default_values["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: + _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 ( @@ -785,9 +800,11 @@ def _update_internal_user_params( ) 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 diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a8def13869f..74ed92ce1f8 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2015,6 +2015,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 @@ -2076,10 +2077,27 @@ 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: + 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) From 9d234a959ab47dd999c93bbacaecc928f8733a41 Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Tue, 25 Nov 2025 11:41:55 +0530 Subject: [PATCH 02/11] Added validator for rooting out past dates in budget_reset_at --- litellm/proxy/_types.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b4f64a76fd1..30f24dcde7a 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 @@ -803,6 +803,23 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): aliases: Optional[dict] = {} object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + @model_validator(mode="after") + def validate_budget_reset_at(self) -> "GenerateRequestBase": + """Validate that budget_reset_at is not in the past""" + if self.budget_reset_at is not None: + # Ensure timezone-aware for comparison + reset_at = self.budget_reset_at + 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"budget_reset_at cannot be in the past. " + f"Provided: {reset_at.isoformat()}, Current time: {current_time.isoformat()}" + ) + return self + class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): index_name: str From 1b865d893a8ecb238ca6783960ea21a58fbe29a0 Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Tue, 2 Dec 2025 16:07:50 +0530 Subject: [PATCH 03/11] Added test cases --- .../test_budget_reset_at.py | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py b/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py new file mode 100644 index 00000000000..3f1e4f4cb7e --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py @@ -0,0 +1,102 @@ +""" +Test cases for 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 +from litellm.proxy.management_endpoints.internal_user_endpoints import new_user + + +class TestBudgetResetAtValidation: + """Test Pydantic validation for budget_reset_at field""" + + def test_past_date_rejection_user_request(self): + """Test that NewUserRequest rejects past dates for 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", + budget_reset_at=past_date, + ) + + assert "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", + budget_reset_at=future_date, + ) + + assert user_request.budget_reset_at == future_date + + +class TestUserCreationWithBudgetResetAt: + """Integration tests for creating users with budget_reset_at""" + + @pytest.mark.asyncio + async def test_create_user_explicit_budget_reset_at_takes_precedence(self): + """Test that explicit budget_reset_at is honored and takes precedence over duration-based computation""" + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=MagicMock(user_id="new_user", spend=0.0) + ) + mock_prisma_client.db.litellm_usertable.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_budgettable.upsert = AsyncMock( + return_value=MagicMock() + ) + mock_prisma_client.insert_data = AsyncMock() + + explicit_date = datetime.now(timezone.utc) + timedelta(days=30) + + user_request = NewUserRequest( + user_id="new_user", + budget_duration="10m", + budget_reset_at=explicit_date, + max_budget=100.0, + ) + + with patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ): + result = await new_user( + data=user_request, + user_api_key_dict=MagicMock(user_id="admin", user_role="proxy_admin"), + ) + + assert result is not None + + # verify that budget_reset_at was set and matches explicit date (not computed from duration) + budget_upsert_calls = ( + mock_prisma_client.db.litellm_budgettable.upsert.call_args_list + ) + if budget_upsert_calls: + for call in budget_upsert_calls: + data_dict = call[1].get("data", {}) + update_data = data_dict.get("update") or data_dict.get("create") + if update_data and "budget_reset_at" in update_data: + stored_date = update_data["budget_reset_at"] + assert isinstance(stored_date, datetime) + assert abs((stored_date - explicit_date).total_seconds()) < 1 # verify dates are effectively equal + break + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 001717bc07aaedff4229b2b0c83eff8dd520ca0b Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Wed, 3 Dec 2025 10:41:59 +0530 Subject: [PATCH 04/11] Fix insert mock --- .../test_budget_reset_at.py | 35 ++++++++----------- 1 file changed, 15 insertions(+), 20 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py b/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py index 3f1e4f4cb7e..a63f00d8c99 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py @@ -48,20 +48,16 @@ class TestBudgetResetAtValidation: class TestUserCreationWithBudgetResetAt: - """Integration tests for creating users with budget_reset_at""" + """Integration test for creating users with budget_reset_at""" @pytest.mark.asyncio async def test_create_user_explicit_budget_reset_at_takes_precedence(self): """Test that explicit budget_reset_at is honored and takes precedence over duration-based computation""" mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( - return_value=MagicMock(user_id="new_user", spend=0.0) + 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) - mock_prisma_client.db.litellm_budgettable.upsert = AsyncMock( - return_value=MagicMock() - ) - mock_prisma_client.insert_data = AsyncMock() explicit_date = datetime.now(timezone.utc) + timedelta(days=30) @@ -82,20 +78,19 @@ class TestUserCreationWithBudgetResetAt: ) assert result is not None + assert mock_prisma_client.insert_data.called - # verify that budget_reset_at was set and matches explicit date (not computed from duration) - budget_upsert_calls = ( - mock_prisma_client.db.litellm_budgettable.upsert.call_args_list - ) - if budget_upsert_calls: - for call in budget_upsert_calls: - data_dict = call[1].get("data", {}) - update_data = data_dict.get("update") or data_dict.get("create") - if update_data and "budget_reset_at" in update_data: - stored_date = update_data["budget_reset_at"] - assert isinstance(stored_date, datetime) - assert abs((stored_date - explicit_date).total_seconds()) < 1 # verify dates are effectively equal - break + # 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") if __name__ == "__main__": From 09a13621c6b31ac573e45a687f327552b0c3658e Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Mon, 2 Mar 2026 14:20:00 +0530 Subject: [PATCH 05/11] Renamed param budget_reset_at to initial_budget_reset_at --- litellm/proxy/_types.py | 13 ++++++------ .../internal_user_endpoints.py | 16 ++++++++++++--- .../key_management_endpoints.py | 3 +++ ..._at.py => test_initial_budget_reset_at.py} | 20 +++++++++---------- 4 files changed, 33 insertions(+), 19 deletions(-) rename tests/test_litellm/proxy/management_endpoints/{test_budget_reset_at.py => test_initial_budget_reset_at.py} (78%) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 30f24dcde7a..d05a2ce0b43 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -786,7 +786,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None + initial_budget_reset_at: Optional[datetime] = None allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} @@ -804,18 +804,18 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @model_validator(mode="after") - def validate_budget_reset_at(self) -> "GenerateRequestBase": - """Validate that budget_reset_at is not in the past""" - if self.budget_reset_at is not None: + def validate_initial_budget_reset_at(self) -> "GenerateRequestBase": + """Validate that initial_budget_reset_at is not in the past""" + if self.initial_budget_reset_at is not None: # Ensure timezone-aware for comparison - reset_at = self.budget_reset_at + reset_at = self.initial_budget_reset_at 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"budget_reset_at cannot be in the past. " + f"initial_budget_reset_at cannot be in the past. " f"Provided: {reset_at.isoformat()}, Current time: {current_time.isoformat()}" ) return self @@ -874,6 +874,7 @@ class GenerateKeyResponse(KeyRequestBase): key: str # type: ignore key_name: Optional[str] = None expires: Optional[datetime] = None + budget_reset_at: Optional[datetime] = None # computed/stored value from DB user_id: Optional[str] = None token_id: Optional[str] = None organization_id: Optional[str] = None diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 681087451bc..1026b967a8f 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -98,6 +98,10 @@ 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: + data_json["budget_reset_at"] = data_json.pop("initial_budget_reset_at") + return data_json @@ -306,7 +310,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"). - - 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. + - 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) @@ -423,6 +427,9 @@ async def new_user( response_dict["key"] = response.get("token", "") + if data.initial_budget_reset_at is not None: + response_dict["initial_budget_reset_at"] = data.initial_budget_reset_at + new_user_response = NewUserResponse(**response_dict) ######################################################### @@ -764,8 +771,8 @@ def _update_internal_user_params( if data.user_role == LitellmUserRoles.INTERNAL_USER: is_internal_user = True - if "budget_reset_at" in non_default_values and non_default_values["budget_reset_at"] is not None: - _reset_at = non_default_values["budget_reset_at"] + 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): @@ -776,6 +783,8 @@ def _update_internal_user_params( except Exception: _reset_at = None non_default_values["budget_reset_at"] = _reset_at + elif "initial_budget_reset_at" in non_default_values: + non_default_values.pop("initial_budget_reset_at") if "budget_duration" in non_default_values: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time @@ -1003,6 +1012,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 74ed92ce1f8..d0715eae959 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -582,6 +582,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 if "budget_duration" in data_json: data_json["key_budget_duration"] = data_json.pop("budget_duration", None) + if "initial_budget_reset_at" in data_json: + data_json["budget_reset_at"] = data_json.pop("initial_budget_reset_at", None) + if user_api_key_dict.user_id is not None: data_json["created_by"] = user_api_key_dict.user_id data_json["updated_by"] = user_api_key_dict.user_id diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py b/tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py similarity index 78% rename from tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py rename to tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py index a63f00d8c99..40875b13736 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py +++ b/tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py @@ -1,5 +1,5 @@ """ -Test cases for budget_reset_at feature in user management endpoints. +Test cases for initial_budget_reset_at feature in user management endpoints. """ import os @@ -19,20 +19,20 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import new_user class TestBudgetResetAtValidation: - """Test Pydantic validation for budget_reset_at field""" + """Test Pydantic validation for initial_budget_reset_at field""" def test_past_date_rejection_user_request(self): - """Test that NewUserRequest rejects past dates for budget_reset_at""" + """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", - budget_reset_at=past_date, + initial_budget_reset_at=past_date, ) - assert "budget_reset_at cannot be in the past" in str(exc_info.value) + 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""" @@ -41,18 +41,18 @@ class TestBudgetResetAtValidation: user_request = NewUserRequest( user_id="test_user", budget_duration="10m", - budget_reset_at=future_date, + initial_budget_reset_at=future_date, ) - assert user_request.budget_reset_at == future_date + assert user_request.initial_budget_reset_at == future_date class TestUserCreationWithBudgetResetAt: - """Integration test for creating users with budget_reset_at""" + """Integration test for creating users with initial_budget_reset_at""" @pytest.mark.asyncio async def test_create_user_explicit_budget_reset_at_takes_precedence(self): - """Test that explicit budget_reset_at is honored and takes precedence over duration-based computation""" + """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=[]) @@ -64,7 +64,7 @@ class TestUserCreationWithBudgetResetAt: user_request = NewUserRequest( user_id="new_user", budget_duration="10m", - budget_reset_at=explicit_date, + initial_budget_reset_at=explicit_date, max_budget=100.0, ) From 7d7c47c676bb0e27acb45d61c419c0b811ad0dc7 Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Mon, 2 Mar 2026 14:37:00 +0530 Subject: [PATCH 06/11] Remove budget_reset_at previously added in response --- litellm/proxy/_types.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d05a2ce0b43..920bb25453c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -874,7 +874,6 @@ class GenerateKeyResponse(KeyRequestBase): key: str # type: ignore key_name: Optional[str] = None expires: Optional[datetime] = None - budget_reset_at: Optional[datetime] = None # computed/stored value from DB user_id: Optional[str] = None token_id: Optional[str] = None organization_id: Optional[str] = None From 7d893e78f59589170fc18c796e9b74f0b24fc4f8 Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Thu, 5 Mar 2026 17:16:35 +0530 Subject: [PATCH 07/11] Incorporate Greptile reviews --- litellm/proxy/_types.py | 46 +++++++++++-------- .../internal_user_endpoints.py | 5 +- .../key_management_endpoints.py | 8 ++-- 3 files changed, 36 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 920bb25453c..a4e11398558 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -786,7 +786,6 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None - initial_budget_reset_at: Optional[datetime] = None allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} @@ -803,23 +802,6 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): aliases: Optional[dict] = {} object_permission: Optional[LiteLLM_ObjectPermissionBase] = None - @model_validator(mode="after") - def validate_initial_budget_reset_at(self) -> "GenerateRequestBase": - """Validate that initial_budget_reset_at is not in the past""" - if self.initial_budget_reset_at is not None: - # Ensure timezone-aware for comparison - reset_at = self.initial_budget_reset_at - 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()}" - ) - return self - class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): index_name: str @@ -1128,10 +1110,26 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase): user_role: Literal["user", "admin"] = "user" +def _validate_initial_budget_reset_at(value: Optional[datetime]) -> Optional[datetime]: + """Shared validator: ensure initial_budget_reset_at is not in the past.""" + 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()}" + ) + return value + + 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, @@ -1148,6 +1146,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": + _validate_initial_budget_reset_at(self.initial_budget_reset_at) + return self + class NewUserResponse(GenerateKeyResponse): max_budget: Optional[float] = None @@ -1165,6 +1168,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( @@ -1183,6 +1187,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": + _validate_initial_budget_reset_at(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 1026b967a8f..196605689fc 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -780,7 +780,10 @@ def _update_internal_user_params( # ensure timezone aware if _reset_at.tzinfo is None: _reset_at = _reset_at.replace(tzinfo=timezone.utc) - except Exception: + 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 elif "initial_budget_reset_at" in non_default_values: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d0715eae959..a0aabfb3121 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -582,9 +582,6 @@ async def _common_key_generation_helper( # noqa: PLR0915 if "budget_duration" in data_json: data_json["key_budget_duration"] = data_json.pop("budget_duration", None) - if "initial_budget_reset_at" in data_json: - data_json["budget_reset_at"] = data_json.pop("initial_budget_reset_at", None) - if user_api_key_dict.user_id is not None: data_json["created_by"] = user_api_key_dict.user_id data_json["updated_by"] = user_api_key_dict.user_id @@ -2091,7 +2088,10 @@ async def generate_key_helper_fn( # noqa: PLR0915 if _reset_candidate.tzinfo is None: _reset_candidate = _reset_candidate.replace(tzinfo=timezone.utc) reset_at = _reset_candidate - except Exception: + 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: From 4ad722453292d3b0344af2aa0579e0e03d24554e Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Thu, 5 Mar 2026 17:29:39 +0530 Subject: [PATCH 08/11] Incorporate Greptile review- remove dead elif block --- litellm/proxy/management_endpoints/internal_user_endpoints.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 196605689fc..75c2110b589 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -786,8 +786,6 @@ def _update_internal_user_params( ) _reset_at = None non_default_values["budget_reset_at"] = _reset_at - elif "initial_budget_reset_at" in non_default_values: - non_default_values.pop("initial_budget_reset_at") if "budget_duration" in non_default_values: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time From 28bf548ea87646998b44aa39a98479d556092cf1 Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Fri, 6 Mar 2026 10:33:49 +0530 Subject: [PATCH 09/11] Incorporate Greptile review- remove unused return, mock license check in test --- litellm/proxy/_types.py | 9 ++++----- .../management_endpoints/test_initial_budget_reset_at.py | 6 ++++++ 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a4e11398558..0b0b4e66ef3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1110,8 +1110,8 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase): user_role: Literal["user", "admin"] = "user" -def _validate_initial_budget_reset_at(value: Optional[datetime]) -> Optional[datetime]: - """Shared validator: ensure initial_budget_reset_at is not in the past.""" +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: @@ -1122,7 +1122,6 @@ def _validate_initial_budget_reset_at(value: Optional[datetime]) -> Optional[dat f"initial_budget_reset_at cannot be in the past. " f"Provided: {reset_at.isoformat()}, Current time: {current_time.isoformat()}" ) - return value class NewUserRequest(GenerateRequestBase): @@ -1148,7 +1147,7 @@ class NewUserRequest(GenerateRequestBase): @model_validator(mode="after") def validate_initial_budget_reset_at(self) -> "NewUserRequest": - _validate_initial_budget_reset_at(self.initial_budget_reset_at) + _assert_initial_budget_reset_at_valid(self.initial_budget_reset_at) return self @@ -1191,7 +1190,7 @@ class UpdateUserRequestNoUserIDorEmail( @model_validator(mode="after") def validate_initial_budget_reset_at(self) -> "UpdateUserRequestNoUserIDorEmail": - _validate_initial_budget_reset_at(self.initial_budget_reset_at) + _assert_initial_budget_reset_at_valid(self.initial_budget_reset_at) return self 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 index 40875b13736..0dbe7f68924 100644 --- 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 @@ -68,9 +68,15 @@ class TestUserCreationWithBudgetResetAt: 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, From 1a635ff79efc0acdc2ee5f7d81ccd7828eb9e4a6 Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Fri, 6 Mar 2026 12:08:18 +0530 Subject: [PATCH 10/11] Incorporate Greptile review- None guard before mapping initial_budget_reset_at, echo stored value in response --- .../proxy/management_endpoints/internal_user_endpoints.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 75c2110b589..7d8d28740ec 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -99,8 +99,10 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d data_json.pop("teams", None) # handled separately - if "initial_budget_reset_at" in data_json: + 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 @@ -427,8 +429,8 @@ async def new_user( response_dict["key"] = response.get("token", "") - if data.initial_budget_reset_at is not None: - response_dict["initial_budget_reset_at"] = data.initial_budget_reset_at + 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) From ff421f6a2c827416e1c055ab56766c3d3437fe5a Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Fri, 6 Mar 2026 12:40:02 +0530 Subject: [PATCH 11/11] Incorporate Greptile review- improve tests --- .../test_initial_budget_reset_at.py | 95 ++++++++++++++++++- 1 file changed, 90 insertions(+), 5 deletions(-) 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 index 0dbe7f68924..e1becc7d0f0 100644 --- 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 @@ -14,11 +14,11 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) -from litellm.proxy._types import NewUserRequest -from litellm.proxy.management_endpoints.internal_user_endpoints import new_user +from litellm.proxy._types import NewUserRequest, UpdateUserRequest, UpdateUserRequestNoUserIDorEmail +from litellm.proxy.management_endpoints.internal_user_endpoints import new_user, user_update -class TestBudgetResetAtValidation: +class TestInitialBudgetResetAtValidation: """Test Pydantic validation for initial_budget_reset_at field""" def test_past_date_rejection_user_request(self): @@ -46,12 +46,49 @@ class TestBudgetResetAtValidation: 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) -class TestUserCreationWithBudgetResetAt: + 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_budget_reset_at_takes_precedence(self): + 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( @@ -99,5 +136,53 @@ class TestUserCreationWithBudgetResetAt: 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"])