scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

Merge pull request #41345 from BerriAI/litellm_remove_duplicate_user_budget_hook

Browse source 
fix(proxy): remove duplicate user budget hook that 429'd zero-cost models
This commit is contained in:
ryan-crabbe-berri 2026-09-16 16:06:41 -07:00 • committed by GitHub
commit be6e7ef173
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 111 additions and 476 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
ARCHITECTURE.md
View file

@ -45,7 +45,7 @@ sequenceDiagram
ProxyServer->>Auth: user_api_key_auth()
Auth->>Redis: Check API key cache
Redis-->>Auth: Key info + spend limits
ProxyServer->>Hooks: max_budget_limiter, parallel_request_limiter
ProxyServer->>Hooks: parallel_request_limiter, cache_control_check
Hooks->>Redis: Check/increment rate limit counters
ProxyServer->>Router: route_request()
Router->>Main: litellm.acompletion()
@ -145,7 +145,6 @@ graph TD
| Hook | File | Purpose |
|------|------|---------|
| `max_budget_limiter` | `proxy/hooks/max_budget_limiter.py` | Enforce budget limits |
| `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
| `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
| `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |

2
litellm/proxy/common_utils/proxy_rate_limit_error.py
View file

@ -11,7 +11,7 @@ exception types:
an upstream LLM provider returns 429.
* :class:`fastapi.HTTPException` (status 429) — raised directly by proxy hooks
such as ``parallel_request_limiter``, ``dynamic_rate_limiter``,
``batch_rate_limiter``, ``max_budget_limiter``, ``max_iterations_limiter``,
``batch_rate_limiter``, ``max_iterations_limiter``,
etc.
* :class:`litellm.llms.base_llm.chat.transformation.BaseLLMException` (status
429) — raised by some provider transports.

4
litellm/proxy/hooks/__init__.py
View file

@ -4,7 +4,6 @@ from typing import Final, Literal
from . import *
from .cache_control_check import _PROXY_CacheControlCheck
from .litellm_skills import SkillsInjectionHook
from .max_budget_limiter import _PROXY_MaxBudgetLimiter
from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler
from .max_iterations_limiter import _PROXY_MaxIterationsHandler
from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler
@ -18,7 +17,6 @@ from .sensitive_data_routing import _PROXY_SensitiveDataRoutingHandler
# transitively through `enterprise.enterprise_hooks` can resolve `PROXY_HOOKS`
# and `get_proxy_hook` from this partially-initialized module without circling.
PROXY_HOOKS: Final = {
"max_budget_limiter": _PROXY_MaxBudgetLimiter,
"parallel_request_limiter": _PROXY_MaxParallelRequestsHandler_v3,
"cache_control_check": _PROXY_CacheControlCheck,
"responses_id_security": ResponsesIDSecurity,
@ -35,7 +33,7 @@ if os.getenv("LEGACY_MULTI_INSTANCE_RATE_LIMITING", "false").lower() == "true":
def get_proxy_hook(
hook_name: Literal["max_budget_limiter", "managed_files", "parallel_request_limiter", "cache_control_check"] | str,
hook_name: Literal["managed_files", "parallel_request_limiter", "cache_control_check"] | str,
):
"""
Factory method to get a proxy hook instance by name

84
litellm/proxy/hooks/max_budget_limiter.py
View file

@ -1,84 +0,0 @@
from typing import Final
from fastapi import HTTPException
from litellm import verbose_logger
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.exceptions import RateLimitType
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
class _PROXY_MaxBudgetLimiter(CustomLogger):
# Class variables or attributes
def __init__(self):
pass
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
try:
verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook")
max_budget: Final = user_api_key_dict.user_max_budget
user_id: Final = user_api_key_dict.user_id
if max_budget is None or user_id is None:
return
from litellm.proxy.proxy_server import general_settings
if (
user_api_key_dict.team_id is not None
and general_settings.get("apply_user_budget_to_team_keys") is not True
):
return
# The reservation path admits at the strict-`<` boundary and
# atomically pre-fills the same counter we'd read here. Re-checking
# with `>=` would reject a request the reservation already admitted
# when the reservation fills the counter to exactly max_budget.
# Imported lazily to avoid a circular import via proxy.utils.
from litellm.proxy.spend_tracking.budget_reservation import (
get_reserved_counter_keys,
)
user_counter_key: Final = f"spend:user:{user_id}"
if user_counter_key in get_reserved_counter_keys(user_api_key_dict.budget_reservation):
return
from litellm.proxy.proxy_server import get_current_spend
curr_spend: Final = await get_current_spend(
counter_key=user_counter_key,
fallback_spend=user_api_key_dict.user_spend or 0.0,
)
verbose_proxy_logger.debug(
"MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f",
user_id,
curr_spend,
max_budget,
)
# CHECK IF REQUEST ALLOWED
if curr_spend >= max_budget:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None)
raise ProxyRateLimitError(
detail="Max budget limit reached.",
rate_limit_type=RateLimitType.BUDGET,
model=resolved_model,
llm_provider=llm_provider,
)
except HTTPException as e:
raise e
except Exception as e:
verbose_logger.exception(
"litellm.proxy.hooks.max_budget_limiter.py::async_pre_call_hook(): Exception occured - %s", e
)

4
litellm/proxy/utils.py
View file

@ -164,7 +164,6 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai
)
from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.parallel_request_limiter import (
_PROXY_MaxParallelRequestsHandler,
)
@ -982,7 +981,6 @@ class ProxyLogging:
dual_cache=DualCache(default_in_memory_ttl=1) # ping redis cache every 1s
)
self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache)
self.max_budget_limiter = _PROXY_MaxBudgetLimiter()
self.cache_control_check = _PROXY_CacheControlCheck()
self.alerting: list[str] | None = None
self.alerting_threshold: float = 300 # default to 5 min. threshold
@ -3580,7 +3578,7 @@ class ProxyLogging:
caps: Final = ProxyLogging._callback_capabilities()
post_call_pipelines: Final = _streamable_post_call_pipelines(request_data, user_api_key_dict)
# Fast path: no real overrides. Internal proxy CustomLogger callbacks
# (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default
# (e.g. _PROXY_CacheControlCheck, ManagedFiles) inherit the default
# ``async for chunk: yield chunk`` body, so wrapping the iterator
# through each of them adds N pass-through trampolines per chunk for
# zero behavior change. Skip the chain entirely and stream through.

2
tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md
View file

@ -21,7 +21,7 @@ on the shared lifecycle (every entity it creates is deleted on teardown).
| Entity | Unit | Pre-existing live | This suite (live) | Status |
|--------|------|-------------------|-------------------|--------|
| API key | `test_budget_reservation.py`, `test_max_budget_limiter.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** |
| API key | `test_budget_reservation.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** |
| Team | `test_team_budget_limits.py` | `otel_tests` | (org test builds a team) | **covered** |
| Internal user | auth unit tests | - | `test_internal_user_budget_blocks` | **covered (new)** |
| Team member | `test_team_member_budget.py` | - | `test_team_member_budget_blocks` | **covered (new)** |

12
tests/logging_callback_tests/test_unit_test_litellm_logging.py
View file

@ -8,8 +8,8 @@ from typing import Literal
import pytest
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler
from litellm._service_logger import ServiceLogging
import asyncio
@ -58,11 +58,11 @@ def test_is_internal_litellm_proxy_callback():
"""
Ensure we can determine if a callback is an internal litellm proxy callback
eg. `_PROXY_MaxBudgetLimiter`, `_PROXY_CacheControlCheck`
eg. `_PROXY_MaxIterationsHandler`, `_PROXY_CacheControlCheck`
"""
logging = setup_logging()
assert logging._is_internal_litellm_proxy_callback(_PROXY_MaxBudgetLimiter) == True
assert logging._is_internal_litellm_proxy_callback(_PROXY_MaxIterationsHandler) == True
# Test non-internal callbacks
def regular_callback():
@ -95,7 +95,7 @@ def test_should_run_sync_callbacks_for_async_calls():
assert logging._should_run_sync_callbacks_for_async_calls() == True
# Test with internal callback only
litellm.success_callback = [_PROXY_MaxBudgetLimiter]
litellm.success_callback = [_PROXY_MaxIterationsHandler]
assert logging._should_run_sync_callbacks_for_async_calls() == False
@ -107,7 +107,7 @@ def test_remove_internal_litellm_callbacks():
callbacks = [
regular_callback,
_PROXY_MaxBudgetLimiter,
_PROXY_MaxIterationsHandler,
_PROXY_CacheControlCheck,
"string_callback",
]
@ -116,5 +116,5 @@ def test_remove_internal_litellm_callbacks():
assert len(filtered) == 2 # Should only keep regular_callback and string_callback
assert regular_callback in filtered
assert "string_callback" in filtered
assert _PROXY_MaxBudgetLimiter not in filtered
assert _PROXY_MaxIterationsHandler not in filtered
assert _PROXY_CacheControlCheck not in filtered

59
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -5605,6 +5605,65 @@ async def test_common_checks_personal_user_budget_blocks_in_gather():
assert "User=u1" in str(over.value)
async def _common_checks_for_over_budget_personal_key(*, model: str) -> bool:
from litellm import Router
from litellm.proxy.auth.auth_checks import _is_model_cost_zero, common_checks
llm_router: Final = Router(
model_list=[
{
"model_name": "free-model",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
"model_info": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0},
},
{
"model_name": "paid-model",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
},
]
)
user: Final = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=1.0)
token: Final = UserAPIKeyAuth(token="k1", user_id="u1")
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
return 5.0 if counter_key == "spend:user:u1" else 0.0
proxy_logging_obj: Final = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
with (
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter),
):
result: Final = await common_checks(
request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]},
team_object=None,
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route="/chat/completions",
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=token,
request=MagicMock(spec=Request),
skip_budget_checks=_is_model_cost_zero(model=model, llm_router=llm_router),
)
await asyncio.sleep(0)
return result
@pytest.mark.asyncio
async def test_common_checks_over_budget_user_can_still_call_zero_cost_model():
"""LIT-7464: an exhausted personal budget must not block a model priced at 0/0,
while the same user is still rejected on a priced model."""
assert await _common_checks_for_over_budget_personal_key(model="free-model") is True
with pytest.raises(litellm.BudgetExceededError) as over:
await _common_checks_for_over_budget_personal_key(model="paid-model")
assert "ExceededBudget: User=u1" in str(over.value)
async def _run_internal_user_budget_alert(
*,
spend: float,

237
tests/test_litellm/proxy/hooks/test_max_budget_limiter.py
View file

@ -1,237 +0,0 @@
"""
Unit tests for the personal-budget pre-call hook.
The reservation path (added in PR #26845) atomically pre-fills the same
`spend:user:{user_id}` counter this hook reads, admitting at a strict-`<`
boundary. Re-checking with `>=` after reservation would reject requests the
reservation already admitted when the reservation fills the counter to
exactly `max_budget` (e.g. requests with no `max_tokens` cap fall back to
reserving the smallest remaining headroom).
These tests pin the skip-when-reserved behavior and guard against drift.
"""
from unittest.mock import AsyncMock, patch
import pytest
from fastapi import HTTPException
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
def _make_user_api_key_auth(
user_id: str = "user-1",
user_max_budget: float = 10.0,
user_spend: float = 0.0,
team_id=None,
budget_reservation=None,
) -> UserAPIKeyAuth:
return UserAPIKeyAuth(
api_key="sk-test",
user_id=user_id,
user_max_budget=user_max_budget,
user_spend=user_spend,
team_id=team_id,
budget_reservation=budget_reservation,
)
@pytest.mark.asyncio
async def test_under_budget_passes():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=3.0),
):
result = await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert result is None
@pytest.mark.asyncio
async def test_over_budget_rejects_without_reservation():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.status_code == 429
assert "Max budget limit reached." in exc_info.value.detail
@pytest.mark.asyncio
async def test_skips_when_user_counter_is_reserved():
"""
Reservation atomically pre-fills `spend:user:{user_id}` and admits the
request. The legacy `>=` check must not double-enforce on the same
counter — that's what produced the boundary regression where a fresh
user with no `max_tokens` cap got 429'd on their first request.
"""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_id="user-1",
user_max_budget=10.0,
budget_reservation={
"reserved_cost": 10.0,
"entries": [
{
"counter_key": "spend:user:user-1",
"entity_type": "User",
"entity_id": "user-1",
"reserved_cost": 10.0,
"applied_adjustment": 0.0,
}
],
"finalized": False,
},
)
# `get_current_spend` would return 10.0 here (counter pre-filled by the
# reservation). The hook must skip without reading it.
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
) as mock_get_spend:
result = await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert result is None
mock_get_spend.assert_not_awaited()
@pytest.mark.asyncio
async def test_does_not_skip_when_reservation_covers_a_different_counter():
"""
A reservation that only covers e.g. `spend:team:{team_id}` (not the user
counter) must not exempt the user-budget check.
"""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_id="user-1",
user_max_budget=10.0,
budget_reservation={
"reserved_cost": 5.0,
"entries": [
{
"counter_key": "spend:team:team-x",
"entity_type": "Team",
"entity_id": "team-x",
"reserved_cost": 5.0,
"applied_adjustment": 0.0,
}
],
"finalized": False,
},
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.status_code == 429
@pytest.mark.asyncio
async def test_team_keys_skip_personal_budget():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_max_budget=10.0,
team_id="team-1",
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=999.0),
) as mock_get_spend:
result = await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert result is None
mock_get_spend.assert_not_awaited()
@pytest.mark.asyncio
async def test_team_keys_enforce_personal_budget_when_flag_enabled():
"""This hook is the third personal-budget gate alongside common_checks and the
reservation path, so apply_user_budget_to_team_keys has to reach it too or an
opted-in deployment enforces in two places out of three."""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_max_budget=10.0,
team_id="team-1",
)
with patch.dict(
"litellm.proxy.proxy_server.general_settings",
{"apply_user_budget_to_team_keys": True},
), patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=999.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.status_code == 429
@pytest.mark.asyncio
async def test_no_max_budget_passes():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-test",
user_id="user-1",
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=999.0),
) as mock_get_spend:
result = await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert result is None
mock_get_spend.assert_not_awaited()

61
tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py
View file

@ -6,7 +6,7 @@ Background
----------
The proxy's internal rate-limit hooks (parallel_request_limiter,
parallel_request_limiter_v3, dynamic_rate_limiter, dynamic_rate_limiter_v3,
batch_rate_limiter, max_budget_limiter, max_iterations_limiter,
batch_rate_limiter, max_iterations_limiter,
max_budget_per_session_limiter) all fire from ``async_pre_call_hook`` —
*before* :func:`litellm.get_llm_provider` runs anywhere else in the request
lifecycle.
@ -50,7 +50,6 @@ from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHand
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
_PROXY_DynamicRateLimitHandlerV3,
)
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.max_budget_per_session_limiter import (
_PROXY_MaxBudgetPerSessionHandler,
)
@ -830,64 +829,6 @@ async def test_batch_rate_limiter_unknown_model_falls_back():
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
# ---------------------------------------------------------------------------
# max_budget_limiter
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_max_budget_limiter_populates_provider():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-budget",
user_id="user-1",
user_max_budget=10.0,
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={"model": "gpt-4o-mini"},
call_type="completion",
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == "openai"
assert exc.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_max_budget_limiter_no_model_falls_back():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-budget",
user_id="user-1",
user_max_budget=10.0,
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
assert exc_info.value.model == ""
# ---------------------------------------------------------------------------
# max_iterations_limiter
# ---------------------------------------------------------------------------

4
tests/test_litellm/proxy/proxy_server/test_routes_config.py
View file

@ -1397,7 +1397,7 @@ def test_get_config_callbacks_excludes_internal_runtime_callbacks(client, auth_a
from litellm.integrations.s3_v2 import S3Logger
from litellm.integrations.sqs import SQSLogger
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import VectorStorePreCallHook
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
from litellm.router import Router
class _InventoryTestGuardrail(CustomGuardrail):
@ -1425,7 +1425,7 @@ def test_get_config_callbacks_excludes_internal_runtime_callbacks(client, auth_a
litellm,
"callbacks",
[
_PROXY_MaxBudgetLimiter(),
_PROXY_CacheControlCheck(),
_PROXY_LiteLLMManagedFiles(internal_usage_cache=MagicMock(), prisma_client=MagicMock()),
ServiceLogging(),
VectorStorePreCallHook(),

2
tests/test_litellm/proxy/test_proxy_logging_hook_detection.py
View file

@ -73,7 +73,7 @@ async def test_post_call_response_headers_hook_returns_early_without_callbacks(
def test_callback_capabilities_skips_default_custom_logger(monkeypatch):
"""
Internal proxy hooks (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit
Internal proxy hooks (e.g. _PROXY_CacheControlCheck, ManagedFiles) inherit
the default ``async_post_call_streaming_iterator_hook`` body. The
capability scanner must NOT report them as iterator overrides — wrapping
the chunk stream through every no-op layer was responsible for ~10x

10
tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py
View file

@ -220,7 +220,7 @@ def test_add_proxy_hooks_registers_callbacks(proxy_logging, monkeypatch):
what gets registered. Verifies that the resulting instances land in
``proxy_logging.proxy_hook_mapping`` keyed by hook name.
"""
hook_keys = ["cache_control_check", "max_budget_limiter"]
hook_keys = ["cache_control_check", "max_iterations_limiter"]
registered: List[Any] = []
from litellm.proxy import utils as utils_mod
@ -362,22 +362,22 @@ def test_add_proxy_hooks_unknown_hook_raises(proxy_logging, monkeypatch):
def test_get_proxy_hook_returns_registered_instance(proxy_logging):
s_cache = MagicMock()
s_budget = MagicMock()
s_iterations = MagicMock()
s_parallel = MagicMock()
proxy_logging.proxy_hook_mapping = {
"cache_control_check": s_cache,
"max_budget_limiter": s_budget,
"max_iterations_limiter": s_iterations,
"max_parallel_request_limiter": s_parallel,
}
snapshot = {
"cache_control_check": proxy_logging.get_proxy_hook("cache_control_check") is s_cache,
"max_budget_limiter": proxy_logging.get_proxy_hook("max_budget_limiter") is s_budget,
"max_iterations_limiter": proxy_logging.get_proxy_hook("max_iterations_limiter") is s_iterations,
"max_parallel_request_limiter": proxy_logging.get_proxy_hook("max_parallel_request_limiter") is s_parallel,
"unknown_returns_none": proxy_logging.get_proxy_hook("unknown") is None,
}
assert snapshot == {
"cache_control_check": True,
"max_budget_limiter": True,
"max_iterations_limiter": True,
"max_parallel_request_limiter": True,
"unknown_returns_none": True,
}

34
tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py
View file

@ -3,7 +3,7 @@
from __future__ import annotations
from typing import Any, Dict
from unittest.mock import AsyncMock, MagicMock
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
@ -400,6 +400,37 @@ def test_has_pre_call_guardrails_counts_a_content_enforcer(proxy_logging, monkey
assert proxy_logging.has_pre_call_guardrails({}) is True
@pytest.mark.asyncio
async def test_registered_hooks_do_not_enforce_user_budget(proxy_logging, monkeypatch):
"""
Personal budget is auth's job (`_user_max_budget_check`), which exempts
zero-cost models. A hook re-checking the same counter without that
exemption is what 429'd free models once a user was over budget.
"""
monkeypatch.setattr(litellm, "callbacks", [])
with patch("litellm.proxy.proxy_server.prisma_client", None):
proxy_logging._add_proxy_hooks(llm_router=None)
ProxyLogging._callback_capabilities_cache.clear()
over_budget_user = UserAPIKeyAuth(
api_key="sk-personal",
user_id="user-over-budget",
user_max_budget=1.0,
user_spend=5.0,
team_id=None,
)
data = {"model": "free-model", "messages": [{"role": "user", "content": "hi"}]}
with patch("litellm.proxy.proxy_server.get_current_spend", new=AsyncMock(return_value=5.0)):
out = await proxy_logging.pre_call_hook(
user_api_key_dict=over_budget_user,
data=data,
call_type="completion",
)
assert out == data
def test_every_pre_call_customlogger_is_deliberately_classified():
"""
A ledger, so a new hook cannot land unclassified.
@ -415,7 +446,6 @@ def test_every_pre_call_customlogger_is_deliberately_classified():
"_ENTERPRISE_BlockedUserList",
}
counts_or_shapes_the_request = {
"_PROXY_MaxBudgetLimiter",
"_PROXY_MaxParallelRequestsHandler_v3",
"_PROXY_MaxIterationsHandler",
"_PROXY_MaxBudgetPerSessionHandler",

69
tests/test_litellm/test_rate_limit_error_unification.py
View file

@ -221,28 +221,6 @@ class TestProxyHookCategoryWiring:
"""End-to-end check that every proxy-side rate limiter raises the unified
class with a sensible category, not a bare HTTPException."""
def test_max_budget_limiter_raises_proxy_rate_limit_error(self):
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
limiter = _PROXY_MaxBudgetLimiter()
# The simplest deterministic path: directly raise from the conditional
# branch by calling into the helper's exception construction. We
# round-trip through the public class to assert the shape.
with pytest.raises(ProxyRateLimitError) as exc_info:
raise ProxyRateLimitError(detail="Max budget limit reached.")
assert exc_info.value.status_code == 429
assert exc_info.value.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT
# And it's also a RateLimitError + HTTPException (the unification).
assert isinstance(exc_info.value, RateLimitError)
assert isinstance(exc_info.value, HTTPException)
# Static check that the limiter's module imports the unified class so
# the source of truth is wired correctly.
from litellm.proxy.hooks import max_budget_limiter
assert hasattr(max_budget_limiter, "ProxyRateLimitError")
assert max_budget_limiter.ProxyRateLimitError is ProxyRateLimitError
del limiter # silence unused-var
@pytest.mark.parametrize(
"module_path",
[
@ -251,7 +229,6 @@ class TestProxyHookCategoryWiring:
"litellm.proxy.hooks.dynamic_rate_limiter",
"litellm.proxy.hooks.dynamic_rate_limiter_v3",
"litellm.proxy.hooks.batch_rate_limiter",
"litellm.proxy.hooks.max_budget_limiter",
"litellm.proxy.hooks.max_budget_per_session_limiter",
"litellm.proxy.hooks.max_iterations_limiter",
],
@ -542,44 +519,6 @@ class TestProxyHooksActuallyRaiseProxyRateLimitError:
assert isinstance(e, RateLimitError)
assert isinstance(e, HTTPException)
@pytest.mark.asyncio
async def test_max_budget_limiter_raises_proxy_rate_limit_error(self):
"""
Drive `_PROXY_MaxBudgetLimiter` past the user budget and assert it
raises the unified class. Mocks `get_current_spend` so we don't need
the proxy DB.
"""
from unittest.mock import patch
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.max_budget_limiter import (
_PROXY_MaxBudgetLimiter,
)
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-test-budget",
user_id="user-budget-1",
user_max_budget=1.0,
user_spend=2.0,
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
return_value=5.0,
):
with pytest.raises(ProxyRateLimitError) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
e = exc_info.value
assert e.status_code == 429
assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT
assert "max budget" in str(e.detail).lower()
@pytest.mark.asyncio
async def test_dynamic_rate_limiter_v1_raises_proxy_rate_limit_error(self):
"""
@ -1156,14 +1095,6 @@ class TestProxyHooksWireTypeCorrectly:
max-iterations) without grepping the error message.
"""
def test_max_budget_limiter_emits_budget_type(self):
e = ProxyRateLimitError(
detail="Max budget limit reached.",
rate_limit_type=RateLimitType.BUDGET,
)
assert e.category == "litellm_rate_limit"
assert e.rate_limit_type == "budget"
def test_max_iterations_limiter_emits_max_iterations_type(self):
e = ProxyRateLimitError(
detail="Max iterations exceeded for session abc.",