From 27bc2a80e26a520012234ad29f236b422bbf6935 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Tue, 15 Sep 2026 03:00:20 -0700 Subject: [PATCH] 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. --- litellm/proxy/auth/auth_checks.py | 21 +++ litellm/proxy/auth/user_api_key_auth.py | 16 +-- litellm/proxy/common_request_processing.py | 70 ++++++++- .../spend_tracking/budget_reservation.py | 79 +++++++++- .../proxy/auth/test_auth_checks.py | 27 ++++ .../spend_tracking/test_budget_reservation.py | 102 ++++++++++++- .../proxy/test_common_request_processing.py | 136 ++++++++++++++++++ 7 files changed, 428 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e2ef7ef35bd..56703d8a971 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9491f77ecfc..095f61c666b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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 diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 656329c612b..481b7d87382 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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, diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 6074a50a69b..40e934c197f 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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 [] diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index b9619c3c07a..09809ee6316 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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 diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 3e0acf917aa..ecea188e864 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -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"}]}], diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index c5141ea02da..193d5fa6a4d 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -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