From b31b29d694ae31d1f8d940794500c22f96148772 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 02:09:39 +0000 Subject: [PATCH] fix(proxy): stop treating empty budget_duration as a daily reset Empty or unparseable budget_duration values were falling through get_next_standardized_reset_time to next-midnight. Treat blank as unset and raise on garbage so spend is not wiped every 24 hours Fixes #39368 Co-authored-by: Zsanz3 --- litellm/litellm_core_utils/duration_parser.py | 20 ++------- litellm/proxy/_types.py | 21 ++++++++++ .../management_endpoints/team_endpoints.py | 24 ++++++++--- .../test_duration_parser.py | 13 +++--- .../test_team_endpoints.py | 42 +++++++++++++++++++ tests/test_litellm/proxy/test__types.py | 11 +++++ 6 files changed, 101 insertions(+), 30 deletions(-) diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 7ac2299b7d8..414369e35c6 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -12,8 +12,6 @@ from datetime import datetime, time, timedelta, timezone, tzinfo from typing import Final from zoneinfo import ZoneInfo -from litellm._logging import verbose_logger - _BUDGET_DURATION_WORD_ALIASES: Final[dict[str, str]] = { "hourly": "1h", "daily": "24h", @@ -135,20 +133,14 @@ def get_next_standardized_reset_time( # Set up timezone and normalize current time current_time, _ = _setup_timezone(current_time, timezone_str) - # Parse duration value, unit = _parse_duration(_normalize_duration(duration)) - if value is None: - verbose_logger.warning( - "Unrecognized budget_duration %r; falling back to a next-midnight reset. " - "Use the format (e.g. '1h', '7d', '30d', '1mo').", - duration, + if value is None or unit not in {"s", "m", "h", "d", "w", "mo"}: + raise ValueError( + f"Invalid budget_duration {duration!r}. Use the format (e.g. '1h', '7d', '30d', '1mo')." ) - return current_time.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) - # Midnight of the current day in the specified timezone base_midnight: Final = current_time.replace(hour=0, minute=0, second=0, microsecond=0) - # Handle different time units if unit == "d": return _handle_day_reset(current_time, base_midnight, value, reset_time_of_day) elif unit == "w": @@ -159,11 +151,7 @@ def get_next_standardized_reset_time( return _handle_minute_reset(current_time, base_midnight, value) elif unit == "s": return _handle_second_reset(current_time, base_midnight, value) - elif unit == "mo": - return _handle_month_reset(current_time, base_midnight, value, reset_time_of_day) - else: - # Unrecognized unit, default to next midnight - return base_midnight + timedelta(days=1) + return _handle_month_reset(current_time, base_midnight, value, reset_time_of_day) def _setup_timezone(current_time: datetime, timezone_str: str = "UTC") -> tuple[datetime, tzinfo]: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b33e2fe7ff6..e418f16d4f5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1121,6 +1121,12 @@ from litellm.types.object_permission import ( # noqa: E402 ) +def _blank_duration_as_unset(v: object) -> object: + if isinstance(v, str) and v.strip() == "": + return None + return v + + class GenerateRequestBase(LiteLLMPydanticObjectBase): """ Overlapping schema between key and user generate/update requests @@ -1166,6 +1172,11 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): return None return v + @field_validator("budget_duration", mode="before") + @classmethod + def empty_budget_duration_is_unset(cls, v: object) -> object: + return _blank_duration_as_unset(v) + class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): index_name: str @@ -1942,6 +1953,11 @@ class NewTeamRequest(TeamBase): return None return v + @field_validator("budget_duration", "team_member_budget_duration", mode="before") + @classmethod + def empty_budget_duration_is_unset(cls, v: object) -> object: + return _blank_duration_as_unset(v) + class GlobalEndUsersSpend(LiteLLMPydanticObjectBase): api_key: str | None = None @@ -2005,6 +2021,11 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): budget_limits: list[BudgetLimitEntry] | None = None # multiple concurrent budget windows default_team_member_models: list[str] | None = None # default allowed_models seeded onto new team members + @field_validator("budget_duration", "team_member_budget_duration", mode="before") + @classmethod + def empty_budget_duration_is_unset(cls, v: object) -> object: + return _blank_duration_as_unset(v) + class PatchTeamRequest(UpdateTeamRequest): """ diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 2ec68a10f65..44733fd3cd9 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1526,12 +1526,14 @@ async def new_team( value=getattr(data, field), ) - # If budget_duration is set, set `budget_reset_at` - if complete_team_data.budget_duration is not None: + usable_budget_duration: Final = _usable_budget_duration(complete_team_data.budget_duration) + if usable_budget_duration is None: + complete_team_data.budget_duration = None + else: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time complete_team_data.budget_reset_at = get_budget_reset_time( - budget_duration=complete_team_data.budget_duration, + budget_duration=usable_budget_duration, ) # If budget_limits is set, initialize reset_at for each window @@ -2403,14 +2405,24 @@ async def patch_team( raise handle_exception_on_proxy(e) +def _usable_budget_duration(duration: str | None) -> str | None: + if duration is None or duration.strip() == "": + return None + return duration + + def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None: """Set budget_reset_at in updated_kv if budget_duration is provided.""" - if data.budget_duration is not None: + usable_budget_duration: Final = _usable_budget_duration(data.budget_duration) + if usable_budget_duration is not None: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - reset_at: Final = get_budget_reset_time(budget_duration=data.budget_duration) + reset_at: Final = get_budget_reset_time(budget_duration=usable_budget_duration) updated_kv["budget_reset_at"] = reset_at - elif "budget_duration" in updated_kv and updated_kv["budget_duration"] is None: + elif data.budget_duration is not None or ( + "budget_duration" in updated_kv and updated_kv["budget_duration"] is None + ): + updated_kv["budget_duration"] = None updated_kv["budget_reset_at"] = None if data.budget_limits is not None and len(data.budget_limits) > 0: diff --git a/tests/test_litellm/litellm_core_utils/test_duration_parser.py b/tests/test_litellm/litellm_core_utils/test_duration_parser.py index cb9f273a0a7..94c429530ee 100644 --- a/tests/test_litellm/litellm_core_utils/test_duration_parser.py +++ b/tests/test_litellm/litellm_core_utils/test_duration_parser.py @@ -1,9 +1,7 @@ import unittest from datetime import datetime, time, timezone -from unittest.mock import patch from zoneinfo import ZoneInfo -import litellm.litellm_core_utils.duration_parser as duration_parser from litellm.litellm_core_utils.duration_parser import ( duration_in_seconds, get_next_standardized_reset_time, @@ -376,13 +374,12 @@ class TestWordFormBudgetDurations(unittest.TestCase): self.assertEqual(duration_in_seconds("weekly"), 604800) self.assertEqual(duration_in_seconds("monthly"), 2592000) - def test_invalid_duration_logs_warning_and_falls_back(self): + def test_unparseable_duration_raises_instead_of_midnight_fallback(self): base_time = datetime(2023, 5, 15, 15, 0, 0, tzinfo=timezone.utc) - with patch.object(duration_parser.verbose_logger, "warning") as mock_warning: - result = get_next_standardized_reset_time("garbage", base_time, "UTC") - self.assertEqual(result, datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc)) - mock_warning.assert_called_once() - self.assertIn("garbage", mock_warning.call_args.args) + for duration in ("", " ", "garbage", "1x"): + with self.assertRaises(ValueError) as ctx: + get_next_standardized_reset_time(duration, base_time, "UTC") + self.assertIn("Invalid budget_duration", str(ctx.exception)) if __name__ == "__main__": diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 46678c8ff6a..63e848398c0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -541,6 +541,19 @@ async def test_update_team_rejects_a_duration_that_never_advances( mock_find_unique.assert_not_awaited() +@pytest.mark.parametrize("blank", ["", " "]) +def test_set_budget_reset_at_treats_a_blank_duration_as_unset(blank): + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + data = UpdateTeamRequest.model_construct(team_id="team-1", budget_duration=blank) + updated_kv = {"budget_duration": blank} + + _set_budget_reset_at(data, updated_kv) + + assert updated_kv["budget_duration"] is None + assert updated_kv["budget_reset_at"] is None + + @pytest.mark.asyncio async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): """ @@ -13127,6 +13140,35 @@ async def test_new_team_explicit_null_budget_duration_beats_configured_default( assert team_data.get("budget_reset_at") is None +@pytest.mark.asyncio +@pytest.mark.parametrize("blank", ["", " "]) +async def test_new_team_blank_budget_duration_is_lifetime_not_a_daily_reset( + mock_db_client, mock_admin_auth, monkeypatch, blank +): + """A form-posted empty budget_duration used to stamp budget_reset_at at + next midnight. It must persist as unset, the same as an explicit null. + """ + from fastapi import Request + + import litellm + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr(litellm, "default_team_settings", None) + monkeypatch.setattr(litellm, "default_team_params", {"budget_duration": "30d"}) + mock_team_create = _wire_new_team_prisma(mock_db_client) + + await new_team( + data=NewTeamRequest(team_alias="blank-duration-team", budget_duration=blank), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + team_data = mock_team_create.call_args.kwargs["data"] + assert team_data.get("budget_duration") is None + assert team_data.get("budget_reset_at") is None + + @pytest.mark.asyncio async def test_new_team_omitted_budget_duration_still_takes_configured_default( mock_db_client, mock_admin_auth, monkeypatch diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index 26bb1533da4..c598ce427ca 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -10,9 +10,11 @@ from litellm.proxy._types import ( LiteLLM_AuditLogs, LiteLLM_TeamMembership, LitellmUserRoles, + NewTeamRequest, OrganizationMemberUpdateRequest, ResetSpendRequest, UpdateKeyRequest, + UpdateTeamRequest, UpdateUserRequest, UserAPIKeyAuth, ) @@ -188,6 +190,15 @@ def test_an_empty_max_budget_from_a_form_post_reads_as_no_budget_not_as_zero(): assert GenerateKeyRequest(max_budget="").max_budget is None +@pytest.mark.parametrize("blank", ("", " ")) +def test_an_empty_budget_duration_from_a_form_post_reads_as_unset(blank): + assert GenerateKeyRequest(budget_duration=blank).budget_duration is None + assert NewTeamRequest(team_alias="my-team", budget_duration=blank).budget_duration is None + assert NewTeamRequest(team_alias="my-team", team_member_budget_duration=blank).team_member_budget_duration is None + assert UpdateTeamRequest(team_id="team-1", budget_duration=blank).budget_duration is None + assert UpdateTeamRequest(team_id="team-1", team_member_budget_duration=blank).team_member_budget_duration is None + + @pytest.mark.parametrize("sent", (0, 0.0, 25.5)) def test_a_max_budget_that_was_actually_sent_is_kept(sent): assert GenerateKeyRequest(max_budget=sent).max_budget == sent