diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index be697ffb792..e9983691253 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -1,5 +1,4 @@ import re -from functools import lru_cache from typing import Final from fastapi import HTTPException, Request, status @@ -16,6 +15,26 @@ from litellm.proxy._types import ( from .auth_checks_organization import _user_is_org_admin + +def _placeholder_to_regex(match: "re.Match[str]") -> str: + placeholder: Final = match.group(0).strip("{}") + if placeholder.endswith(":path"): + # allow "/" in the placeholder value, but don't eat the route suffix after ":" + return r"[^:]+" + return r"[^/]+" + + +def _route_pattern_regex(pattern: str) -> str: + """Anchored regex for a route template, so a precompiled copy cannot drift from the live matcher.""" + expanded: Final = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) + return f"^{expanded}$" + + +_EXACT_INFO_ROUTES: Final = frozenset(LiteLLMRoutes.info_routes.value) +_TEMPLATED_INFO_ROUTE_PATTERNS: Final = tuple( + re.compile(_route_pattern_regex(route)) for route in LiteLLMRoutes.info_routes.value if "{" in route +) + # Management write routes denied to PROXY_ADMIN_VIEW_ONLY. Adding a new write # endpoint to a management router REQUIRES adding it here too — the surrounding # check falls through to "allow" if the route is not matched, which previously @@ -438,7 +457,6 @@ class RouteChecks: return RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.management_routes.value) @staticmethod - @lru_cache(maxsize=2048) def is_info_route(route: str) -> bool: """ Check if route is an info route @@ -446,11 +464,15 @@ class RouteChecks: Pattern-aware, like ``is_management_route``, so an info route carrying a path parameter is as reachable as one without: the incoming route holds a resolved id, never the ``{...}`` template. - Cached because this runs per request off ``_check_end_user_budget`` and the allowlist is a - module constant, so an uncached call rebuilds and matches 25 regexes to answer the same - question. ``normalize_request_route`` in ``auth_utils`` is bounded the same way. + Matched off the precomputed sets rather than ``check_route_access`` because this runs per + request off ``_check_end_user_budget``, and rebuilding 25 regexes to answer a question whose + allowlist is a module constant costs more than the answer is worth. Caching instead would key + on a route that carries resolved ids, so its working set is unbounded on exactly the traffic + that needs it most. """ - return RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.info_routes.value) + return route in _EXACT_INFO_ROUTES or any( + pattern.match(route) is not None for pattern in _TEMPLATED_INFO_ROUTE_PATTERNS + ) @staticmethod def _is_azure_openai_route(route: str) -> bool: @@ -490,17 +512,7 @@ class RouteChecks: if not isinstance(route, str): return False - def _placeholder_to_regex(match: re.Match) -> str: - placeholder: Final = match.group(0).strip("{}") - if placeholder.endswith(":path"): - # allow "/" in the placeholder value, but don't eat the route suffix after ":" - return r"[^:]+" - return r"[^/]+" - - pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) - # Anchor the pattern to match the entire string - pattern = f"^{pattern}$" - if re.match(pattern, route): + if re.match(_route_pattern_regex(pattern), route): return True return False diff --git a/litellm/proxy/management_endpoints/key_budget_resolver.py b/litellm/proxy/management_endpoints/key_budget_resolver.py index 923f0a57b5f..601c17f4ba5 100644 --- a/litellm/proxy/management_endpoints/key_budget_resolver.py +++ b/litellm/proxy/management_endpoints/key_budget_resolver.py @@ -157,8 +157,11 @@ _USER_ON_TEAM_KEY_NOTE: Final = KeyBudgetNote( ) _THROTTLE_NOTE: Final = KeyBudgetNote( code="throttled_instead_of_blocked", - severity="warning", - text="this key opted into throttle_on_budget_exceeded, so exceeding it slows requests instead of blocking", + severity="info", + text=( + "past the limit the key's own tpm and rpm limits are scaled down by " + "litellm.budget_exceeded_throttle_percentage; no other budget on this key throttles" + ), ) @@ -516,10 +519,10 @@ def _effective_comparison(plan: _PlannedBudget, reservation_enabled: bool) -> Bu Reservation runs ahead of the read-time check and blocks once spend reaches the cap, so for every scope it covers it, not the read-time comparison, decides when a request stops going through. - It builds no counter at all for a non-positive cap, though, which leaves the read-time operator - in charge there and is the one case where reporting the tightened one would invent a denial. + It builds no counter at all for a non-positive cap, and it releases the one it built for a key + that throttles instead of blocking, so neither of those is tightened by it. """ - if plan.enforcement == "soft" or not reservation_enabled: + if plan.enforcement != "hard" or not reservation_enabled: return plan.comparison if plan.max_budget is None or plan.max_budget <= 0: return plan.comparison @@ -855,18 +858,19 @@ def _plan_key(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded token: Final = context.valid_token + throttled: Final = should_throttle_budget_exceeded(token) hard: Final = _PlannedBudget( scope="key", entity_id=token.key_alias, entity_label=token.key_alias, - enforcement="hard", + enforcement="throttled" if throttled else "hard", max_budget=token.max_budget, comparison=">=", source=_key_max_budget_source(context), spend_source=_CounterSpend(counter_key=key_spend_counter(token.token), fallback_spend=token.spend or 0.0), budget_duration=token.budget_duration, budget_reset_at=token.budget_reset_at, - notes=(_THROTTLE_NOTE,) if should_throttle_budget_exceeded(token) else (), + notes=(_THROTTLE_NOTE,) if throttled else (), ) soft: Final = _PlannedBudget( scope="key", diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 442c90666ef..0c4fc6658c3 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3818,7 +3818,9 @@ async def key_budgets_fn( - entity_type: Litellm_EntityType - The entity a `BudgetExceededError` from this scope names, so a denial message maps back to a row here - entity_id / entity_label: str | None - Which entity is limited, and its human-facing alias - - enforcement: str - `hard` blocks the request, `soft` only raises an alert + - enforcement: str - `hard` blocks the request, `soft` only raises an alert, `throttled` + scales the key's rate limits down instead of denying anything. Only the key's own + `max_budget` can be `throttled`; every other scope on the same key still blocks - max_budget: float | None - The limit in effect. `null` means this scope applies to the key but places no limit on it - spend: float | None - Spend as the enforcing check reads it, from the same cross-pod diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index bc887ae7ccf..574ac81290c 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -126,7 +126,7 @@ BudgetScope = Literal[ "end_user_model", ] -BudgetEnforcement = Literal["hard", "soft"] +BudgetEnforcement = Literal["hard", "soft", "throttled"] BudgetComparison = Literal[">=", ">"] diff --git a/tests/test_litellm/proxy/auth/test_info_routes.py b/tests/test_litellm/proxy/auth/test_info_routes.py index b2512607c9f..54430c53d74 100644 --- a/tests/test_litellm/proxy/auth/test_info_routes.py +++ b/tests/test_litellm/proxy/auth/test_info_routes.py @@ -196,15 +196,36 @@ def test_templated_route_matching_does_not_widen_unrelated_key_routes(): assert RouteChecks.is_info_route("/key/some-hash/budgets") is True -def test_info_route_matching_is_cached_without_answering_for_the_wrong_route(): - """Pattern matching runs per request off the end-user budget check, so it is memoised; the cache must not blur routes.""" - RouteChecks.is_info_route.cache_clear() +def test_info_route_matching_answers_exactly_what_the_generic_matcher_would(): + """ + This runs per request off the end-user budget check, so it matches precomputed sets instead of + rebuilding 25 regexes. That is only safe while it agrees with the generic matcher on every route. + """ + corpus = { + entry + for member in LiteLLMRoutes + if isinstance(member.value, list) + for entry in member.value + if isinstance(entry, str) + } | { + "/key/hash-a/budgets", + "/key/budgets", + "/key/hash-a/regenerate", + "/chat/completions", + "/v1/responses/resp_9f2", + "/key//budgets", + "/key/a/b/budgets", + "/keyX/a/budgets", + "/key/a/budgetsX", + "", + "/", + } - assert RouteChecks.is_info_route("/key/hash-a/budgets") is True - assert RouteChecks.is_info_route("/key/hash-a/regenerate") is False - assert RouteChecks.is_info_route("/chat/completions") is False - assert RouteChecks.is_info_route("/key/hash-b/budgets") is True - assert RouteChecks.is_info_route.cache_info().hits == 0, "four distinct routes cannot share an answer" - - assert RouteChecks.is_info_route("/key/hash-a/regenerate") is False - assert RouteChecks.is_info_route.cache_info().hits == 1 + disagreements = { + route + for route in corpus + if RouteChecks.is_info_route(route) + != RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.info_routes.value) + } + assert disagreements == set() + assert len(corpus) > 100, "a corpus this small would not have exercised the templated entry" diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 6776c51e4d8..717c0150b77 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -16463,7 +16463,9 @@ from litellm.proxy._types import LiteLLM_ProjectTableCachedObj # noqa: E402 from litellm.proxy.auth.auth_checks import TeamMemberBudget # noqa: E402 from litellm.proxy.management_endpoints.key_budget_resolver import ( # noqa: E402 _match_model_budget_key, + _RecordedSpend as _BudgetsRecordedSpend, _MODEL_BUDGET_COLD_NOTE, + _THROTTLE_NOTE, _read_end_user_model_spend, _read_key_model_spend, _request_models, @@ -17237,6 +17239,104 @@ async def test_key_budgets_call_an_unreadable_counter_unreadable_rather_than_rep assert project_entry.spend_state == "live", "the recorded-spend scopes do not go through the counter" +@pytest.mark.asyncio +async def test_key_budgets_say_a_throttling_key_throttles_rather_than_calling_it_a_blocker(): + """`hard` plus a note meant every client that skipped the prose reported a denial that never happened.""" + token = _budgets_token( + tpm_limit=1000, + metadata={"tags": ["prod"], "throttle_on_budget_exceeded": True}, + ) + world = _fully_populated_world() + world["team_member"] = TeamMemberBudget(max_budget=50.0, recorded_spend=8.0, source="budget_table:budget-member") + with patch("litellm.budget_exceeded_throttle_percentage", 0.5), _budgets_world(**world): + budgets = await resolve_key_budgets( + valid_token=token, + end_user_id=None, + deps=_budgets_deps( + read_spend=_RecordingSpendReader(_BUDGETS_SPEND_AT_LIMIT), + general_settings={"apply_user_budget_to_team_keys": True}, + ), + ) + + key_entry = next(e for e in budgets if e.scope == "key" and e.enforcement != "soft") + assert key_entry.enforcement == "throttled" + assert key_entry.status == "exceeded", "spend really is past the cap; what changes is what happens next" + assert key_entry.comparison == ">=", "throttling starts at the same point a block would have" + assert _THROTTLE_NOTE in key_entry.notes + + +@pytest.mark.parametrize( + ("enforcement", "expected"), + [("hard", ">="), ("soft", ">"), ("throttled", ">")], +) +def test_key_budgets_apply_the_reservation_operator_only_to_budgets_that_block(enforcement, expected): + """ + The reservation only decides when a request stops going through for a budget that denies one. It + never runs for a soft budget and it releases the entry it built for a throttling key, so neither + can be reported as blocking earlier than the read-time check does. + """ + from litellm.proxy.management_endpoints.key_budget_resolver import _effective_comparison, _PlannedBudget + + plan = _PlannedBudget( + scope="key", + entity_id="k", + entity_label=None, + enforcement=enforcement, + max_budget=100.0, + comparison=">", + source="key.max_budget", + spend_source=_BudgetsRecordedSpend(0.0), + ) + + assert _effective_comparison(plan=plan, reservation_enabled=True) == expected + + +@pytest.mark.asyncio +async def test_key_budgets_keep_every_other_scope_blocking_on_a_throttling_key(): + """Only the key's own counter is released for a throttling key; the rest still raise.""" + token = _budgets_token( + tpm_limit=1000, + metadata={"tags": ["prod"], "throttle_on_budget_exceeded": True}, + ) + world = _fully_populated_world() + world["team_member"] = TeamMemberBudget(max_budget=50.0, recorded_spend=8.0, source="budget_table:budget-member") + with patch("litellm.budget_exceeded_throttle_percentage", 0.5), _budgets_world(**world): + budgets = await resolve_key_budgets( + valid_token=token, + end_user_id="end-user-budgets", + deps=_budgets_deps( + read_spend=_RecordingSpendReader(_BUDGETS_SPEND_AT_LIMIT), + general_settings={"apply_user_budget_to_team_keys": True}, + ), + ) + + throttled = {e.scope for e in budgets if e.enforcement == "throttled"} + assert throttled == {"key"} + for scope in ("key_window", "team", "team_member", "user", "organization", "tag", "end_user"): + entry = next(e for e in budgets if e.scope == scope and e.enforcement != "soft") + assert entry.enforcement == "hard", scope + + +@pytest.mark.asyncio +async def test_key_budgets_call_a_key_hard_when_throttling_cannot_actually_engage(): + """The flag alone does nothing: without a rate limit to scale or a configured percentage it still blocks.""" + flagged = _budgets_token(tpm_limit=1000, metadata={"throttle_on_budget_exceeded": True}) + + with patch("litellm.budget_exceeded_throttle_percentage", None), _budgets_world(**_fully_populated_world()): + without_percentage = await resolve_key_budgets(valid_token=flagged, end_user_id=None, deps=_budgets_deps()) + + no_rate_limit = _budgets_token(metadata={"throttle_on_budget_exceeded": True}) + with patch("litellm.budget_exceeded_throttle_percentage", 0.5), _budgets_world(**_fully_populated_world()): + without_rate_limit = await resolve_key_budgets( + valid_token=no_rate_limit, end_user_id=None, deps=_budgets_deps() + ) + + for budgets in (without_percentage, without_rate_limit): + entry = next(e for e in budgets if e.scope == "key" and e.enforcement != "soft") + assert entry.enforcement == "hard" + assert entry.notes == () + + def test_key_budgets_classify_every_note_code_and_leave_none_to_a_default(): """ Severity is the only thing a client has for a code its build predates, so every code is classified @@ -17256,13 +17356,13 @@ def test_key_budgets_classify_every_note_code_and_leave_none_to_a_default(): "rolling_window": "info", "user_budget_not_applied_to_team_key": "info", "model_budget_fails_open": "info", + "throttled_instead_of_blocked": "info", # only the note states it "custom_auth_may_override_end_user_cap": "warning", "end_user_route_only": "warning", "per_model_counters": "warning", "project_spend_not_tracked": "warning", "request_tags_add_budgets": "warning", - "throttled_instead_of_blocked": "warning", }