scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(proxy): preserve discovery after budget exhaustion (#48809) (#49463)

Browse source 
* fix(proxy): preserve discovery after budget exhaustion (#48809)

* refactor(proxy): decide budget exemption once (#48809)

* refactor(proxy): centralize request budget policy (#48809)

* fix(proxy): keep custom pass-through budget checks (#48809)
This commit is contained in:
ishaan-j-berri 2026-10-31 12:00:51 -07:00 • committed by GitHub
parent 6eddc4a979
commit 941f52f26f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 213 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/proxy/auth/auth_checks.py
View file

@ -1086,13 +1086,31 @@ def model_max_budget_of(holder: ModelMaxBudgetHolder) -> Mapping[str, object] |
def route_skips_budget_checks(route: str) -> bool:
return route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES and (
route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route)
if route in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES:
return False
if route in MODEL_DISCOVERY_ROUTES:
return True
if RouteChecks.is_llm_api_route(route=route):
return False
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass-through endpoints import the auth stack
InitPassThroughEndpointHelpers,
)
return not InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route)
def request_skips_budget_checks(route: str, model: str | list[str] | None, llm_router: Router | None) -> bool:
return route_skips_budget_checks(route=route) or is_model_cost_zero(model=model, llm_router=llm_router)
def request_skips_budget_checks(
route: str,
model: str | list[str] | None,
llm_router: Router | None,
*,
request_body: Mapping[str, object] | None = None,
) -> bool:
if route_skips_budget_checks(route=route):
return True
if request_body is not None and is_mcp_discovery_request(route=route, request_body=request_body):
return True
return is_model_cost_zero(model=model, llm_router=llm_router)
async def common_checks(
@ -1142,10 +1160,11 @@ async def common_checks(
team_id=valid_token.team_id if valid_token is not None else None,
)
skip_all_budget_checks: Final = (
skip_budget_checks
or route_skips_budget_checks(route=route)
or is_mcp_discovery_request(route=route, request_body=_typed_request_body(request_body))
skip_all_budget_checks: Final = skip_budget_checks or request_skips_budget_checks(
route=route,
model=_model,
llm_router=llm_router,
request_body=_typed_request_body(request_body),
)
membership_user_id: Final = (

32
litellm/proxy/auth/user_api_key_auth.py
View file

@ -80,7 +80,6 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
get_team_object,
get_user_object,
get_user_role,
is_mcp_discovery_request,
is_model_cost_zero,
is_user_proxy_admin,
is_valid_fallback_model,
@ -88,6 +87,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
key_access_group_grants_model,
key_model_aliases_for_auth_check,
model_max_budget_of,
request_skips_budget_checks,
resolve_and_validate_end_user_id,
resolve_default_end_user_budget,
route_skips_budget_checks,
@ -2497,15 +2497,9 @@ async def user_api_key_auth_builder(
llm_router=llm_router,
team_id=valid_token.team_id,
)
skip_budget_checks = is_mcp_discovery_request(route=route, request_body=request_data)
if not skip_budget_checks and model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import is_model_cost_zero
skip_budget_checks = ( # rebind-ok: pre-existing rebinding on a rename-only line
is_model_cost_zero(model=model, llm_router=llm_router)
)
if skip_budget_checks:
verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model)
skip_budget_checks = request_skips_budget_checks(
route=route, model=model, llm_router=llm_router, request_body=request_data
)
# Fetch project object for JWT path if project_id is set
_jwt_project_obj = None
@ -3019,15 +3013,9 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
llm_router=llm_router,
team_id=valid_token.team_id,
)
skip_budget_checks = is_mcp_discovery_request(route=route, request_body=request_data)
if not skip_budget_checks and model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import is_model_cost_zero
skip_budget_checks = is_model_cost_zero( # rebind-ok: pre-existing rebinding on a rename-only line
model=model, llm_router=llm_router
)
if skip_budget_checks:
verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model)
skip_budget_checks: Final = request_skips_budget_checks(
route=route, model=model, llm_router=llm_router, request_body=request_data
)
# Check 3. Check if user is in their team budget
if not skip_budget_checks and valid_token.team_member_spend is not None:
@ -4088,11 +4076,7 @@ def _should_skip_budget_checks(
llm_router=llm_router,
team_id=team_id,
)
if is_mcp_discovery_request(route=route, request_body=request_data):
return True
if model is not None and llm_router is not None:
return is_model_cost_zero(model=model, llm_router=llm_router)
return False
return request_skips_budget_checks(route=route, model=model, llm_router=llm_router, request_body=request_data)
def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> Principal:

44
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -11202,9 +11202,47 @@ def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None:
assert route_skips_budget_checks(route="/v1/chat/completions") is False
def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() -> None:
assert request_skips_budget_checks(route="/v1/models", model=None, llm_router=None) is True
assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False
@pytest.mark.parametrize(
"route, request_body, model, expected",
[
("/v1/models", None, None, True),
("/v1/models", None, "paid", True),
("/v1/chat/completions", None, None, False),
("/v1/chat/completions", None, "free", True),
("/v1/chat/completions", None, "paid", False),
("/mcp", {"method": "tools/list"}, "paid", True),
("/mcp", {"method": "tools/call"}, "paid", False),
("/mcp", None, "paid", False),
("/mcp/tools/call", {"method": "initialize"}, "paid", False),
("/health/services", None, "paid", False),
],
)
def test_request_skips_budget_checks_for_non_spending_requests(
route: str, request_body: Mapping[str, object] | None, model: str | None, expected: bool
) -> None:
router: Final = litellm.Router(
model_list=[
{
"model_name": "free",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"input_cost_per_token": 0,
"output_cost_per_token": 0,
},
},
{
"model_name": "paid",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"input_cost_per_token": 1,
"output_cost_per_token": 1,
},
},
]
)
assert (
request_skips_budget_checks(route=route, model=model, llm_router=router, request_body=request_body) is expected
)
def _agent_model_ceiling_resolver(

141
tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py
View file

@ -5,7 +5,7 @@ import os
import subprocess
import sys
import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from contextlib import contextmanager, nullcontext
from datetime import datetime, timedelta, timezone
from functools import partial
@ -17,6 +17,7 @@ from typing import Final
from unittest.mock import ANY, AsyncMock, MagicMock, patch
import fastapi
import orjson
import pytest
from fastapi import HTTPException, Request, status
@ -24,6 +25,8 @@ from starlette.datastructures import UploadFile
import litellm
import litellm.proxy.proxy_server
from litellm.proxy import utils as proxy_utils
from litellm.proxy.common_utils import user_api_key_cache as key_cache_utils
from litellm.caching.dual_cache import DualCache
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.proxy.common_utils.user_api_key_cache import (
@ -85,6 +88,7 @@ from litellm.proxy.auth.user_api_key_auth import (
user_api_key_auth,
user_api_key_auth_websocket_for_model,
)
from litellm.proxy.pass_through_endpoints import pass_through_endpoints as pass_through
from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata
from tests.unit.proxy.db.fake_prisma_engine import engine_call
from tests._master_key import MASTER_KEY
@ -11173,6 +11177,30 @@ async def test_temp_budget_increase_applied_for_cached_key():
assert cached_after.max_budget == 2.0
BUDGETED_PASS_THROUGH_ROUTE: Final = "/custom-provider/infer"
@pytest.fixture
def registered_budgeted_route() -> Iterator[None]:
pass_through.InitPassThroughEndpointHelpers.add_exact_path_route(
app=fastapi.FastAPI(),
path=BUDGETED_PASS_THROUGH_ROUTE,
target="https://provider.example/infer",
custom_headers=None,
forward_headers=None,
merge_query_params=None,
dependencies=None,
cost_per_request=1.0,
endpoint_id="budget-regression",
auth=True,
methods=["POST"],
)
try:
yield
finally:
pass_through.InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id="budget-regression")
@pytest.mark.asyncio
@pytest.mark.parametrize(
"team_member_spend, expect_blocked",
@ -11182,7 +11210,11 @@ async def test_temp_budget_increase_applied_for_cached_key():
(2.39, False),
],
)
async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spend, expect_blocked):
@pytest.mark.usefixtures("registered_budgeted_route")
@pytest.mark.parametrize(
"route", ["/v1/messages", "/health/services", "/health/test_connection", BUDGETED_PASS_THROUGH_ROUTE]
)
async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spend, expect_blocked, route):
"""A team member counter sitting exactly at the cap (where a resized reservation
lands it) must be rejected by the cached-key auth path like every other budget check."""
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
@ -11227,8 +11259,8 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe
)
mock_request = MagicMock()
mock_request.url.path = "/v1/messages"
mock_request.method = "POST"
mock_request.url.path = route
mock_request.method = "POST" if route == "/v1/messages" else "GET"
mock_request.headers = {"authorization": f"Bearer {api_key}"}
mock_request.query_params = {}
mock_request.state = SimpleNamespace()
@ -11276,6 +11308,107 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe
assert f"TeamMember={user_id}:{team_id}" in exc_info.value.message
@pytest.mark.asyncio
@pytest.mark.parametrize("key_max_budget, member_max_budget", [(None, 1.0), (1.0, None)])
@pytest.mark.parametrize(
"route, expect_blocked",
[
("/models", False),
("/v1/models", False),
("/utils/model_info", False),
("/model/info", False),
("/key/info", False),
("/user/info", False),
("/v1/chat/completions", True),
("/v1/messages", True),
],
)
async def test_cached_key_budget_preserves_non_spending_access(
key_max_budget: float | None,
member_max_budget: float | None,
route: str,
expect_blocked: bool,
monkeypatch: pytest.MonkeyPatch,
) -> None:
api_key: Final = "sk-budget-route-fixture"
hashed_token: Final = proxy_utils.hash_token(api_key)
team_id: Final = "budget-route-team"
user_id: Final = "budget-route-user"
cache: Final = DualCache()
await cache_key_object(
hashed_token=hashed_token,
user_api_key_obj=UserAPIKeyAuth(
token=hashed_token,
user_id=user_id,
team_id=team_id,
team_member_spend=2.0,
spend=2.0,
max_budget=key_max_budget,
),
user_api_key_cache=cache,
proxy_logging_obj=None,
)
await cache.async_set_cache(key=f"team_id:{team_id}", value=LiteLLM_TeamTableCachedObj(team_id=team_id))
await cache.async_set_cache(
key=user_id,
value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER),
)
await cache.async_set_cache(
key=key_cache_utils.team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
value=LiteLLM_TeamMembership(
user_id=user_id,
team_id=team_id,
spend=2.0,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=member_max_budget),
),
)
await cache.async_set_cache(key=f"spend:team_member:{user_id}:{team_id}", value=2.0)
await cache.async_set_cache(key=f"spend:key:{hashed_token}", value=2.0)
request: Final = Request(
scope={
"type": "http",
"path": route,
"method": ("POST" if route in ("/v1/chat/completions", "/v1/messages") else "GET"),
"headers": [(b"authorization", f"Bearer {api_key}".encode())],
"query_string": b"",
}
)
proxy_logging: Final = MagicMock()
proxy_logging.budget_alerts = AsyncMock()
proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
monkeypatch.setattr(
litellm.proxy.proxy_server,
"general_settings",
{"disable_budget_reservation": True},
)
monkeypatch.setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(litellm.proxy.proxy_server, "spend_counter_cache", cache)
monkeypatch.setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging)
async def authenticate() -> UserAPIKeyAuth:
return await user_api_key_auth_builder(
request=request,
api_key=f"Bearer {api_key}",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={},
)
if expect_blocked:
with pytest.raises(ProxyException) as exc_info:
await authenticate()
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
return
result: Final = await authenticate()
assert result.user_id == user_id
assert result.team_id == team_id
@pytest.mark.asyncio
@pytest.mark.parametrize(
"expiry_offset, expect_blocked",