mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
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:
parent
0333ca1f34
commit
27bc2a80e2
7 changed files with 428 additions and 23 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 []
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"}]}],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue