diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 95bc77e3eba..759ffbffcb1 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -167,6 +167,15 @@ def _shift_months(moment: datetime, months: int) -> datetime: return moment.replace(year=year, month=month, day=min(moment.day, get_last_day_of_month(year, month))) +def is_annual_budget_duration(duration: str) -> bool: + return _parse_duration(normalize_duration(duration)) == (12, "mo") + + +def get_annual_budget_window_start(reset_at: datetime, timezone_str: str) -> datetime: + boundary: Final = _setup_timezone(reset_at, timezone_str)[0] + return _shift_months(boundary, -12).astimezone(timezone.utc) + + def get_budget_window_start(duration: str, reset_at: datetime) -> datetime: """Start of the budget period that ends at `reset_at`, under the same rules `get_next_standardized_reset_time` used to pick it: `30d` and `Nmo` span calendar @@ -390,14 +399,22 @@ def _handle_month_reset( Args: current_time: Current datetime base_midnight: Midnight of current day - value: Number of months (currently only supports 1 month resets) + value: Number of months, 1 for monthly or 12 for annual resets reset_time_of_day: Wall-clock time the reset lands on Returns: datetime: First day of the next reset month at `reset_time_of_day` """ + if value == 12: + first_of_year: Final = base_midnight.replace(month=1, day=1) + annual_candidate: Final = _apply_time_of_day(first_of_year, reset_time_of_day) + return ( + annual_candidate.replace(year=annual_candidate.year + 1) + if annual_candidate <= current_time + else annual_candidate + ) if value != 1: - raise ValueError("Monthly resets currently only support 1 month intervals") + raise ValueError("Calendar resets support 1 month or 12 month intervals") first_of_this_month: Final = base_midnight.replace(day=1) candidate: Final = _apply_time_of_day(first_of_this_month, reset_time_of_day) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index ae91cbb1f2a..92960bf9d21 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -30,7 +30,11 @@ from litellm.constants import ( RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN, RESET_BUDGET_JOB_NAME, ) -from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.litellm_core_utils.duration_parser import ( + duration_in_seconds, + get_annual_budget_window_start, + is_annual_budget_duration, +) from litellm.proxy._types import ( DB_RETRY_SAFE_ERROR_TYPES, LiteLLM_BudgetTableFull, @@ -1514,16 +1518,21 @@ class ResetBudgetJob: reset_at_str: Final = window.get("reset_at") if not reset_at_str: return False - reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None) - if reset_at > now: + budget_duration: Final = window["budget_duration"] + annual: Final = isinstance(budget_duration, str) and is_annual_budget_duration(budget_duration) + saved_reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")) + reset_at: Final = _as_utc(saved_reset_at) if annual else saved_reset_at.replace(tzinfo=None) + cutoff: Final = _as_utc(now) if annual else now + if reset_at > cutoff: return False new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_store) try: await spend_counter_store.set(counter_key, new_value) except Exception as redis_err: verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err) - budget_duration: Final = window["budget_duration"] - next_reset_at: Final = compute_budget_reset_at(budget_duration=budget_duration, settings=reset_settings) + next_reset_at: Final = compute_budget_reset_at( + budget_duration=budget_duration, settings=reset_settings, now=now if annual else None + ) window["reset_at"] = next_reset_at.isoformat() await ResetBudgetJob._roll_window_spend_row( prisma_client=prisma_client, @@ -1531,6 +1540,7 @@ class ResetBudgetJob: entity_id=entity_id, budget_duration=budget_duration, next_reset_at=next_reset_at, + reset_settings=reset_settings, ) return True @@ -1541,6 +1551,7 @@ class ResetBudgetJob: entity_id: str, budget_duration: str, next_reset_at: datetime, + reset_settings: BudgetResetSettings | None = None, ) -> None: """Move this window's LiteLLM_BudgetWindowSpend row onto the window that just started, so the maintained total the read path uses starts @@ -1551,7 +1562,11 @@ class ResetBudgetJob: windows from having their counters reset. """ try: - window_start: Final = next_reset_at - timedelta(seconds=duration_in_seconds(budget_duration)) + window_start: Final = ( + get_annual_budget_window_start(next_reset_at, (reset_settings or get_budget_reset_settings()).timezone) + if is_annual_budget_duration(budget_duration) + else next_reset_at - timedelta(seconds=duration_in_seconds(budget_duration)) + ) except Exception as e: # noqa: BLE001 # duration_in_seconds raises bare exceptions on bad input verbose_proxy_logger.warning("Unparseable budget_duration %s: %s", budget_duration, e) return @@ -1708,7 +1723,9 @@ class ResetBudgetJob: item.spend = _carried_spend(item.spend, _rollover_cap(item.max_budget)) if _rollover_enabled() else 0.0 if hasattr(item, "budget_duration") and item.budget_duration is not None: item.budget_reset_at = compute_budget_reset_at( - budget_duration=item.budget_duration, settings=reset_settings + budget_duration=item.budget_duration, + settings=reset_settings, + now=current_time if is_annual_budget_duration(item.budget_duration) else None, ) return item except Exception as e: diff --git a/litellm/proxy/common_utils/timezone_utils.py b/litellm/proxy/common_utils/timezone_utils.py index 7a40931f50d..1a44585f75e 100644 --- a/litellm/proxy/common_utils/timezone_utils.py +++ b/litellm/proxy/common_utils/timezone_utils.py @@ -63,11 +63,13 @@ def get_budget_reset_settings() -> BudgetResetSettings: ) -def compute_budget_reset_at(budget_duration: str, settings: BudgetResetSettings) -> datetime: +def compute_budget_reset_at( + budget_duration: str, settings: BudgetResetSettings, now: datetime | None = None +) -> datetime: """Compute the next reset time for a budget duration using injected settings.""" return get_next_standardized_reset_time( duration=budget_duration, - current_time=datetime.now(timezone.utc), + current_time=now if now is not None else datetime.now(timezone.utc), timezone_str=settings.timezone, reset_time_of_day=settings.reset_time_of_day, ) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 096578264e3..7355da72c36 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -18,7 +18,11 @@ import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.litellm_core_utils.duration_parser import ( + duration_in_seconds, + get_annual_budget_window_start, + is_annual_budget_duration, +) from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.proxy._types import ( Litellm_EntityType, @@ -30,6 +34,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings, get_budget_reset_settings from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, UserApiKeyCache, @@ -1287,12 +1292,18 @@ def fixed_budget_window(window: object) -> tuple[str, datetime] | None: return duration, start -def get_budget_window_start(window: object) -> datetime | None: +def get_budget_window_start(window: object, reset_settings: BudgetResetSettings | None = None) -> datetime | None: window_dict: Final = _coerce_window(window) budget_duration: Final = window_dict.get("budget_duration") if budget_duration is None: return None try: + if is_annual_budget_duration(str(budget_duration)): + annual_reset_at: Final = _coerce_datetime(window_dict.get("reset_at")) + if annual_reset_at is not None: + return get_annual_budget_window_start( + annual_reset_at, (reset_settings or get_budget_reset_settings()).timezone + ) duration_seconds: Final = duration_in_seconds(str(budget_duration)) except Exception: return None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7ff2b9e236e..262cf88f696 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -175,7 +175,11 @@ from litellm.litellm_core_utils.core_helpers import ( require_attributes, ) from litellm.litellm_core_utils.deployment_pricing import pricing_models -from litellm.litellm_core_utils.duration_parser import get_budget_window_start +from litellm.litellm_core_utils.duration_parser import ( + get_annual_budget_window_start, + get_budget_window_start, + is_annual_budget_duration, +) from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int from litellm.litellm_core_utils.litellm_logging import ( BilledRoundsDispatch, @@ -208,6 +212,7 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolv from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_settings from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db import tool_registry_writer from litellm.proxy.db.create_views import ( @@ -9535,7 +9540,11 @@ def _as_aware(moment: datetime) -> datetime: def _reset_window(budget_duration: str, budget_reset_at: datetime) -> tuple[datetime, datetime] | None: window_end: Final = _as_aware(budget_reset_at) - window_start: Final = get_budget_window_start(budget_duration, window_end) + window_start: Final = ( + get_annual_budget_window_start(window_end, get_budget_reset_settings().timezone) + if is_annual_budget_duration(budget_duration) + else get_budget_window_start(budget_duration, window_end) + ) return (window_start, window_end) if window_start < window_end else None diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx index 67d3f45ec95..efc596075fa 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx @@ -47,6 +47,7 @@ const RESET_OPTIONS: readonly { value: string; labelKey: ParseKeys }[] = [ { value: "24h", labelKey: "budgets.form.resetOptions.perDay" }, { value: "7d", labelKey: "budgets.form.resetOptions.perWeek" }, { value: "30d", labelKey: "budgets.form.resetOptions.perMonth" }, + { value: "12mo", labelKey: "budgets.form.resetOptions.perYear" }, ]; const numberProps = void }>({ diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.integration.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.integration.test.tsx index 0770150ab91..9dbfd61a79f 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.integration.test.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.integration.test.tsx @@ -25,7 +25,10 @@ describe("BudgetModal", () => { createMock.mockResolvedValue(undefined); }); - it("submits the spend limit and reset period alongside the rate limits", async () => { + it.each([ + ["per month", "30d"], + ["per year", "12mo"], + ])("submits the spend limit %s alongside the rate limits", async (label, duration) => { const user = userEvent.setup(); renderModal(); @@ -33,7 +36,7 @@ describe("BudgetModal", () => { expect(resetPeriod()).toBeVisible(); fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } }); fireEvent.change(screen.getByLabelText(/^Spend limit/), { target: { value: "42.567" } }); - await chooseSelectOption(user, resetPeriod(), "per month"); + await chooseSelectOption(user, resetPeriod(), label); fireEvent.change(screen.getByLabelText("Tokens per minute"), { target: { value: "500" } }); fireEvent.change(screen.getByLabelText("Requests per minute"), { target: { value: "7" } }); fireEvent.change(screen.getByLabelText("Batch tokens per day"), { target: { value: "90000" } }); @@ -43,7 +46,7 @@ describe("BudgetModal", () => { const expected = { budget_id: "budget-alpha", max_budget: 42.57, - budget_duration: "30d", + budget_duration: duration, tpm_limit: 500, rpm_limit: 7, tpd_limit: 90000, diff --git a/litellm/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.test.tsx b/litellm/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.test.tsx index 184e4297a67..45020685e9e 100644 --- a/litellm/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.test.tsx +++ b/litellm/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.test.tsx @@ -2,7 +2,7 @@ import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; -import BudgetDurationDropdown from "./budget_duration_dropdown"; +import BudgetDurationDropdown, { getBudgetDurationLabel } from "./budget_duration_dropdown"; describe("BudgetDurationDropdown", () => { it("labels an unset period as Never resets in the trigger and option", async () => { @@ -39,3 +39,13 @@ describe("BudgetDurationDropdown", () => { expect(await screen.findByRole("option", { name: "Not set" })).toBeInTheDocument(); }); }); + +it("saves and displays an annual budget period", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + render(); + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByRole("option", { name: "annually" })); + expect(onChange).toHaveBeenCalledWith("12mo", expect.anything()); + expect(getBudgetDurationLabel("12mo")).toBe("annually"); +}); diff --git a/litellm/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx b/litellm/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx index 7e5aeff8ffe..af3e33b624c 100644 --- a/litellm/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx +++ b/litellm/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx @@ -12,6 +12,7 @@ const DURATION_LABELS: Record = { "24h": "budgetDurationDropdown.periods.daily", "7d": "budgetDurationDropdown.periods.weekly", "30d": "budgetDurationDropdown.periods.monthly", + "12mo": "budgetDurationDropdown.periods.annual", }; interface BudgetDurationDropdownProps { @@ -51,6 +52,7 @@ const BudgetDurationDropdown: React.FC = ({ {t("budgetDurationDropdown.periods.daily")} {t("budgetDurationDropdown.periods.weekly")} {t("budgetDurationDropdown.periods.monthly")} + {t("budgetDurationDropdown.periods.annual")} ); @@ -64,6 +66,7 @@ export const getBudgetDurationLabel = (value: string | null | undefined, t: TFun "24h": "budgetDurationDropdown.periods.daily", "7d": "budgetDurationDropdown.periods.weekly", "30d": "budgetDurationDropdown.periods.monthly", + "12mo": "budgetDurationDropdown.periods.annual", }; return Object.hasOwn(budgetDurationMap, value) ? t(budgetDurationMap[value]) : value; diff --git a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json index 7064b6c476d..0e8c1c092e7 100644 --- a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json +++ b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json @@ -3072,7 +3072,8 @@ "hourly": "hourly", "daily": "daily", "weekly": "weekly", - "monthly": "monthly" + "monthly": "monthly", + "annual": "annually" }, "notSet": "Not set" }, @@ -3406,7 +3407,8 @@ "inTotal": "in total", "perDay": "per day", "perWeek": "per week", - "perMonth": "per month" + "perMonth": "per month", + "perYear": "per year" }, "noLimit": "No limit", "optional": "Optional", diff --git a/tests/unit/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py index f14015effe6..2f74762edaa 100644 --- a/tests/unit/litellm_core_utils/test_duration_parser.py +++ b/tests/unit/litellm_core_utils/test_duration_parser.py @@ -10,6 +10,7 @@ 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_annual_budget_window_start, get_budget_window_start, get_next_standardized_reset_time, ) @@ -195,6 +196,79 @@ class TestStandardizedResetTime(unittest.TestCase): self.assertEqual(result, expected) +@pytest.mark.parametrize( + "now, zone, reset_time, expected", + [ + ( + datetime(2026, 11, 12, tzinfo=timezone.utc), + "UTC", + time(0), + datetime(2027, 1, 1, tzinfo=timezone.utc), + ), + ( + datetime(2028, 2, 29, tzinfo=timezone.utc), + "UTC", + time(0), + datetime(2029, 1, 1, tzinfo=timezone.utc), + ), + ( + datetime(2027, 1, 1, tzinfo=timezone.utc), + "UTC", + time(0), + datetime(2028, 1, 1, tzinfo=timezone.utc), + ), + ( + datetime(2027, 1, 1, 8, tzinfo=timezone.utc), + "UTC", + time(12), + datetime(2027, 1, 1, 12, tzinfo=timezone.utc), + ), + ( + datetime(2027, 1, 1, 12, tzinfo=timezone.utc), + "UTC", + time(12), + datetime(2028, 1, 1, 12, tzinfo=timezone.utc), + ), + ( + datetime(2027, 1, 1, 2, tzinfo=timezone.utc), + "America/New_York", + time(0), + datetime(2027, 1, 1, tzinfo=ZoneInfo("America/New_York")), + ), + ], +) +def test_annual_reset_and_window_start( + now: datetime, zone: str, reset_time: time, expected: datetime +) -> None: + reset: Final = get_next_standardized_reset_time("12mo", now, zone, reset_time) + assert reset == expected + assert get_budget_window_start("12mo", reset) == expected.replace(year=expected.year - 1) + + +@pytest.mark.parametrize( + "zone,reset_at,expected", + ( + ("UTC", "2029-01-01T00:00:00+00:00", "2028-01-01T00:00:00+00:00"), + ("UTC", "2101-01-01T00:00:00+00:00", "2100-01-01T00:00:00+00:00"), + ("America/New_York", "2029-01-01T00:00:00-05:00", "2028-01-01T05:00:00+00:00"), + ("Asia/Kathmandu", "2029-01-01T12:30:00+05:45", "2028-01-01T06:45:00+00:00"), + ("Pacific/Apia", "2012-01-01T00:00:00+14:00", "2011-01-01T10:00:00+00:00"), + ), +) +@pytest.mark.parametrize("storage", ("offset", "utc", "naive_utc")) +def test_annual_window_survives_timestamp_storage( + zone: str, reset_at: str, expected: str, storage: str +) -> None: + reset: Final = datetime.fromisoformat(reset_at) + stored: Final = { + "offset": reset, + "utc": reset.astimezone(timezone.utc), + "naive_utc": reset.astimezone(timezone.utc).replace(tzinfo=None), + }[storage] + + assert get_annual_budget_window_start(stored, zone) == datetime.fromisoformat(expected) + + class TestResetTimeOfDay(unittest.TestCase): """A configurable reset_time_of_day shifts day/week/month resets off midnight.""" diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index f0ce610be01..544b6fb2082 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -28,11 +28,71 @@ from litellm.proxy.common_utils.reset_budget_job import ( _write_team_windows, ) from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings +from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.utils import PrismaClient from tests.unit.proxy.db.fake_prisma_engine import engine_call from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters +@pytest.mark.asyncio +@pytest.mark.parametrize( + "reset_at,now", + ( + ("2027-01-01T00:00:00-05:00", "2027-01-01T04:59:59+00:00"), + ("2027-01-01T00:00:00+05:45", "2026-12-31T18:14:59+00:00"), + ), +) +async def test_annual_window_reset_waits_for_the_saved_instant(reset_at: str, now: str) -> None: + store: Final = FakeSpendCounters().store() + counter: Final = "spend:key:offset-budget:window:12mo" + await store.set(counter, 10.0) + window: Final = {"budget_duration": "12mo", "reset_at": reset_at, "max_budget": 10.0} + client: Final = MagicMock(spec=PrismaClient) + changed: Final = await ResetBudgetJob._reset_expired_window( + window, + counter, + store, + datetime.fromisoformat(now), + BudgetResetSettings(timezone="America/New_York"), + client, + reset_budget_job_module.Litellm_EntityType.KEY, + "offset-budget", + ) + + assert changed is False + assert (await store.read(counter))[0] == 10.0 + assert window["reset_at"] == reset_at + + +@pytest.mark.asyncio +@pytest.mark.parametrize("today", ["2028-03-01", "2029-03-01"]) +async def test_annual_reset_persists_the_same_calendar_start_as_admission(today: str) -> None: # test-quality-ok: the SQL write is the reset job's observable persistence contract + client: Final = MagicMock(spec=PrismaClient) + client.db = MagicMock() + client.db.execute_raw = AsyncMock(return_value=1) + + with patch("time.time", return_value=datetime.fromisoformat(today).replace(tzinfo=timezone.utc).timestamp()): + await ResetBudgetJob._roll_window_spend_row( + prisma_client=client, + entity_type=reset_budget_job_module.Litellm_EntityType.KEY, + entity_id="annual-key", + budget_duration="12mo", + next_reset_at=datetime(2029, 1, 1, tzinfo=timezone.utc), + reset_settings=BudgetResetSettings(), + ) + + assert client.db.execute_raw.await_count == 1 + assert client.db.execute_raw.await_args.args[1:5] == ( + "key", + "annual-key", + "12mo", + datetime(2028, 1, 1), + ) + assert get_budget_window_start( + {"budget_duration": "12mo", "reset_at": datetime(2029, 1, 1, tzinfo=timezone.utc)} + ) == client.db.execute_raw.await_args.args[4].replace(tzinfo=timezone.utc) + + # Mock classes for testing class MockTable: """A single prisma table: records reads/writes and replays canned rows.""" diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 087434c981f..4b9ebf6235d 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -2,7 +2,9 @@ from __future__ import annotations import math from datetime import datetime, timedelta, timezone +from types import SimpleNamespace from typing import Final +from unittest.mock import patch import pytest @@ -23,6 +25,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( project_cache_key, team_membership_reservation_cache_key, ) +from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings from litellm.proxy.spend_tracking.budget_reservation import ( _get_team_member_budget_counter, estimate_request_max_cost, @@ -465,6 +468,88 @@ def test_budget_window_start_reads_json_encoded_windows(window: str, expected: d assert get_budget_window_start(window) == expected +@pytest.mark.parametrize("today", ["2028-03-01", "2029-03-01"]) +@pytest.mark.parametrize( + ("duration", "reset_at", "expected"), + [ + ("12mo", "2029-01-01T00:00:00Z", datetime(2028, 1, 1, tzinfo=timezone.utc)), + ("12mo", "2028-01-01T00:00:00", datetime(2027, 1, 1, tzinfo=timezone.utc)), + ], +) +def test_annual_window_keeps_all_spend_from_its_reset_boundary( + today: str, duration: str, reset_at: str, expected: datetime +) -> None: + with patch("time.time", return_value=datetime.fromisoformat(today).replace(tzinfo=timezone.utc).timestamp()): + assert get_budget_window_start({"budget_duration": duration, "reset_at": reset_at}) == expected + + +class _AnnualWindowDatabase: + def __init__(self) -> None: + self.db = self + self.litellm_budgetwindowspend = self + self.litellm_spendlogs = self + self.start: datetime | None = None + self.spend = 0.0 + self.scans = 0 + + async def execute_raw(self, sql: str, *parameters: object) -> int: + start: Final = parameters[3] + assert isinstance(start, datetime) + self.start = start + self.spend = 0.0 + return 1 + + async def find_unique(self, **query: object) -> SimpleNamespace: + return SimpleNamespace(window_start=self.start, spend=self.spend) + + async def group_by(self, **query: object) -> list[dict[str, object]]: + self.scans += 1 + return [{"api_key": "calendar-key", "_sum": {"spend": 0.0}}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "duration,zone,reset_at", + ( + ("12mo", "UTC", "2029-01-01T00:00:00Z"), + ("12mo", "America/New_York", "2029-01-01T00:00:00-05:00"), + ("12mo", "Pacific/Apia", "2012-01-01T00:00:00+14:00"), + ), +) +async def test_cold_annual_counter_preserves_reset_writer_spend( + monkeypatch: pytest.MonkeyPatch, duration: str, zone: str, reset_at: str +) -> None: + from litellm.proxy.auth.auth_checks import _virtual_key_multi_budget_check + from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob + from litellm.types.integrations.slack_alerting import Litellm_EntityType + + database: Final = _AnnualWindowDatabase() + store: Final = SpendCounterStore(connect(None).spend_counters()) + monkeypatch.setattr(litellm, "timezone", zone, raising=False) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "spend_counter_store", store) + monkeypatch.setattr(proxy_server, "spend_counter_cache", DualCache()) + await ResetBudgetJob._roll_window_spend_row( + prisma_client=database, # pyright: ignore[reportArgumentType] # in-memory external database boundary + entity_type=Litellm_EntityType.KEY, + entity_id="calendar-key", + budget_duration=duration, + next_reset_at=datetime.fromisoformat(reset_at.replace("Z", "+00:00")), + reset_settings=BudgetResetSettings(timezone=zone), + ) + database.spend = 10.0 + token: Final = UserAPIKeyAuth( + token="calendar-key", + budget_limits=[{"budget_duration": duration, "reset_at": reset_at, "max_budget": 10.0}], + ) + + with pytest.raises(litellm.BudgetExceededError) as failure: + await _virtual_key_multi_budget_check(token) + assert failure.value.current_cost == 10.0 + assert store.get_local(f"spend:key:calendar-key:window:{duration}") == 10.0 + assert database.scans == 0 + + @pytest.mark.asyncio async def test_reservation_counts_a_pdf_with_the_deployment_behind_a_gpt_named_group() -> None: from litellm.proxy.spend_tracking.input_tokens import count_input_tokens_for_model diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 5a6c9e13ab6..09a896ec053 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -13933,12 +13933,13 @@ async def _drain(queue): @pytest.mark.asyncio -async def test_key_window_spend_row_is_enqueued_with_the_actual_cost(): +@pytest.mark.parametrize("duration", ("30d", "12mo")) +async def test_key_window_spend_row_is_enqueued_with_the_actual_cost(duration: str): from litellm.proxy.proxy_server import increment_spend_counters - reset_at = datetime.now(timezone.utc) + timedelta(days=10) + reset_at: Final = datetime(2029, 1, 1, tzinfo=timezone.utc) key_obj = MagicMock() - key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}] + key_obj.budget_limits = [{"budget_duration": duration, "max_budget": 100.0, "reset_at": reset_at.isoformat()}] with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) @@ -13947,9 +13948,10 @@ async def test_key_window_spend_row_is_enqueued_with_the_actual_cost(): assert len(enqueued) == 1 assert enqueued[0]["entity_type"] == "key" assert enqueued[0]["entity_id"] == "hashed-token" - assert enqueued[0]["window_duration"] == "30d" + assert enqueued[0]["window_duration"] == duration assert enqueued[0]["spend"] == pytest.approx(0.25) - assert enqueued[0]["window_start"] == (reset_at - timedelta(days=30)).astimezone(timezone.utc).replace( + expected_start: Final = datetime(2028, 1, 1, tzinfo=timezone.utc) if duration == "12mo" else reset_at - timedelta(days=30) + assert enqueued[0]["window_start"] == expected_start.astimezone(timezone.utc).replace( tzinfo=None ).isoformat(timespec="microseconds")