From 25d52df97b284324582a09d0de33d333fe2ca856 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Sat, 10 Oct 2026 17:04:24 +0200 Subject: [PATCH] fix(proxy): count a monthly budget window from its own start when reporting its reset duration_in_seconds measures Nmo from now, so a window opened on January 15 and read on February 10 reported February 12, then null while the counter still refused requests until February 15. duration_in_seconds_from measures the month from the window start, and duration_in_seconds now shares that month arithmetic --- litellm/litellm_core_utils/duration_parser.py | 54 +++++++-------- .../proxy/hooks/model_max_budget_limiter.py | 6 +- .../test_duration_parser.py | 29 +++++++++ ...test_unit_test_max_model_budget_limiter.py | 65 +++++++++++++++++++ 4 files changed, 121 insertions(+), 33 deletions(-) diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index fe8b71e7623..9e6ad7493f6 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -74,38 +74,29 @@ def duration_in_seconds(duration: str) -> int: elif unit == "w": return value * 604800 elif unit == "mo": - now: Final = time_module.time() - current_time: Final = datetime.fromtimestamp(now) - - # Calculate target month and year, handling overflow past December - total_months: Final = current_time.month - 1 + value # 0-indexed months - target_year: Final = current_time.year + total_months // 12 - target_month: Final = total_months % 12 + 1 # back to 1-indexed - - # Determine the day to set for next month - target_day = current_time.day - last_day_of_target_month: Final = get_last_day_of_month(target_year, target_month) - - target_day = min(target_day, last_day_of_target_month) - - next_month: Final = datetime( - year=target_year, - month=target_month, - day=target_day, - hour=current_time.hour, - minute=current_time.minute, - second=current_time.second, - microsecond=current_time.microsecond, - ) - - # Calculate the duration until the first day of the next month - duration_until_next_month: Final = next_month - current_time - return int(duration_until_next_month.total_seconds()) + return _months_in_seconds(start=datetime.fromtimestamp(time_module.time()), months=value) else: raise ValueError(f"Unsupported duration unit, passed duration: {duration}") +def duration_in_seconds_from(start: datetime, duration: str) -> int: + """How long a period of `duration` lasts when it begins at `start`. + + `duration_in_seconds` measures `Nmo` from now, so a monthly window that opened in a + month of a different length needs its own start to know when it ends. Every other + unit is a fixed length and ignores `start`. + """ + value, unit = _extract_from_regex(duration=_normalize_duration(duration)) + if unit == "mo": + return _months_in_seconds(start=start, months=value) + return duration_in_seconds(duration) + + +def _months_in_seconds(start: datetime, months: int) -> int: + return int((_shift_months(start, months) - start).total_seconds()) + + def get_next_standardized_reset_time( duration: str, current_time: datetime, @@ -166,8 +157,9 @@ def get_next_standardized_reset_time( return base_midnight + timedelta(days=1) -def _subtract_months(moment: datetime, months: int) -> datetime: - total_months: Final = moment.year * 12 + moment.month - 1 - months +def _shift_months(moment: datetime, months: int) -> datetime: + """`moment` moved by `months` calendar months, back when negative, clamped to the target month's last day.""" + total_months: Final = moment.year * 12 + moment.month - 1 + months year, month_index = divmod(total_months, 12) month: Final = month_index + 1 return moment.replace(year=year, month=month, day=min(moment.day, get_last_day_of_month(year, month))) @@ -182,9 +174,9 @@ def get_budget_window_start(duration: str, reset_at: datetime) -> datetime: return reset_at - timedelta(days=1) match unit: case "mo": - return _subtract_months(reset_at, value) + return _shift_months(reset_at, -value) case "d": - return _subtract_months(reset_at, 1) if value == 30 else reset_at - timedelta(days=value) + return _shift_months(reset_at, -1) if value == 30 else reset_at - timedelta(days=value) case "w": return reset_at - timedelta(weeks=value) case "h": diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index d0b39d98dd1..a518d60de05 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -14,7 +14,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span -from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.litellm_core_utils.duration_parser import duration_in_seconds, duration_in_seconds_from from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.router_strategy.budget_limiter import RouterBudgetLimiting @@ -303,7 +303,9 @@ def _readings(values: Sequence[float | None], window_count: int) -> tuple[_Windo def _window_reset_at(started_at: float | None, budget_duration: str | None, now: float) -> str | None: if started_at is None or budget_duration is None: return None - ends_at: Final = started_at + duration_in_seconds(budget_duration) + ends_at: Final = started_at + duration_in_seconds_from( + start=datetime.fromtimestamp(started_at), duration=budget_duration + ) if ends_at <= now: return None return datetime.fromtimestamp(ends_at, tz=timezone.utc).isoformat() diff --git a/tests/unit/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py index f5e50d67e71..9212da633e0 100644 --- a/tests/unit/litellm_core_utils/test_duration_parser.py +++ b/tests/unit/litellm_core_utils/test_duration_parser.py @@ -7,6 +7,7 @@ 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, + duration_in_seconds_from, get_budget_window_start, get_next_standardized_reset_time, ) @@ -373,5 +374,33 @@ class TestGetBudgetWindowStart(unittest.TestCase): ) +class TestDurationInSecondsFrom(unittest.TestCase): + def test_a_month_lasts_as_long_as_the_month_it_starts_in(self): + cases: Final = ( + (datetime(2026, 1, 15, 12), "1mo", 31), + (datetime(2026, 2, 10, 12), "1mo", 28), + (datetime(2026, 1, 31, 12), "1mo", 28), + (datetime(2028, 1, 31, 12), "1mo", 29), + (datetime(2026, 12, 15, 12), "1mo", 31), + (datetime(2026, 1, 15, 12), "2mo", 59), + ) + for start, duration, days in cases: + with self.subTest(start=start, duration=duration): + self.assertEqual(duration_in_seconds_from(start, duration), days * 86400) + + def test_fixed_length_units_do_not_depend_on_the_start(self): + for duration in ("45s", "30m", "1h", "1d", "7d", "2w", "30d", "daily", "monthly"): + with self.subTest(duration=duration): + self.assertEqual( + duration_in_seconds_from(datetime(2026, 2, 10, 12), duration), + duration_in_seconds(duration), + ) + + def test_a_month_starting_now_matches_duration_in_seconds(self): + now: Final = datetime(2026, 1, 31, 9, 30, 0, 250000) + with patch.object(duration_parser.time_module, "time", return_value=now.timestamp()): + self.assertEqual(duration_in_seconds("1mo"), duration_in_seconds_from(now, "1mo")) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index b17b5c34e37..b523e5fced0 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -1144,6 +1144,71 @@ async def test_usage_reports_when_the_open_window_resets(seconds_since_window_op assert usage["gpt-4"]["budget_reset_at"] == expected_reset_at +_DAY: Final = 86400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "opened_at,read_after_days,window_days", + [ + (datetime(2026, 1, 15, 12), 26, 31), + (datetime(2026, 1, 15, 12), 29, 31), + (datetime(2026, 1, 31, 12), 20, 28), + (datetime(2026, 12, 15, 12), 30, 31), + ], + ids=["read_in_a_shorter_month", "read_past_the_shorter_month_length", "clamped_to_february", "across_new_year"], +) +async def test_a_monthly_window_resets_a_calendar_month_after_it_opened(opened_at, read_after_days, window_days): + """ + A month is only as long as the month the window opened in, which is the lifetime the + counter was given when it opened. Measuring the month from the read instead made a + window opened on January 15 and read on February 10 report February 12, and then null + while the counter still refused requests until February 15. + """ + dual_cache = DualCache() + started_at: Final = opened_at.timestamp() + await dual_cache.async_set_cache( + key=model_budget_start_time_cache_key( + entity_type=Litellm_EntityType.USER, entity_id="user-1", budget_model="gpt-4", budget_duration="1mo" + ), + value=started_at, + ) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id="user-1", + model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, + cache=dual_cache, + now=datetime.fromtimestamp(started_at + read_after_days * _DAY, tz=timezone.utc), + ) + + assert usage["gpt-4"]["budget_reset_at"] == ( + datetime.fromtimestamp(started_at + window_days * _DAY, tz=timezone.utc).isoformat() + ) + + +@pytest.mark.asyncio +async def test_a_monthly_window_reports_no_reset_once_its_month_has_passed(): + dual_cache = DualCache() + started_at: Final = datetime(2026, 1, 15, 12).timestamp() + await dual_cache.async_set_cache( + key=model_budget_start_time_cache_key( + entity_type=Litellm_EntityType.USER, entity_id="user-1", budget_model="gpt-4", budget_duration="1mo" + ), + value=started_at, + ) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id="user-1", + model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, + cache=dual_cache, + now=datetime.fromtimestamp(started_at + 31 * _DAY, tz=timezone.utc), + ) + + assert usage["gpt-4"]["budget_reset_at"] is None + + @pytest.mark.asyncio async def test_each_budgeted_model_reports_its_own_window_reset(): """Two models on one key own separate windows, so one model's reset must never be reported for the other."""