diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index c3a00add875..14c8977dce6 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -140,8 +140,15 @@ class ProxyRateLimitError(HTTPException, RateLimitError): # type: ignore[misc] ] = RateLimitErrorCategory.LITELLM_RATE_LIMIT, rate_limit_type: Optional[Union[str, RateLimitType]] = None, model: Optional[str] = None, - llm_provider: str = "litellm_proxy", + llm_provider: Optional[str] = "litellm_proxy", ): + # Normalize None → safe defaults so callers (and the resolver helper + # in `rate_limiter_utils`) can pass `None` without producing an + # instance whose `.llm_provider` attribute is `None` — that would + # break Prometheus' `_get_exception_class_name` (it calls + # `.capitalize()` on the provider string). + model = model or "" + llm_provider = llm_provider or "litellm_proxy" message = _coerce_message(detail) stringified_headers: Optional[Dict[str, str]] = ( {k: str(v) for k, v in headers.items()} if headers else None diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 43d82ade391..57e4241374a 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -33,6 +33,7 @@ from litellm.exceptions import RateLimitErrorCategory, RateLimitType from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -106,6 +107,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): descriptors: List["RateLimitDescriptor"], batch_usage: BatchFileUsage, limit_type: str, + requested_model: Optional[str] = None, ) -> NoReturn: """Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded.""" from datetime import datetime @@ -150,6 +152,9 @@ class _PROXY_BatchRateLimiter(CustomLogger): f"Limit resets at: {reset_time_formatted}" ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( + requested_model + ) raise ProxyRateLimitError( detail=detail, headers={ @@ -166,6 +171,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): if limit_type == "tokens" else RateLimitType.REQUESTS ), + model=resolved_model, + llm_provider=llm_provider, ) async def _check_and_increment_batch_counters( @@ -207,6 +214,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) if rate_limit_response["overall_code"] == "OVER_LIMIT": + requested_model = data.get("model") if data else None for status in rate_limit_response["statuses"]: if status["code"] == "OVER_LIMIT": self._raise_rate_limit_error( @@ -214,6 +222,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): descriptors, batch_usage, status["rate_limit_type"], + requested_model=requested_model, ) async def count_input_file_usage( diff --git a/litellm/proxy/hooks/dynamic_rate_limiter.py b/litellm/proxy/hooks/dynamic_rate_limiter.py index 156ae3d1147..b9e2bd12ecf 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter.py @@ -11,15 +11,17 @@ from litellm import ModelResponse, Router from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth from litellm.exceptions import RateLimitType +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import ( + convert_priority_to_percent, + resolve_llm_provider_for_rate_limit, +) from litellm.types.router import ModelGroupInfo from litellm.types.utils import CallTypesLiteral from litellm.utils import get_utc_datetime -from .rate_limiter_utils import convert_priority_to_percent - class DynamicRateLimiterCache: """ @@ -218,6 +220,9 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): ) ### CHECK TPM ### if available_tpm is not None and available_tpm == 0: + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( + data.get("model") + ) raise ProxyRateLimitError( detail={ "error": "Key={} over available TPM={}. Model TPM={}, Active keys={}".format( @@ -228,10 +233,14 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): ) }, rate_limit_type=RateLimitType.TOKENS, - model=data.get("model"), + model=resolved_model, + llm_provider=llm_provider, ) ### CHECK RPM ### elif available_rpm is not None and available_rpm == 0: + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( + data.get("model") + ) raise ProxyRateLimitError( detail={ "error": "Key={} over available RPM={}. Model RPM={}, Active keys={}".format( @@ -242,7 +251,8 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): ) }, rate_limit_type=RateLimitType.REQUESTS, - model=data.get("model"), + model=resolved_model, + llm_provider=llm_provider, ) elif available_rpm is not None or available_tpm is not None: ## UPDATE CACHE WITH ACTIVE PROJECT diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 6aa7d98e349..493afe6105a 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -23,7 +23,10 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( RateLimitDescriptorRateLimitObject, _PROXY_MaxParallelRequestsHandler_v3, ) -from litellm.proxy.hooks.rate_limiter_utils import convert_priority_to_percent +from litellm.proxy.hooks.rate_limiter_utils import ( + convert_priority_to_percent, + resolve_llm_provider_for_rate_limit, +) from litellm.proxy.utils import InternalUsageCache from litellm.types.router import ModelGroupInfo from litellm.types.utils import CallTypesLiteral @@ -491,6 +494,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) if atomic_response["overall_code"] == "OVER_LIMIT": + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model) for status in atomic_response["statuses"]: if status["code"] != "OVER_LIMIT": continue @@ -513,7 +517,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): rate_limit_type=map_v3_rate_limit_type( status["rate_limit_type"] ), - model=model, + model=resolved_model, + llm_provider=llm_provider, ) if descriptor_key == "priority_model": verbose_proxy_logger.debug( @@ -540,7 +545,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): rate_limit_type=map_v3_rate_limit_type( status["rate_limit_type"] ), - model=model, + model=resolved_model, + llm_provider=llm_provider, ) # Fail-closed guard: overall_code says OVER_LIMIT but no status @@ -570,11 +576,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): rate_limit_type=map_v3_rate_limit_type( offending["rate_limit_type"] if offending else None ), - model=model, headers={ "retry-after": str(self.v3_limiter.window_size), "x-litellm-priority": priority or "default", }, + model=resolved_model, + llm_provider=llm_provider, ) # If priority is NOT enforced (saturation below threshold) but diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 32b43c26247..769348a0b88 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -4,9 +4,10 @@ from litellm import verbose_logger from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth from litellm.exceptions import RateLimitType +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit class _PROXY_MaxBudgetLimiter(CustomLogger): @@ -65,9 +66,14 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): # CHECK IF REQUEST ALLOWED if curr_spend >= max_budget: + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( + data.get("model") if data else None + ) raise ProxyRateLimitError( detail="Max budget limit reached.", rate_limit_type=RateLimitType.BUDGET, + model=resolved_model, + llm_provider=llm_provider, ) except HTTPException as e: raise e diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index d3dcb5200e3..20bfeb3a6d5 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -20,9 +20,10 @@ from typing import TYPE_CHECKING, Any, Optional, Union from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth from litellm.exceptions import RateLimitType +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache @@ -112,6 +113,9 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): ) if current_spend >= max_budget: + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( + data.get("model") if data else None + ) raise ProxyRateLimitError( detail=( f"Session budget exceeded for session {session_id}. " @@ -119,6 +123,8 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): f"max_budget_per_session: ${max_budget:.2f}." ), rate_limit_type=RateLimitType.BUDGET, + model=resolved_model, + llm_provider=llm_provider, ) return None diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index 785c3894183..525214ff6be 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -16,9 +16,10 @@ from typing import TYPE_CHECKING, Any, Optional, Union from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth from litellm.exceptions import RateLimitType +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache @@ -116,12 +117,17 @@ class _PROXY_MaxIterationsHandler(CustomLogger): current_count = await self._increment_and_get(cache_key) if current_count > max_iterations: + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( + data.get("model") if data else None + ) raise ProxyRateLimitError( detail=( f"Max iterations exceeded for session {session_id}. " f"Current count: {current_count}, max_iterations: {max_iterations}." ), rate_limit_type=RateLimitType.MAX_ITERATIONS, + model=resolved_model, + llm_provider=llm_provider, ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index a68eff6143a..b622241dfa5 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -12,12 +12,13 @@ from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth +from litellm.exceptions import RateLimitType from litellm.proxy.auth.auth_utils import ( get_key_model_rpm_limit, get_key_model_tpm_limit, ) -from litellm.exceptions import RateLimitType from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -87,6 +88,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): self.raise_rate_limit_error( additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current limits: max_parallel_requests: {max_parallel_requests}, tpm_limit: {tpm_limit}, rpm_limit: {rpm_limit}", rate_limit_type=triggered_type, + requested_model=data.get("model") if data else None, ) new_val = { "current_requests": 1, @@ -117,10 +119,16 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): triggered_type = RateLimitType.TOKENS else: triggered_type = RateLimitType.REQUESTS + requested_model = data.get("model") if data else None + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( + requested_model + ) raise ProxyRateLimitError( detail=f"LiteLLM Rate Limit Handler for rate limit type = {rate_limit_type}. {CommonProxyErrors.max_parallel_request_limit_reached.value}. current rpm: {current['current_rpm']}, rpm limit: {rpm_limit}, current tpm: {current['current_tpm']}, tpm limit: {tpm_limit}, current max_parallel_requests: {current['current_requests']}, max_parallel_requests: {max_parallel_requests}", headers={"retry-after": str(self.time_to_next_minute())}, rate_limit_type=triggered_type, + model=resolved_model, + llm_provider=llm_provider, ) await self.internal_usage_cache.async_batch_set_cache( @@ -147,6 +155,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): self, additional_details: Optional[str] = None, rate_limit_type: Optional[RateLimitType] = None, + requested_model: Optional[str] = None, ) -> NoReturn: """ Raise a 429 with a retry-after header for litellm-proxy parallel-request limits. @@ -163,6 +172,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): (the global-limit branch in ``async_pre_call_hook`` and the all-zeros base case in ``check_key_in_limits``). Callers that know the dimension exactly should pass it explicitly. + + ``requested_model`` is resolved via :func:`get_llm_provider` so the + raised exception carries ``llm_provider`` (and a stripped ``model``) + for downstream loggers (Prometheus failure metric, observability + callbacks). Falls back to ``llm_provider="litellm_proxy"`` when the + model is missing or unparseable — see + :func:`resolve_llm_provider_for_rate_limit`. """ # additional_details is optional; build the detail with a None-guard # so callers that pass nothing don't get the literal string "None" @@ -170,10 +186,15 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): error_message = "Max parallel request limit reached" if additional_details is not None: error_message = error_message + " " + additional_details + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( + requested_model + ) raise ProxyRateLimitError( detail=error_message, headers={"retry-after": str(self.time_to_next_minute())}, rate_limit_type=rate_limit_type or RateLimitType.CONCURRENT_REQUESTS, + model=resolved_model, + llm_provider=llm_provider, ) async def get_all_cache_objects( @@ -265,7 +286,8 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # if above -> raise error if current_global_requests >= global_max_parallel_requests: self.raise_rate_limit_error( - additional_details=f"Hit Global Limit: Limit={global_max_parallel_requests}, current: {current_global_requests}" + additional_details=f"Hit Global Limit: Limit={global_max_parallel_requests}, current: {current_global_requests}", + requested_model=data.get("model") if data else None, ) # if below -> increment else: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 0928cab7c00..7083aade341 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -36,6 +36,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ( ProxyRateLimitError, map_v3_rate_limit_type, ) +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject from litellm.types.utils import ModelResponse, Usage @@ -1839,6 +1840,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, response: RateLimitResponse, descriptors: List[RateLimitDescriptor], + requested_model: Optional[str] = None, ) -> None: """Handle rate limit exceeded by raising :class:`ProxyRateLimitError` (a 429).""" for status in response["statuses"]: @@ -1871,6 +1873,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): f"Limit resets at: {reset_time_formatted}" ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( + requested_model + ) raise ProxyRateLimitError( detail=detail, headers={ @@ -1879,6 +1884,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): "reset_at": reset_time_formatted, }, rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]), + model=resolved_model, + llm_provider=llm_provider, ) async def async_pre_call_hook( @@ -1979,6 +1986,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._handle_rate_limit_error( response=response, descriptors=descriptors, + requested_model=requested_model, ) else: # add descriptors to request headers @@ -2024,6 +2032,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._handle_rate_limit_error( response=tpm_response, descriptors=descriptors, + requested_model=requested_model, ) else: data["_litellm_rate_limit_descriptors"] = descriptors diff --git a/litellm/proxy/hooks/rate_limiter_utils.py b/litellm/proxy/hooks/rate_limiter_utils.py index 927bac0de58..2ed96f717ae 100644 --- a/litellm/proxy/hooks/rate_limiter_utils.py +++ b/litellm/proxy/hooks/rate_limiter_utils.py @@ -2,11 +2,56 @@ Shared utility functions for rate limiter hooks. """ -from typing import Optional, Union +from typing import Optional, Tuple, Union +import litellm +from litellm._logging import verbose_proxy_logger from litellm.types.router import ModelGroupInfo from litellm.types.utils import PriorityReservationDict +PROXY_LLM_PROVIDER_FALLBACK = "litellm_proxy" + + +def resolve_llm_provider_for_rate_limit( + model: Optional[str], +) -> Tuple[str, str]: + """ + Resolve ``(model, llm_provider)`` for a request being rejected by an + internal proxy-side rate-limit hook. + + These hooks fire from ``async_pre_call_hook`` — well before + :func:`litellm.get_llm_provider` is invoked anywhere else in the request + lifecycle — so the raised 429 would otherwise have an empty + ``llm_provider`` field, making the resulting Prometheus + ``litellm_proxy_failed_requests_metric`` show up with + ``exception_class="RateLimitError"`` and no provider attribution. + + Wrapped defensively: if ``model`` is missing, malformed, or + ``get_llm_provider`` raises (unknown alias, router-only model, etc.) we + fall back to ``("", "litellm_proxy")`` so we never break the request path + by piling a second exception on top of the rate-limit one we're trying to + raise. + """ + if not model: + return "", PROXY_LLM_PROVIDER_FALLBACK + try: + resolved_model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=model, + ) + return ( + resolved_model or model, + custom_llm_provider or PROXY_LLM_PROVIDER_FALLBACK, + ) + except Exception as e: + verbose_proxy_logger.debug( + "rate_limiter_utils.resolve_llm_provider_for_rate_limit: " + "could not resolve provider for model=%s, falling back to %s. err=%s", + model, + PROXY_LLM_PROVIDER_FALLBACK, + str(e), + ) + return model, PROXY_LLM_PROVIDER_FALLBACK + def convert_priority_to_percent( value: Union[float, PriorityReservationDict], model_info: Optional[ModelGroupInfo] diff --git a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py new file mode 100644 index 00000000000..5240132bc13 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py @@ -0,0 +1,964 @@ +""" +Regression tests for the "provider field missing" bug on proxy-side +rate-limit errors. + +Background +---------- +The proxy's internal rate-limit hooks (parallel_request_limiter, +parallel_request_limiter_v3, dynamic_rate_limiter, dynamic_rate_limiter_v3, +batch_rate_limiter, max_budget_limiter, max_iterations_limiter, +max_budget_per_session_limiter) all fire from ``async_pre_call_hook`` — +*before* :func:`litellm.get_llm_provider` runs anywhere else in the request +lifecycle. + +Until now, those hooks raised a bare ``HTTPException(429, ...)`` which carries +no ``llm_provider`` / ``model`` attribute. Downstream: + +- The Prometheus ``litellm_proxy_failed_requests_metric`` reads + ``exception.llm_provider`` via ``_get_exception_class_name`` — it came back + empty, so dashboards showed ``exception_class="HTTPException"`` with no + provider attribution. +- Observability callbacks that ``isinstance(e, RateLimitError)`` for + category routing missed these entirely. + +The fix wraps every internal raise site in +:class:`ProxyRateLimitError` (an ``HTTPException`` *and* a +``litellm.RateLimitError``), and resolves ``model`` / ``llm_provider`` from +``data["model"]`` via :func:`get_llm_provider`. When the model is missing or +unparseable we fall back to ``llm_provider="litellm_proxy"`` so we never break +the request path with a second exception. + +These tests pin both the happy path (provider correctly resolved) and the +fallback path (unknown model, missing model) for every limiter. +""" + +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.caching.caching import DualCache +from litellm.exceptions import RateLimitError +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.batch_rate_limiter import ( + BatchFileUsage, + _PROXY_BatchRateLimiter, +) +from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler +from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3, +) +from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter +from litellm.proxy.hooks.max_budget_per_session_limiter import ( + _PROXY_MaxBudgetPerSessionHandler, +) +from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler +from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, +) +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, +) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import ( + PROXY_LLM_PROVIDER_FALLBACK, + resolve_llm_provider_for_rate_limit, +) +from litellm.proxy.utils import InternalUsageCache +from litellm.types.agents import AgentResponse + + +# --------------------------------------------------------------------------- +# Helper class itself +# --------------------------------------------------------------------------- + + +class TestProxyRateLimitErrorClass: + """Pin the dual ``HTTPException`` + ``RateLimitError`` shape.""" + + def test_is_both_http_exception_and_rate_limit_error(self): + e = ProxyRateLimitError( + detail="boom", + model="gpt-4o-mini", + llm_provider="openai", + ) + # FastAPI handler keys off HTTPException to render the 429. + assert isinstance(e, HTTPException) + # Prometheus / observability key off RateLimitError + .llm_provider. + assert isinstance(e, RateLimitError) + assert e.status_code == 429 + assert e.model == "gpt-4o-mini" + assert e.llm_provider == "openai" + # ProxyRateLimitError prefixes message via RateLimitError.__init__. + assert "boom" in e.message + assert e.detail == "boom" + + def test_dict_detail_is_stringified_for_message(self): + # Some hooks pass a dict detail (e.g. dynamic_rate_limiter v1) — the + # `message` attr (read by RateLimitError.__str__ and observability + # callbacks) must still be a string. + e = ProxyRateLimitError( + detail={"error": "over rpm"}, + model="claude-3-5-sonnet", + llm_provider="anthropic", + ) + assert isinstance(e.message, str) + assert "over rpm" in e.message + + def test_defaults_to_litellm_proxy_provider(self): + e = ProxyRateLimitError(detail="x") + assert e.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + assert e.model == "" + + def test_none_provider_normalized_to_fallback(self): + e = ProxyRateLimitError( + detail="x", + model=None, + llm_provider=None, + ) + assert e.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + assert e.model == "" + + +class TestResolveLLMProviderForRateLimit: + @pytest.mark.parametrize( + "model, expected_provider", + [ + ("gpt-4o-mini", "openai"), + ("anthropic/claude-3-5-sonnet", "anthropic"), + ("bedrock/meta.llama3-1-70b-instruct-v1:0", "bedrock"), + ], + ) + def test_known_models_resolve_provider(self, model, expected_provider): + resolved_model, provider = resolve_llm_provider_for_rate_limit(model) + assert provider == expected_provider + assert resolved_model # non-empty + + @pytest.mark.parametrize("model", [None, "", "totally-not-a-real-model-name"]) + def test_missing_or_unknown_model_falls_back(self, model): + # Must never raise — the resolver wraps `get_llm_provider` defensively + # because raising here would mask the rate-limit error we're trying + # to surface to the user. + resolved_model, provider = resolve_llm_provider_for_rate_limit(model) + assert provider == PROXY_LLM_PROVIDER_FALLBACK + # Resolver returns the input model verbatim on the unknown branch so + # the `.model` attribute is never silently swapped to a different one. + if not model: + assert resolved_model == "" + else: + assert resolved_model == model + + def test_get_llm_provider_raising_is_swallowed(self): + # If get_llm_provider itself blows up (unexpected error), we still + # fall back rather than letting the secondary exception escape. + with patch.object( + litellm, + "get_llm_provider", + side_effect=RuntimeError("boom"), + ): + resolved_model, provider = resolve_llm_provider_for_rate_limit("anything") + assert provider == PROXY_LLM_PROVIDER_FALLBACK + assert resolved_model == "anything" + + +# --------------------------------------------------------------------------- +# parallel_request_limiter v1 +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_parallel_request_limiter_v1_populates_provider_when_at_rpm_limit(): + """ + Trip the per-key RPM cap and assert the raised exception carries + ``model`` / ``llm_provider`` resolved from ``data["model"]``. + """ + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-rl-test", + max_parallel_requests=10, + rpm_limit=1, + tpm_limit=10, + ) + data = {"model": "gpt-4o-mini"} + + # First request consumes the budget. + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data=data, + call_type="completion", + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data=data, + call_type="completion", + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert isinstance(exc, RateLimitError) + assert exc.llm_provider == "openai" + assert exc.model == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_parallel_request_limiter_v1_zero_limit_path_populates_provider(): + """ + When tpm_limit / rpm_limit is 0 the limiter takes the + ``raise_rate_limit_error`` path. That path receives ``requested_model`` + via the call-site change and must pass it through. + """ + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-rl-zero", + max_parallel_requests=0, + rpm_limit=10, + tpm_limit=10, + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "anthropic/claude-3-5-sonnet"}, + call_type="completion", + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert isinstance(exc, RateLimitError) + assert exc.llm_provider == "anthropic" + assert exc.model == "claude-3-5-sonnet" + + +@pytest.mark.asyncio +async def test_parallel_request_limiter_v1_global_limit_populates_provider(): + """global_max_parallel_requests path also threads the model through.""" + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-global") + + # Pre-fill the global counter so the next call exceeds it. + await handler.internal_usage_cache.async_set_cache( + key="global_max_parallel_requests", + value=5, + local_only=True, + litellm_parent_otel_span=None, + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={ + "model": "bedrock/meta.llama3-1-70b-instruct-v1:0", + "metadata": {"global_max_parallel_requests": 1}, + }, + call_type="completion", + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert exc.llm_provider == "bedrock" + assert exc.model == "meta.llama3-1-70b-instruct-v1:0" + + +@pytest.mark.asyncio +async def test_parallel_request_limiter_v1_unknown_model_falls_back(): + """ + When ``data["model"]`` is unparseable, the resolver falls back to + ``litellm_proxy`` — and crucially does *not* leak a secondary exception. + """ + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-rl-unknown", + max_parallel_requests=10, + rpm_limit=1, + tpm_limit=10, + ) + data = {"model": "totally-not-a-real-model"} + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data=data, + call_type="completion", + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data=data, + call_type="completion", + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert exc.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + # Resolver returns the input verbatim so we don't silently relabel the + # model in the user-facing 429 detail. + assert exc.model == "totally-not-a-real-model" + + +@pytest.mark.asyncio +async def test_parallel_request_limiter_v1_missing_model_falls_back(): + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-rl-no-model", + max_parallel_requests=10, + rpm_limit=1, + tpm_limit=10, + ) + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + exc = exc_info.value + assert exc.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + assert exc.model == "" + + +# --------------------------------------------------------------------------- +# parallel_request_limiter v3 +# --------------------------------------------------------------------------- + + +def _v3_over_limit_response(rate_limit_type: str = "rpm") -> dict: + return { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 1, + "limit_remaining": -1, + "rate_limit_type": rate_limit_type, + } + ], + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model, expected_provider", + [ + ("gpt-4o-mini", "openai"), + ("anthropic/claude-3-5-sonnet", "anthropic"), + ], +) +async def test_parallel_request_limiter_v3_populates_provider(model, expected_provider): + handler = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + + descriptors = [{"key": "key", "value": "v", "rate_limit": {"requests_per_unit": 1}}] + over = _v3_over_limit_response() + + with pytest.raises(HTTPException) as exc_info: + handler._handle_rate_limit_error( + response=over, + descriptors=descriptors, + requested_model=model, + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert isinstance(exc, RateLimitError) + assert exc.llm_provider == expected_provider + # v3 may strip the "anthropic/" prefix in the resolved model — accept + # either; we only care that the provider field is correct and the model + # is non-empty. + assert exc.model + + +@pytest.mark.asyncio +async def test_parallel_request_limiter_v3_unknown_model_falls_back(): + handler = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + descriptors = [{"key": "key", "value": "v", "rate_limit": {"requests_per_unit": 1}}] + + with pytest.raises(HTTPException) as exc_info: + handler._handle_rate_limit_error( + response=_v3_over_limit_response(), + descriptors=descriptors, + requested_model="totally-bogus", + ) + + assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + assert exc_info.value.model == "totally-bogus" + + +@pytest.mark.asyncio +async def test_parallel_request_limiter_v3_missing_model_falls_back(): + handler = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + descriptors = [{"key": "key", "value": "v", "rate_limit": {"requests_per_unit": 1}}] + + with pytest.raises(HTTPException) as exc_info: + handler._handle_rate_limit_error( + response=_v3_over_limit_response(), + descriptors=descriptors, + requested_model=None, + ) + + assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + assert exc_info.value.model == "" + + +# --------------------------------------------------------------------------- +# dynamic_rate_limiter v1 +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v1_tpm_zero_populates_provider(): + handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=DualCache()) + handler.check_available_usage = AsyncMock(return_value=(0, 5, 100, 5, 1)) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-dyn") + user_api_key_dict.metadata = {} + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "gpt-4o-mini"}, + call_type="completion", + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert isinstance(exc, RateLimitError) + assert exc.llm_provider == "openai" + assert exc.model == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v1_rpm_zero_populates_provider(): + handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=DualCache()) + handler.check_available_usage = AsyncMock(return_value=(5, 0, 5, 100, 1)) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-dyn") + user_api_key_dict.metadata = {} + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "anthropic/claude-3-5-sonnet"}, + call_type="completion", + ) + + exc = exc_info.value + assert exc.llm_provider == "anthropic" + assert exc.model == "claude-3-5-sonnet" + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v1_unknown_model_falls_back(): + handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=DualCache()) + handler.check_available_usage = AsyncMock(return_value=(0, 5, 100, 5, 1)) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-dyn") + user_api_key_dict.metadata = {} + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "no-such-model"}, + call_type="completion", + ) + + assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + assert exc_info.value.model == "no-such-model" + + +# --------------------------------------------------------------------------- +# dynamic_rate_limiter v3 — exercise just the raise path via the helper, not +# the full Redis/Lua stack. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_model_capacity_path_populates_provider(): + """ + The v3 dynamic limiter has three raise sites: model_saturation_check, + priority_model, and the fail-closed unknown-descriptor branch. We patch + the atomic increment to short-circuit straight into the model_saturation + path — that's the most common production trip — and confirm the + raised exception carries provider info. + """ + from litellm.types.router import ModelGroupInfo + + handler = _PROXY_DynamicRateLimitHandlerV3(internal_usage_cache=DualCache()) + handler.v3_limiter.atomic_check_and_increment_by_n = AsyncMock( + return_value={ + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "model_saturation_check", + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "rpm", + } + ], + } + ) + handler._create_priority_based_descriptors = MagicMock(return_value=[]) + handler._create_model_tracking_descriptor = MagicMock( + return_value={ + "key": "model_saturation_check", + "value": "gpt-4o-mini", + "rate_limit": {"requests_per_unit": 100}, + } + ) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-dyn-v3") + user_api_key_dict.metadata = {} + model_info = ModelGroupInfo(model_group="gpt-4o-mini", providers=["openai"]) + + with pytest.raises(HTTPException) as exc_info: + await handler._check_rate_limits( + model="gpt-4o-mini", + model_group_info=model_info, + user_api_key_dict=user_api_key_dict, + priority="default", + saturation=1.0, + data={"model": "gpt-4o-mini"}, + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert isinstance(exc, RateLimitError) + assert exc.llm_provider == "openai" + assert exc.model == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_unknown_descriptor_path_populates_provider(): + """Fail-closed unknown-descriptor branch must still attribute provider.""" + from litellm.types.router import ModelGroupInfo + + handler = _PROXY_DynamicRateLimitHandlerV3(internal_usage_cache=DualCache()) + handler.v3_limiter.atomic_check_and_increment_by_n = AsyncMock( + return_value={ + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "something_we_dont_handle", + "current_limit": 1, + "limit_remaining": 0, + "rate_limit_type": "rpm", + } + ], + } + ) + handler._create_priority_based_descriptors = MagicMock(return_value=[]) + handler._create_model_tracking_descriptor = MagicMock( + return_value={ + "key": "model_saturation_check", + "value": "gpt-4o-mini", + "rate_limit": {"requests_per_unit": 1}, + } + ) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-dyn-v3-unknown") + user_api_key_dict.metadata = {} + model_info = ModelGroupInfo(model_group="gpt-4o-mini", providers=["openai"]) + + with pytest.raises(HTTPException) as exc_info: + await handler._check_rate_limits( + model="gpt-4o-mini", + model_group_info=model_info, + user_api_key_dict=user_api_key_dict, + priority="default", + saturation=1.0, + data={"model": "gpt-4o-mini"}, + ) + + assert exc_info.value.llm_provider == "openai" + + +# --------------------------------------------------------------------------- +# batch_rate_limiter +# --------------------------------------------------------------------------- + + +def _batch_over_limit_response() -> dict: + return { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 10, + "limit_remaining": -5, + "rate_limit_type": "requests", + } + ], + } + + +@pytest.mark.asyncio +async def test_batch_rate_limiter_populates_provider(): + """ + batch_rate_limiter trips when the file's request/token count exceeds the + remaining window. The raise must thread `data["model"]` through the + helper. + """ + parallel_limiter = MagicMock() + parallel_limiter.window_size = 60 + parallel_limiter._create_rate_limit_descriptors = MagicMock( + return_value=[ + {"key": "key", "value": "v", "rate_limit": {"requests_per_unit": 10}} + ] + ) + parallel_limiter.atomic_check_and_increment_by_n = AsyncMock( + return_value=_batch_over_limit_response() + ) + + handler = _PROXY_BatchRateLimiter( + internal_usage_cache=InternalUsageCache(DualCache()), + parallel_request_limiter=parallel_limiter, + ) + + with pytest.raises(HTTPException) as exc_info: + await handler._check_and_increment_batch_counters( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-batch"), + data={"model": "gpt-4o-mini"}, + batch_usage=BatchFileUsage(total_tokens=100, request_count=15), + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert isinstance(exc, RateLimitError) + assert exc.llm_provider == "openai" + assert exc.model == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_batch_rate_limiter_unknown_model_falls_back(): + parallel_limiter = MagicMock() + parallel_limiter.window_size = 60 + parallel_limiter._create_rate_limit_descriptors = MagicMock( + return_value=[ + {"key": "key", "value": "v", "rate_limit": {"requests_per_unit": 10}} + ] + ) + parallel_limiter.atomic_check_and_increment_by_n = AsyncMock( + return_value=_batch_over_limit_response() + ) + + handler = _PROXY_BatchRateLimiter( + internal_usage_cache=InternalUsageCache(DualCache()), + parallel_request_limiter=parallel_limiter, + ) + + with pytest.raises(HTTPException) as exc_info: + await handler._check_and_increment_batch_counters( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-batch"), + data={"model": "fake-model-xyz"}, + batch_usage=BatchFileUsage(total_tokens=100, request_count=15), + ) + + assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + + +# --------------------------------------------------------------------------- +# max_budget_limiter +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_max_budget_limiter_populates_provider(): + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-budget", + user_id="user-1", + user_max_budget=10.0, + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=10.0), + ): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "gpt-4o-mini"}, + call_type="completion", + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert isinstance(exc, RateLimitError) + assert exc.llm_provider == "openai" + assert exc.model == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_max_budget_limiter_no_model_falls_back(): + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-budget", + user_id="user-1", + user_max_budget=10.0, + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=10.0), + ): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + assert exc_info.value.model == "" + + +# --------------------------------------------------------------------------- +# max_iterations_limiter +# --------------------------------------------------------------------------- + + +def _make_iter_agent(max_iterations: int) -> AgentResponse: + return AgentResponse( + agent_id="agent-iter", + agent_name="iter-agent", + litellm_params={"max_iterations": max_iterations}, + agent_card_params={"name": "iter-agent", "version": "1.0.0"}, + ) + + +@pytest.mark.asyncio +async def test_max_iterations_limiter_populates_provider(): + local_cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-iter", agent_id="agent-iter") + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_iter_agent(max_iterations=1) + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4o-mini", + "metadata": {"session_id": "session-iter-1"}, + }, + call_type="completion", + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4o-mini", + "metadata": {"session_id": "session-iter-1"}, + }, + call_type="completion", + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert isinstance(exc, RateLimitError) + assert exc.llm_provider == "openai" + assert exc.model == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_max_iterations_limiter_unknown_model_falls_back(): + local_cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-iter", agent_id="agent-iter") + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_iter_agent(max_iterations=1) + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "no-such-model", + "metadata": {"session_id": "session-iter-2"}, + }, + call_type="completion", + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "no-such-model", + "metadata": {"session_id": "session-iter-2"}, + }, + call_type="completion", + ) + + assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + + +# --------------------------------------------------------------------------- +# max_budget_per_session_limiter +# --------------------------------------------------------------------------- + + +def _make_session_budget_agent(max_budget: float) -> AgentResponse: + return AgentResponse( + agent_id="agent-session-budget", + agent_name="session-budget-agent", + litellm_params={"max_budget_per_session": max_budget}, + agent_card_params={"name": "session-budget-agent", "version": "1.0.0"}, + ) + + +@pytest.mark.asyncio +async def test_max_budget_per_session_limiter_populates_provider(): + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-session-budget", agent_id="agent-session-budget" + ) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_session_budget_agent( + max_budget=1.0 + ) + with patch.object( + handler, "_get_current_spend", new=AsyncMock(return_value=5.0) + ): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={ + "model": "anthropic/claude-3-5-sonnet", + "metadata": {"session_id": "session-budget-1"}, + }, + call_type="completion", + ) + + exc = exc_info.value + assert exc.status_code == 429 + assert isinstance(exc, RateLimitError) + assert exc.llm_provider == "anthropic" + + +@pytest.mark.asyncio +async def test_max_budget_per_session_limiter_unknown_model_falls_back(): + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-session-budget", agent_id="agent-session-budget" + ) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_session_budget_agent( + max_budget=1.0 + ) + with patch.object( + handler, "_get_current_spend", new=AsyncMock(return_value=5.0) + ): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={ + "model": "no-such-model", + "metadata": {"session_id": "session-budget-2"}, + }, + call_type="completion", + ) + + assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK + + +# --------------------------------------------------------------------------- +# Prometheus integration: failure metric reads exception.llm_provider +# via _get_exception_class_name. With the fix, this returns +# "Openai.RateLimitError" instead of plain "HTTPException" for proxy-side +# 429s on a known model. Pin that contract — that's what dashboards see. +# --------------------------------------------------------------------------- + + +def test_prometheus_exception_class_name_back_compat_for_proxy_rate_limit_error(): + """ + `_get_exception_class_name` deliberately returns the literal string + ``"HTTPException"`` for every ``ProxyRateLimitError`` instance so that + pre-existing dashboards / alerts (which key off the historical value) + keep working after the unified rate-limit error class landed in #27687. + + Provider attribution is now surfaced separately via the + ``rate_limit_category`` / ``rate_limit_type`` labels — this test pins + the back-compat shim itself. + """ + from litellm.integrations.prometheus import PrometheusLogger + + exc = ProxyRateLimitError( + detail="over limit", + model="gpt-4o-mini", + llm_provider="openai", + ) + assert PrometheusLogger._get_exception_class_name(exc) == "HTTPException" + + # Same back-compat path even when the resolver fell back to litellm_proxy. + exc_no_model = ProxyRateLimitError(detail="over limit") + assert PrometheusLogger._get_exception_class_name(exc_no_model) == "HTTPException" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-vv", "-x"]))