fix(proxy): reserve hook-added tag budgets and honor auth's check scope

Reserve the estimated request cost against the tags a pre-call hook added,
folded into the request's reservation, so a burst of requests cannot all
pass the read check on the same stale spend the way body tags already
cannot. Skip the post-hook check where the auth wrapper runs no
common_checks (no-auth dev mode, custom auth without
custom_auth_run_common_checks), sharing that predicate with auth.
This commit is contained in:
Yucheng He 2026-09-15 03:00:20 -07:00
parent 0333ca1f34
commit 27bc2a80e2
7 changed files with 428 additions and 23 deletions

View file

@ -860,6 +860,27 @@ def route_skips_budget_checks(route: str) -> bool:
)
_AUTHN_FLAGS: Final = ("enable_jwt_auth", "enable_oauth2_auth", "enable_oauth2_proxy_auth")
def auth_skips_common_checks(
general_settings: Mapping[str, object], master_key: str | None, custom_auth_configured: bool
) -> bool:
"""
Whether ``user_api_key_auth`` runs no ``common_checks`` at all for this deployment.
That is the case in no-auth dev mode (no master key and no JWT or OAuth2
auth configured, so the proxy is unauthenticated by configuration) and behind
a custom auth hook that did not opt in with ``custom_auth_run_common_checks``.
Post-auth checks that mirror ``common_checks`` skip themselves on the same terms.
"""
no_auth_mode: Final = master_key is None and not any(general_settings.get(flag, False) for flag in _AUTHN_FLAGS)
custom_auth_opted_out: Final = custom_auth_configured and not general_settings.get(
"custom_auth_run_common_checks", False
)
return no_auth_mode or custom_auth_opted_out
async def common_checks(
request_body: dict,
team_object: LiteLLM_TeamTable | None,

View file

@ -50,6 +50,7 @@ from litellm.proxy.auth.auth_checks import (
_virtual_key_max_budget_alert_check,
_virtual_key_max_budget_check,
_virtual_key_soft_budget_check,
auth_skips_common_checks,
can_key_call_model,
common_checks,
get_end_user_object,
@ -2539,22 +2540,11 @@ async def _run_centralized_common_checks(
if isinstance(endpoint, dict) and endpoint.get("path", "") == route and endpoint.get("auth") is not True:
return
# No-auth dev mode: master_key unset AND no JWT/OAuth2 auth
# configured. The builder returns an INTERNAL_USER token for any
# api_key; the proxy is unauthenticated by configuration.
# Running common_checks would block every admin route on these
# deployments where that was previously not the contract. If any
# authn is enabled (JWT, OAuth2, OAuth2-proxy), authz must run.
if master_key is None and not (
general_settings.get("enable_jwt_auth", False)
or general_settings.get("enable_oauth2_auth", False)
or general_settings.get("enable_oauth2_proxy_auth", False)
if auth_skips_common_checks(
general_settings=general_settings, master_key=master_key, custom_auth_configured=user_custom_auth is not None
):
return
if user_custom_auth is not None and not general_settings.get("custom_auth_run_common_checks", False):
return
parent_otel_span: Final = user_api_key_auth_obj.parent_otel_span
# In the integrated auth flow ``_user_api_key_auth_builder`` has already
# resolved the end-user id and attached it here. Reuse that to avoid a

View file

@ -51,6 +51,7 @@ from litellm.litellm_core_utils.streaming_handler import (
)
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
auth_skips_common_checks,
can_key_call_resolved_model,
route_skips_budget_checks,
tag_max_budget_check_for_tags,
@ -76,6 +77,7 @@ from litellm.proxy.common_utils.sse_keepalive import (
from litellm.proxy.dd_span_tagger import DDSpanTagger
from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_added_tags
from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails
from litellm.router import Router
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
@ -679,18 +681,23 @@ async def _enforce_tag_budgets_for_added_tags(
route: str,
llm_router: Router | None,
proxy_logging_obj: ProxyLogging,
) -> None:
"""Budget-check the tags that ``pre_call_hook`` added to the request.
general_settings: Mapping[str, object],
) -> tuple[str, ...]:
"""Budget-check the tags that ``pre_call_hook`` added to the request and return them.
Tag budgets are enforced in auth against the tags in the request body, and
guardrails run after auth, so a tag a guardrail sets is only checked here,
on the same routes auth checks.
on the same routes auth checks and only when auth ran its checks at all.
"""
added_tags: Final = tuple(tag for tag in _tags_on_request(data) if tag not in tags_before_pre_call_hook)
if not added_tags or route_skips_budget_checks(route=route):
return
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
return ()
from litellm.proxy.proxy_server import master_key, prisma_client, user_api_key_cache, user_custom_auth
if auth_skips_common_checks(
general_settings=general_settings, master_key=master_key, custom_auth_configured=user_custom_auth is not None
):
return ()
await tag_max_budget_check_for_tags(
tags=added_tags,
prisma_client=prisma_client,
@ -699,6 +706,7 @@ async def _enforce_tag_budgets_for_added_tags(
model=_request_model(data),
llm_router=llm_router,
)
return added_tags
async def _parse_event_data_for_error(event_line: str | bytes) -> int | None:
@ -2094,15 +2102,63 @@ class ProxyBaseLLMRequestProcessing:
if "messages" in self.data and self.data["messages"]:
logging_obj.update_messages(self.data["messages"])
await _enforce_tag_budgets_for_added_tags(
request_route: Final = get_request_route(request=request)
added_tags: Final = await _enforce_tag_budgets_for_added_tags(
data=self.data,
tags_before_pre_call_hook=self._tags_before_pre_call_hook,
route=get_request_route(request=request),
route=request_route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
general_settings=general_settings,
)
if added_tags and general_settings.get("disable_budget_reservation") is not True:
await self._reserve_budget_for_added_tags(
added_tags=added_tags,
route=request_route,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
)
return self.data, logging_obj
async def _reserve_budget_for_added_tags(
self,
added_tags: Sequence[str],
route: str,
llm_router: Router | None,
user_api_key_dict: UserAPIKeyAuth,
proxy_logging_obj: ProxyLogging,
fail_closed_budget_enforcement: bool,
) -> None:
"""Reserve the added tags' budgets the way auth reserved the body tags, so a burst cannot overshoot them.
The entries join the request's reservation, on the auth object and in the
request metadata, so the success, failure and cancel paths settle them together.
"""
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
reservation: Final = await reserve_budget_for_added_tags(
tags=added_tags,
request_body=self.data,
route=route,
llm_router=llm_router,
valid_token=user_api_key_dict,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
fail_closed_budget_enforcement=fail_closed_budget_enforcement,
)
if reservation is None:
return
existing: Final = user_api_key_dict.budget_reservation
if existing is not None:
existing["entries"].extend(reservation["entries"])
return
user_api_key_dict.budget_reservation = reservation # rebind-ok: the failure and cancel paths read it here
_, metadata_bucket = get_or_create_metadata_bucket(self.data)
metadata_bucket["user_api_key_budget_reservation"] = reservation
async def _pre_call_with_fallbacks(
self,
request: Request,

View file

@ -235,7 +235,68 @@ async def reserve_budget_for_request(
)
if not counters:
return None
return await _reserve_counters(
counters=counters,
request_body=request_body,
route=route,
llm_router=llm_router,
valid_token=valid_token,
fail_closed_budget_enforcement=fail_closed_budget_enforcement,
raw_body=raw_body,
)
async def reserve_budget_for_added_tags(
tags: Sequence[str],
request_body: dict[str, object], # mutable-ok: the request payload the proxy threads through the pipeline
route: str,
llm_router: Router | None,
valid_token: UserAPIKeyAuth,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
fail_closed_budget_enforcement: bool = False,
) -> dict[str, object] | None: # mutable-ok: the reservation dict the settlement paths stamp in place
"""
Reserve the request's estimated cost against ``tags`` a pre-call hook added.
Auth reserved the body tags before the hook ran, so without this a burst of
requests all read the same spend for a hook-added tag and all get through.
Same route and model guards as ``reserve_budget_for_request``; the caller
folds the result into the request's reservation so one settlement covers both.
"""
if not RouteChecks.is_llm_api_route(route=route) or _is_unbilled_route(route):
return None
if get_model_from_request(request_body, route, llm_router=llm_router) is None:
return None
counters: Final = await _tag_budget_counters(
tag_names=_dedupe_tags(list(tags)),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if not counters:
return None
return await _reserve_counters(
counters=counters,
request_body=request_body,
route=route,
llm_router=llm_router,
valid_token=valid_token,
fail_closed_budget_enforcement=fail_closed_budget_enforcement,
raw_body=None,
)
async def _reserve_counters(
counters: Sequence[_BudgetCounter],
request_body: dict[str, object], # mutable-ok: the request payload the proxy threads through the pipeline
route: str,
llm_router: Router | None,
valid_token: UserAPIKeyAuth,
fail_closed_budget_enforcement: bool,
raw_body: bytes | None,
) -> dict[str, object] | None: # mutable-ok: the reservation dict the settlement paths stamp in place
input_token_counts: Final = await count_request_input_tokens(
request_body=request_body,
route=route,
@ -582,10 +643,24 @@ async def _get_tag_budget_counters(
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> list[_BudgetCounter]:
from litellm.proxy.auth.auth_checks import get_tag_objects_batch
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
tag_names: Final = _dedupe_tags(get_tags_from_request_body(request_body=request_body))
return await _tag_budget_counters(
tag_names=_dedupe_tags(get_tags_from_request_body(request_body=request_body)),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
async def _tag_budget_counters(
tag_names: Sequence[str],
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> list[_BudgetCounter]:
from litellm.proxy.auth.auth_checks import get_tag_objects_batch
if not tag_names:
return []

View file

@ -2481,6 +2481,33 @@ def test_route_skips_budget_checks_matches_auth_scope(route, expected):
assert route_skips_budget_checks(route=route) is expected
@pytest.mark.parametrize(
("general_settings", "master_key", "custom_auth_configured", "expected"),
[
({}, None, False, True),
({"enable_jwt_auth": True}, None, False, False),
({"enable_oauth2_auth": True}, None, False, False),
({"enable_oauth2_proxy_auth": True}, None, False, False),
({}, "sk-master", False, False),
({}, "sk-master", True, True),
({"custom_auth_run_common_checks": True}, "sk-master", True, False),
({"custom_auth_run_common_checks": False}, "sk-master", True, True),
],
)
def test_auth_skips_common_checks_names_the_deployments_that_never_run_them(
general_settings, master_key, custom_auth_configured, expected
):
"""No-auth dev mode and a custom auth hook without the opt-in run no common_checks, so no budget checks."""
from litellm.proxy.auth.auth_checks import auth_skips_common_checks
assert (
auth_skips_common_checks(
general_settings=general_settings, master_key=master_key, custom_auth_configured=custom_auth_configured
)
is expected
)
@pytest.mark.asyncio
async def test_get_team_object_raises_404_when_not_found():
from unittest.mock import AsyncMock, MagicMock

View file

@ -2,8 +2,9 @@ from __future__ import annotations
import json
import math
from types import MappingProxyType
from types import MappingProxyType, SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -15,6 +16,7 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.spend_tracking.budget_reservation import (
count_request_input_tokens,
estimate_request_max_cost,
reserve_budget_for_added_tags,
reserve_budget_for_request,
)
from litellm.proxy.utils import ProxyLogging
@ -136,6 +138,104 @@ async def test_repeated_token_counting_never_touches_a_tiny_budget(
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reserved_cost)
HOOK_TAG: Final = "hook-added-tag"
BODY_TAG: Final = "body-tag"
CHAT_BODY: Final[dict[str, object]] = {
"model": "gpt-4o",
"messages": ANTHROPIC_MESSAGES,
"max_tokens": 5,
"metadata": {"tags": [BODY_TAG]},
}
def _budgeted_tag_prisma(tag_names: tuple[str, ...], max_budget: float) -> MagicMock:
"""A tag table where every named tag carries ``max_budget`` and no spend yet."""
def _row(tag_name: str) -> MagicMock:
row = MagicMock()
row.tag_name = tag_name
row.dict = MagicMock(
return_value={
"tag_name": tag_name,
"spend": 0.0,
"models": [],
"litellm_budget_table": {"max_budget": max_budget},
}
)
return row
async def find_many(**kwargs: object) -> list[object]:
where = kwargs.get("where")
if not isinstance(where, dict):
return [SimpleNamespace(tag_name=name) for name in tag_names]
return [_row(name) for name in where["tag_name"]["in"] if name in tag_names]
prisma = MagicMock()
prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=find_many)
return prisma
async def _reserve_added_tags(
route: str, prisma: MagicMock, tags: tuple[str, ...] = (HOOK_TAG,)
) -> dict[str, object] | None:
return await reserve_budget_for_added_tags(
tags=tags,
request_body=dict(CHAT_BODY),
route=route,
llm_router=None,
valid_token=UserAPIKeyAuth(token="hashed-hook-tag-key", max_budget=100.0, spend=0.0),
prisma_client=prisma,
user_api_key_cache=UserApiKeyCache(),
proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()),
)
@pytest.mark.asyncio
async def test_reserve_budget_for_added_tags_reserves_only_the_hook_tags_counter(spend_counter_cache: DualCache):
"""The body tag and the key were reserved at auth; the post-hook reservation touches only the added tag."""
reservation: Final = await _reserve_added_tags(
"/v1/chat/completions", _budgeted_tag_prisma((HOOK_TAG, BODY_TAG), max_budget=1.0)
)
assert reservation is not None
entries: Final = reservation["entries"]
assert isinstance(entries, list)
assert [entry["counter_key"] for entry in entries] == [f"spend:tag:{HOOK_TAG}"]
reserved_cost: Final = reservation["reserved_cost"]
assert isinstance(reserved_cost, float) and reserved_cost > 0
assert spend_counter_cache.in_memory_cache.get_cache(key=f"spend:tag:{HOOK_TAG}") == pytest.approx(reserved_cost)
assert spend_counter_cache.in_memory_cache.get_cache(key=f"spend:tag:{BODY_TAG}") is None
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:hashed-hook-tag-key") is None
@pytest.mark.asyncio
async def test_reserve_budget_for_added_tags_rejects_a_tag_with_no_room_for_the_estimate(
spend_counter_cache: DualCache,
):
"""Two requests race past the read check; the second reservation finds the estimate no longer fits."""
prisma: Final = _budgeted_tag_prisma((HOOK_TAG,), max_budget=0.000001)
first: Final = await _reserve_added_tags("/v1/chat/completions", prisma)
assert first is not None
assert first["reserved_cost"] == pytest.approx(0.000001)
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _reserve_added_tags("/v1/chat/completions", prisma)
assert exc_info.value.entity_id == HOOK_TAG
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ("/guardrails/apply_guardrail", "/v1/models", *TOKEN_COUNTING_ROUTES[:2]))
async def test_reserve_budget_for_added_tags_skips_routes_auth_never_reserves(spend_counter_cache: DualCache, route):
assert await _reserve_added_tags(route, _budgeted_tag_prisma((HOOK_TAG,), max_budget=1.0)) is None
assert spend_counter_cache.in_memory_cache.get_cache(key=f"spend:tag:{HOOK_TAG}") is None
@pytest.mark.asyncio
async def test_reserve_budget_for_added_tags_ignores_tags_without_a_budget(spend_counter_cache: DualCache):
assert await _reserve_added_tags("/v1/chat/completions", _budgeted_tag_prisma((), max_budget=1.0)) is None
BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6"
CONVERSE_BODY: Final = {
"messages": [{"role": "user", "content": [{"text": "Reply with one word: pong"}]}],

View file

@ -396,8 +396,13 @@ class TestProxyBaseLLMRequestProcessing:
)
tag_check = AsyncMock()
monkeypatch.setattr(litellm.proxy.common_request_processing, "tag_max_budget_check_for_tags", tag_check)
monkeypatch.setattr(
litellm.proxy.common_request_processing, "reserve_budget_for_added_tags", AsyncMock(return_value=None)
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-master")
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_auth", None)
return mock_request, mock_proxy_logging_obj, tag_check
@pytest.mark.asyncio
@ -607,6 +612,137 @@ class TestProxyBaseLLMRequestProcessing:
tag_check.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("master_key", "user_custom_auth", "general_settings", "checked"),
[
(None, None, {}, False),
(None, None, {"enable_jwt_auth": True}, True),
("sk-master", object(), {}, False),
("sk-master", object(), {"custom_auth_run_common_checks": True}, True),
("sk-master", None, {}, True),
],
)
async def test_common_processing_pre_call_logic_enforces_hook_added_tags_only_where_auth_runs_common_checks(
self, monkeypatch, master_key, user_custom_auth, general_settings, checked
):
"""A deployment whose auth wrapper skips common_checks (no-auth dev mode, custom auth without opt-in)
never budget-checked tags before, so a hook-added tag must not start 429ing it."""
processing_obj = ProxyBaseLLMRequestProcessing(data={})
async def mock_pre_call_hook(user_api_key_dict, data, call_type):
data["metadata"]["tags"].append("guardrail-tag")
return data
mock_request, mock_proxy_logging_obj, tag_check = self._tag_budget_rig(
monkeypatch, request_data={"model": "live-mini", "metadata": {"tags": []}}, pre_call_hook=mock_pre_call_hook
)
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", master_key)
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_auth", user_custom_auth)
await processing_obj.common_processing_pre_call_logic(
request=mock_request,
general_settings=general_settings,
user_api_key_dict=ProxyUserAPIKeyAuth(token="test-token"),
proxy_logging_obj=mock_proxy_logging_obj,
proxy_config=MagicMock(spec=ProxyConfig),
route_type="acompletion",
)
assert tag_check.await_count == (1 if checked else 0)
@staticmethod
def _reservation(counter_key: str) -> dict:
return {
"reserved_cost": 0.5,
"entries": [
{"counter_key": counter_key, "entity_type": "Tag", "entity_id": counter_key, "reserved_cost": 0.5}
],
"finalized": False,
"input_cost": 0.1,
"input_tokens": 3,
}
async def _run_with_hook_added_tag(self, monkeypatch, user_api_key_dict, general_settings: dict):
processing_obj = ProxyBaseLLMRequestProcessing(data={})
async def mock_pre_call_hook(user_api_key_dict, data, call_type):
data["metadata"]["tags"].append("guardrail-tag")
return data
mock_request, mock_proxy_logging_obj, _ = self._tag_budget_rig(
monkeypatch, request_data={"model": "live-mini", "metadata": {"tags": []}}, pre_call_hook=mock_pre_call_hook
)
reserve = AsyncMock(return_value=self._reservation("spend:tag:guardrail-tag"))
monkeypatch.setattr(litellm.proxy.common_request_processing, "reserve_budget_for_added_tags", reserve)
await processing_obj.common_processing_pre_call_logic(
request=mock_request,
general_settings=general_settings,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=mock_proxy_logging_obj,
proxy_config=MagicMock(spec=ProxyConfig),
route_type="acompletion",
)
return processing_obj, reserve
@pytest.mark.asyncio
async def test_common_processing_pre_call_logic_folds_the_hook_tag_reservation_into_the_auth_reservation(
self, monkeypatch
):
"""Auth reserved the body tags before the hook ran. The hook-added tag gets its own reservation
so a burst cannot overshoot it, and it must join the same reservation object auth left on the
key, since that object is what the success, failure and cancel paths settle."""
auth_reservation = self._reservation("spend:key:test-token")
user_api_key_dict = ProxyUserAPIKeyAuth(token="test-token")
user_api_key_dict.budget_reservation = auth_reservation
processing_obj, reserve = await self._run_with_hook_added_tag(
monkeypatch, user_api_key_dict, general_settings={"fail_closed_budget_enforcement": True}
)
reserve.assert_awaited_once()
assert reserve.call_args.kwargs["tags"] == ("guardrail-tag",)
assert reserve.call_args.kwargs["request_body"] is processing_obj.data
assert reserve.call_args.kwargs["route"] == "/v1/chat/completions"
assert reserve.call_args.kwargs["valid_token"] is user_api_key_dict
assert reserve.call_args.kwargs["fail_closed_budget_enforcement"] is True
assert user_api_key_dict.budget_reservation is auth_reservation
assert [entry["counter_key"] for entry in auth_reservation["entries"]] == [
"spend:key:test-token",
"spend:tag:guardrail-tag",
]
assert "user_api_key_budget_reservation" not in processing_obj.data["metadata"]
@pytest.mark.asyncio
async def test_common_processing_pre_call_logic_attaches_the_hook_tag_reservation_when_auth_reserved_nothing(
self, monkeypatch
):
"""With no reservation from auth, the hook tag's reservation has to be placed where the
settlement paths look: on the auth object (failure, cancel) and in the request metadata (success)."""
user_api_key_dict = ProxyUserAPIKeyAuth(token="test-token")
processing_obj, reserve = await self._run_with_hook_added_tag(
monkeypatch, user_api_key_dict, general_settings={}
)
reservation = reserve.return_value
assert user_api_key_dict.budget_reservation is reservation
assert processing_obj.data["metadata"]["user_api_key_budget_reservation"] is reservation
@pytest.mark.asyncio
async def test_common_processing_pre_call_logic_skips_the_hook_tag_reservation_when_reservation_is_disabled(
self, monkeypatch
):
"""disable_budget_reservation turns off auth's reservation too, so only the read check runs."""
user_api_key_dict = ProxyUserAPIKeyAuth(token="test-token")
_, reserve = await self._run_with_hook_added_tag(
monkeypatch, user_api_key_dict, general_settings={"disable_budget_reservation": True}
)
reserve.assert_not_awaited()
assert user_api_key_dict.budget_reservation is None
@pytest.mark.asyncio
async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails(
self, monkeypatch