Honor explicit budget reset datetime for user

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] = {}

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
 

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)