diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 6662cf4a2c2..4a78e4e8e77 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1128,6 +1128,7 @@ class BudgetExceededError(Exception): llm_provider: str | None = None, entity_type: str | None = None, entity_id: str | None = None, + budget_reset_at: datetime.datetime | None = None, ): self.current_cost = current_cost self.max_budget = max_budget @@ -1135,6 +1136,7 @@ class BudgetExceededError(Exception): self.llm_provider = llm_provider or "" self.entity_type = entity_type self.entity_id = entity_id + self.budget_reset_at = budget_reset_at # Surface unified rate-limit fields without joining the RateLimitError # hierarchy so existing `except BudgetExceededError:` handlers keep # working; custom callbacks reading StandardLoggingPayload pick these diff --git a/litellm/models/budget.py b/litellm/models/budget.py index 61123810fd1..40188e131ab 100644 --- a/litellm/models/budget.py +++ b/litellm/models/budget.py @@ -70,3 +70,10 @@ class LiteLLM_TeamMemberTable(LiteLLM_BudgetTable): budget_id: str | None = None model_config = ConfigDict(protected_namespaces=()) + + +def budget_table_reset_at(budget_table: LiteLLM_BudgetTable | None) -> datetime | None: + """The budget's next reset, known only when the row was loaded with it.""" + if isinstance(budget_table, LiteLLM_BudgetTableFull): + return budget_table.budget_reset_at + return None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9d98e2c9973..00d687a8fd7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -53,6 +53,7 @@ from litellm.llms.base_llm.managed_resources.utils import ( extract_provider_resource_id_from_unified_id, is_base64_encoded_unified_id, ) +from litellm.models.budget import budget_table_reset_at from litellm.models.project import LiteLLM_ProjectTable from litellm.proxy._types import ( MODEL_LIST_ROUTES, @@ -149,7 +150,7 @@ from litellm.proxy.guardrails.tool_name_extraction import ( extract_request_tool_names, ) from litellm.proxy.route_llm_request import route_request -from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_reset_at, get_budget_window_start from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics @@ -1451,6 +1452,7 @@ async def common_checks( message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", entity_type=Litellm_EntityType.USER.value, entity_id=user_object.user_id, + budget_reset_at=user_object.budget_reset_at, ) # Each scope reads a distinct counter key with no cross-scope ordering @@ -1987,6 +1989,7 @@ async def check_end_user_budget( message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}", entity_type=Litellm_EntityType.END_USER.value, entity_id=end_user_obj.user_id, + budget_reset_at=budget_table_reset_at(end_user_obj.litellm_budget_table), ) @@ -6785,6 +6788,7 @@ async def virtual_key_max_budget_check( message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}", entity_type=Litellm_EntityType.KEY.value, entity_id=valid_token.token, + budget_reset_at=valid_token.budget_reset_at, ) @@ -6832,6 +6836,7 @@ async def _virtual_key_multi_budget_check( ), entity_type=Litellm_EntityType.KEY.value, entity_id=valid_token.token, + budget_reset_at=get_budget_window_reset_at(window), # pyright: ignore[reportUnknownArgumentType] # key budget_limits rows are untyped dicts ) @@ -7146,6 +7151,11 @@ async def _check_team_member_budget( message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}", entity_type=Litellm_EntityType.TEAM_MEMBER.value, entity_id=f"{valid_token.user_id}:{team_object.team_id}", + budget_reset_at=budget_table_reset_at( + member_budget_row + if member_budget_row is not None and member_budget_row.max_budget is not None + else team_default_budget + ), ) @@ -7261,6 +7271,7 @@ async def team_max_budget_check( message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}", entity_type=Litellm_EntityType.TEAM.value, entity_id=team_object.team_id, + budget_reset_at=team_object.budget_reset_at, ) @@ -7304,6 +7315,7 @@ async def _team_multi_budget_check( ), entity_type=Litellm_EntityType.TEAM.value, entity_id=team_object.team_id, + budget_reset_at=get_budget_window_reset_at(window), ) @@ -7433,6 +7445,7 @@ async def _project_max_budget_check( message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}", entity_type=Litellm_EntityType.PROJECT.value, entity_id=project_object.project_id, + budget_reset_at=budget_table_reset_at(project_object.litellm_budget_table), ) @@ -7641,6 +7654,7 @@ async def organization_max_budget_check( message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}", entity_type=Litellm_EntityType.ORGANIZATION.value, entity_id=org_id, + budget_reset_at=budget_table_reset_at(org_table.litellm_budget_table), ) @@ -7713,6 +7727,7 @@ async def tag_max_budget_check_for_tags( message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}", entity_type=Litellm_EntityType.TAG.value, entity_id=tag_name, + budget_reset_at=budget_table_reset_at(tag_object.litellm_budget_table), ) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 2ec2274db2b..4a03c7c28f5 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -3,7 +3,9 @@ Handles Authentication Errors """ import logging +import math from collections.abc import Mapping +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Final from fastapi import HTTPException, Request, status @@ -48,6 +50,19 @@ else: Span = Any +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +def budget_retry_after_seconds(budget_reset_at: datetime, now: datetime) -> int: + """Whole seconds until the budget resets, never below 1 so a client does not retry at once.""" + reset_at: Final = ( + budget_reset_at if budget_reset_at.tzinfo is not None else budget_reset_at.replace(tzinfo=timezone.utc) + ) + current: Final = now if now.tzinfo is not None else now.replace(tzinfo=timezone.utc) + return max(1, math.ceil((reset_at - current).total_seconds())) + + def _as_proxy_exception(e: Exception) -> ProxyException: """Convert an authentication failure into the ProxyException the client receives.""" if isinstance(e, litellm.BudgetExceededError): @@ -56,6 +71,11 @@ def _as_proxy_exception(e: Exception) -> ProxyException: type=ProxyErrorTypes.budget_exceeded, param=None, code=getattr(e, "status_code", status.HTTP_429_TOO_MANY_REQUESTS), + headers=( + {"retry-after": str(budget_retry_after_seconds(e.budget_reset_at, _utc_now()))} + if e.budget_reset_at is not None + else None + ), ) if isinstance(e, ModelAccessDeniedHTTPException): return ModelAccessDeniedProxyException( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 82135cb1b74..5146d3c1a32 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -42,6 +42,7 @@ from litellm.integrations.otel.runtime import phase_event, phase_span, seed_requ from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias, resolve_served_model +from litellm.models.budget import budget_table_reset_at from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import * from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers @@ -3125,6 +3126,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e ), entity_type=Litellm_EntityType.TEAM_MEMBER.value, entity_id=_entity_id, + budget_reset_at=budget_table_reset_at(team_member_info.litellm_budget_table), ) # Check 3. If token is expired diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index dcfd6148f54..ff85b74db74 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -67,6 +67,7 @@ class _BudgetCounter: spend_log_entity_id: str | None = None window_duration: str | None = None window_start: datetime | None = None + budget_reset_at: datetime | None = None _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { @@ -234,6 +235,7 @@ def _raise_counter_budget_exceeded( ), entity_type=_COUNTER_ENTITY_TYPES.get(counter.entity_type), entity_id=counter.spend_log_entity_id or counter.entity_id, + budget_reset_at=counter.budget_reset_at, ) @@ -495,6 +497,7 @@ async def _get_budget_counters( fallback_spend=float(valid_token.spend or 0.0), entity_type="Key", entity_id=valid_token.token, + budget_reset_at=valid_token.budget_reset_at, ) ) counters.extend( @@ -518,6 +521,7 @@ async def _get_budget_counters( fallback_spend=float(team_object.spend or 0.0), entity_type="Team", entity_id=team_id, + budget_reset_at=team_object.budget_reset_at, ) ) counters.extend( @@ -546,6 +550,7 @@ async def _get_budget_counters( fallback_spend=float(user_object.spend or 0.0), entity_type="User", entity_id=user_object.user_id, + budget_reset_at=user_object.budget_reset_at, ) ) @@ -629,6 +634,7 @@ async def _get_end_user_budget_counter( fallback_spend=fallback_spend, entity_type="EndUser", entity_id=end_user_id, + budget_reset_at=_reset_at_of(_get_value(end_user_object, "litellm_budget_table")), ) @@ -669,6 +675,7 @@ async def _get_tag_budget_counters( fallback_spend=_to_float(_get_value(tag_object, "spend")) or 0.0, entity_type="Tag", entity_id=tag_name, + budget_reset_at=_reset_at_of(budget_table), ) ) return counters @@ -725,6 +732,17 @@ def _dedupe_tags(tags: list[str]) -> list[str]: return deduped_tags +async def _cached_team_member_default_budget( + team_object: LiteLLM_TeamTable, user_api_key_cache: UserApiKeyCache +) -> object: + default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") + if not isinstance(default_budget_id, str): + return None + return _CACHED_VALUE.validate_python( + await user_api_key_cache.async_get_cache(key=f"team_member_default_budget:{default_budget_id}") + ) + + @with_service_target(AUTH_OBJECTS_TARGET) async def _get_team_member_budget_counter( valid_token: UserAPIKeyAuth, @@ -749,20 +767,20 @@ async def _get_team_member_budget_counter( member_budget_row: Final = team_membership.litellm_budget_table if team_membership is not None else None now: Final = datetime.now(timezone.utc) + member_row_sets_cap: Final = member_budget_row is not None and member_budget_row.max_budget is not None + default_budget: Final = ( + None if member_row_sets_cap else await _cached_team_member_default_budget(team_object, user_api_key_cache) + ) + team_member_reset_at: Final = _reset_at_of(member_budget_row if member_row_sets_cap else default_budget) team_member_budget: float | None = None if member_budget_row is not None and member_budget_row.max_budget is not None: team_member_budget = member_budget_row.effective_max_budget(now=now) else: - default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") - if isinstance(default_budget_id, str): - default_budget: Final = _CACHED_VALUE.validate_python( - await user_api_key_cache.async_get_cache(key=f"team_member_default_budget:{default_budget_id}") + default_cap: Final = _to_float(_get_value(default_budget, "max_budget")) + if default_cap is not None and default_cap > 0: + team_member_budget = default_cap + ( + member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 ) - default_cap: Final = _to_float(_get_value(default_budget, "max_budget")) - if default_cap is not None and default_cap > 0: - team_member_budget = default_cap + ( - member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 - ) if team_member_budget is None or team_member_budget <= 0: return None @@ -775,6 +793,7 @@ async def _get_team_member_budget_counter( fallback_spend=float(team_member_spend or 0.0), entity_type="TeamMember", entity_id=f"{valid_token.user_id}:{team_object.team_id}", + budget_reset_at=team_member_reset_at, ) @@ -814,6 +833,7 @@ async def _get_org_budget_counter( fallback_spend=org_spend, entity_type="Organization", entity_id=org_id, + budget_reset_at=_reset_at_of(org_budget_table), ) @@ -847,6 +867,7 @@ async def _get_project_budget_counter( fallback_spend=_to_float(_get_value(project_object, "spend")) or 0.0, entity_type="Project", entity_id=valid_token.project_id, + budget_reset_at=_reset_at_of(project_budget_table), ) @@ -886,6 +907,7 @@ def _get_budget_limit_counters( spend_log_entity_id=entity_id, window_duration=str(budget_duration), window_start=window_start, + budget_reset_at=get_budget_window_reset_at(window_dict), ) ) return counters @@ -1282,6 +1304,17 @@ def get_budget_window_start(window: object) -> datetime | None: return reset_at - timedelta(seconds=duration_seconds) +def get_budget_window_reset_at(window: object) -> datetime | None: + reset_at: Final = _coerce_datetime(_coerce_window(window).get("reset_at")) + if reset_at is None: + return None + return reset_at if reset_at.tzinfo is not None else reset_at.replace(tzinfo=timezone.utc) + + +def _reset_at_of(budget_table: object) -> datetime | None: + return _coerce_datetime(_get_value(budget_table, "budget_reset_at")) + + def _coerce_datetime(value: object) -> datetime | None: if value is None: return None diff --git a/tests/unit/models/test_models.py b/tests/unit/models/test_models.py index e2c0edf6d9d..a73ce3668d1 100644 --- a/tests/unit/models/test_models.py +++ b/tests/unit/models/test_models.py @@ -14,6 +14,7 @@ from litellm.models.budget import ( LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, LiteLLM_TeamMemberTable, + budget_table_reset_at, ) from litellm.models.config import LiteLLM_Config from litellm.models.credentials import CreateCredentialItem, CredentialItem @@ -690,3 +691,12 @@ class TestAutoRouterSession: def test_a_partial_comparison_against_one_baseline_keeps_its_label(self) -> None: assert self._row({"anthropic/claude-opus-5": 3}, estimated_turns=1).baseline_model == "anthropic/claude-opus-5" + + +def test_budget_table_reset_at_is_known_only_on_the_full_row(): + reset_at = datetime(2026, 11, 3, tzinfo=timezone.utc) + full_row = LiteLLM_BudgetTableFull(max_budget=1.0, budget_reset_at=reset_at, created_at=reset_at) + + assert budget_table_reset_at(full_row) == reset_at + assert budget_table_reset_at(LiteLLM_BudgetTable(max_budget=1.0)) is None + assert budget_table_reset_at(None) is None diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index 9d4fe67614e..16c44a56faf 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -3,6 +3,7 @@ import sys, os, asyncio, time, random, uuid import hashlib +from datetime import datetime, timezone import json import traceback from dotenv import load_dotenv @@ -991,6 +992,25 @@ async def test_virtual_key_max_budget_check( assert alert_called, "Budget alert should be triggered" +@pytest.mark.asyncio +async def test_virtual_key_max_budget_check_carries_the_budget_reset_at(): + from litellm.proxy.auth.auth_checks import virtual_key_max_budget_check + + reset_at: Final = datetime(2026, 11, 3, tzinfo=timezone.utc) + valid_token = UserAPIKeyAuth(token="test-token", spend=10.0, max_budget=10.0, budget_reset_at=reset_at) + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + + async def mock_budget_alert(*args, **kwargs): + return None + + proxy_logging_obj.budget_alerts = mock_budget_alert + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await virtual_key_max_budget_check(valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, user_obj=None) + + assert exc_info.value.budget_reset_at == reset_at + + @pytest.mark.parametrize( "model, team_models, expect_to_work", [ diff --git a/tests/unit/proxy/auth/test_auth_exception_handler.py b/tests/unit/proxy/auth/test_auth_exception_handler.py index 2523284ec2c..9c4fdf83ed8 100644 --- a/tests/unit/proxy/auth/test_auth_exception_handler.py +++ b/tests/unit/proxy/auth/test_auth_exception_handler.py @@ -1,5 +1,6 @@ import asyncio import json +from datetime import datetime, timedelta, timezone from collections.abc import Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -38,7 +39,12 @@ from litellm.proxy._types import ( ProxyException, UserAPIKeyAuth, ) -from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler, _as_proxy_exception +from litellm.proxy.auth import auth_exception_handler +from litellm.proxy.auth.auth_exception_handler import ( + UserAPIKeyAuthExceptionHandler, + _as_proxy_exception, + budget_retry_after_seconds, +) from litellm.proxy.auth.model_access_denied import ModelAccessDeniedHTTPException @@ -1312,3 +1318,47 @@ def test_auth_failure_error_body_carries_a_json_null_param(auth_failure: Excepti assert body["error"]["type"] == ProxyErrorTypes.auth_error.value assert body["error"]["param"] is None + + +FIXED_NOW: Final = datetime(2026, 11, 2, 12, 0, tzinfo=timezone.utc) + + +@pytest.mark.parametrize( + "budget_reset_at, expected", + [ + (FIXED_NOW + timedelta(seconds=90.4), 91), + (FIXED_NOW + timedelta(seconds=60), 60), + (FIXED_NOW - timedelta(seconds=5), 1), + (datetime(2026, 11, 2, 12, 0, 30), 30), + ], +) +def test_budget_retry_after_seconds(budget_reset_at: datetime, expected: int): + assert budget_retry_after_seconds(budget_reset_at, FIXED_NOW) == expected + + +@pytest.mark.parametrize( + "budget_reset_at, expected_header", + [ + (FIXED_NOW + timedelta(days=1), "86400"), + (FIXED_NOW - timedelta(minutes=1), "1"), + (None, None), + ], +) +@pytest.mark.asyncio +async def test_handle_authentication_error_budget_exceeded_retry_after( + monkeypatch: pytest.MonkeyPatch, budget_reset_at: datetime | None, expected_header: str | None +): + monkeypatch.setattr(auth_exception_handler, "_utc_now", lambda: FIXED_NOW) + handler = UserAPIKeyAuthExceptionHandler() + budget_error = litellm.BudgetExceededError( + message="Budget exceeded", + current_cost=100, + max_budget=100, + budget_reset_at=budget_reset_at, + ) + + with pytest.raises(ProxyException) as exc_info: + await handler.handle_authentication_error(budget_error, MagicMock(), {}, "/test", None, "test-key") + + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert exc_info.value.headers.get("retry-after") == expected_header diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index a87ac7fa7c0..7a32bb02526 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -24,7 +24,10 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_reservation_cache_key, ) from litellm.proxy.spend_tracking.budget_reservation import ( + _BudgetCounter, _get_team_member_budget_counter, + _raise_counter_budget_exceeded, + get_budget_window_reset_at, estimate_request_max_cost, get_budget_window_start, release_unbound_budget_reservation, @@ -588,3 +591,38 @@ async def test_group_reserves_on_its_current_deployments_after_another_vendors_d by_own_name: Final = litellm.token_counter(model="prod", messages=_MULTILINGUAL_MESSAGES) assert openai_only != by_own_name assert (await _input_tokens(router, "prod"))["prod"] == by_own_name + + +def test_counter_budget_exceeded_carries_the_counter_reset_at(): + reset_at: Final = datetime(2026, 11, 3, tzinfo=timezone.utc) + counter = _BudgetCounter( + counter_key="spend:key:k", + max_budget=1.0, + fallback_spend=1.0, + entity_type="Key", + entity_id="k", + budget_reset_at=reset_at, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + _raise_counter_budget_exceeded(counter=counter, current_cost=1.5) + + assert exc_info.value.budget_reset_at == reset_at + + +@pytest.mark.parametrize( + "window, expected", + [ + ( + {"budget_duration": "1d", "max_budget": 1.0, "reset_at": "2026-11-03T00:00:00Z"}, + datetime(2026, 11, 3, tzinfo=timezone.utc), + ), + ( + {"budget_duration": "1d", "max_budget": 1.0, "reset_at": datetime(2026, 11, 3)}, + datetime(2026, 11, 3, tzinfo=timezone.utc), + ), + ({"budget_duration": "1d", "max_budget": 1.0}, None), + ], +) +def test_get_budget_window_reset_at(window: dict[str, object], expected: datetime | None): + assert get_budget_window_reset_at(window) == expected