From f899d2f1ecda4fd892d675616024279c3f8bacff Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:45:56 +0000 Subject: [PATCH] fix(rate_limiter): attach priority rate limit headers on /v1/messages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 22 ++++------ .../hooks/parallel_request_limiter_v3.py | 41 ++++++++----------- .../hooks/test_dynamic_rate_limiter_v3.py | 40 ++++++++++++++++++ 3 files changed, 63 insertions(+), 40 deletions(-) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 4492f42782c..0d0ca56622b 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -32,6 +32,9 @@ from litellm.proxy.hooks.rate_limiter_utils import ( resolve_llm_provider_for_rate_limit, ) from litellm.proxy.utils import InternalUsageCache +from litellm.router_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, +) from litellm.types.router import ModelGroupInfo from litellm.types.utils import CallTypesLiteral @@ -657,21 +660,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): data=data, user_api_key_dict=user_api_key_dict, response=response ) - # Add additional priority-specific headers - if isinstance(response, ModelResponse): - priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) - - # Get existing additional headers - additional_headers: Final = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {} - - # Add priority information - additional_headers["x-litellm-priority"] = priority or "default" - additional_headers["x-litellm-rate-limiter-version"] = "v3" - - # Update response - if not hasattr(response, "_hidden_params"): - response._hidden_params = {} - response._hidden_params["additional_headers"] = additional_headers + priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) + additional_headers: Final = ensure_response_additional_headers(response) + additional_headers["x-litellm-priority"] = priority or "default" + additional_headers["x-litellm-rate-limiter-version"] = "v3" return response diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 94ef08782d9..e05a710bbc8 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -43,6 +43,9 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ( map_v3_rate_limit_type, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.router_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, +) from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject from litellm.types.utils import ( @@ -3292,36 +3295,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ Post-call hook to update rate limit headers in the response. + + Works for both object responses (``ModelResponse`` and friends) and + the dict responses returned by /v1/messages, which hold their + ``_hidden_params`` under a key instead of an attribute. """ try: - from pydantic import BaseModel - stash: Final = get_request_stash() litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None + if litellm_proxy_rate_limit_response is None: + return - if litellm_proxy_rate_limit_response is not None: - # Update response headers - if hasattr(response, "_hidden_params"): - _hidden_params = getattr(response, "_hidden_params") - else: - _hidden_params = None - - if _hidden_params is not None and ( - isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict) - ): - if isinstance(_hidden_params, BaseModel): - _hidden_params = _hidden_params.model_dump() - - _additional_headers: Final = self._merge_ratelimit_statuses_into_additional_headers( - additional_headers=_hidden_params.get("additional_headers", {}) or {}, - statuses=litellm_proxy_rate_limit_response["statuses"], - ) - - setattr( - response, - "_hidden_params", - {**_hidden_params, "additional_headers": _additional_headers}, - ) + additional_headers: Final = ensure_response_additional_headers(response) + additional_headers.update( + self._merge_ratelimit_statuses_into_additional_headers( + additional_headers=additional_headers, + statuses=litellm_proxy_rate_limit_response["statuses"], + ) + ) except Exception as e: verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e) diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 6c717d6f71c..0e4724abd84 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1863,3 +1863,43 @@ async def test_tpm_only_model_enforces_priority_and_model_capacity(): ) assert capacity_blocked.value.status_code == 429 assert "Model capacity reached" in capacity_blocked.value.detail["error"] + + +@pytest.mark.asyncio +async def test_priority_headers_attached_to_dict_response(): + """ + /v1/messages returns a plain dict rather than a ModelResponse, and used to be skipped + by the success hook, so priority rate limit headers never reached the client. + """ + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + RateLimitResponse, + RateLimitStatus, + get_or_create_request_stash, + ) + + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=75, + limit_remaining=74, + rate_limit_type="requests", + descriptor_key="priority_model", + ) + ], + ) + + user_api_key_dict = UserAPIKeyAuth(metadata={"priority": "premium"}) + response: dict = {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + await handler.async_post_call_success_hook( + data={"model": "claude"}, user_api_key_dict=user_api_key_dict, response=response + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + assert additional_headers["x-ratelimit-priority_model-limit-requests"] == 75 + assert additional_headers["x-ratelimit-priority_model-remaining-requests"] == 74 + assert additional_headers["x-litellm-priority"] == "premium" + assert additional_headers["x-litellm-rate-limiter-version"] == "v3"