From bf1cc8fa5b20f2ebc0ede2a05263962c1dc7b519 Mon Sep 17 00:00:00 2001 From: ecoyoung Date: Thu, 8 Oct 2026 09:52:24 +0800 Subject: [PATCH] fix(retries): honor provider retry-after headers --- litellm/main.py | 49 ++++++++++++++------- litellm/utils.py | 24 ++++++----- tests/unit/test_completion_retry_after.py | 52 +++++++++++++++++++++++ 3 files changed, 99 insertions(+), 26 deletions(-) create mode 100644 tests/unit/test_completion_retry_after.py diff --git a/litellm/main.py b/litellm/main.py index 37b290f51ba..8c747192ae8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -81,6 +81,7 @@ from litellm.litellm_core_utils.chat_completion_agentic_loop import ( ) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.dd_tracing import tracer +from litellm.litellm_core_utils.exception_mapping_utils import _get_response_headers from litellm.litellm_core_utils.get_litellm_params import ( AWS_CREDENTIAL_KWARGS_KEYS, OPTIONAL_KWARGS_KEYS, @@ -6107,6 +6108,18 @@ def completion( ) +def _retry_after_wait(retry_state: Any, fallback_wait: Callable[[Any], float]) -> float: + """Use a provider retry hint when available, otherwise keep the configured backoff.""" + if retry_state.outcome is not None: + exception = retry_state.outcome.exception() + response_headers = _get_response_headers(exception) + retry_after = litellm.utils._get_retry_after_from_exception_header(response_headers) + if retry_after > 0: + return min(retry_after, 60) + + return fallback_wait(retry_state) + + def completion_with_retries(*args, **kwargs): """ Executes a litellm.completion() with 3 retries @@ -6122,14 +6135,16 @@ def completion_with_retries(*args, **kwargs): kwargs["num_retries"] = 0 retry_strategy: Final[RetryStrategy] = kwargs.pop("retry_strategy", "constant_retry") original_function: Final = kwargs.pop("original_function", completion) - if retry_strategy == "exponential_backoff_retry": - retryer = tenacity.Retrying( - wait=tenacity.wait_exponential(multiplier=1, max=10), - stop=tenacity.stop_after_attempt(num_retries), - reraise=True, - ) - else: - retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + fallback_wait = ( + tenacity.wait_exponential(multiplier=1, max=10) + if retry_strategy == "exponential_backoff_retry" + else tenacity.wait_none() + ) + retryer = tenacity.Retrying( + wait=partial(_retry_after_wait, fallback_wait=fallback_wait), + stop=tenacity.stop_after_attempt(num_retries), + reraise=True, + ) return retryer(original_function, *args, **kwargs) @@ -6148,14 +6163,16 @@ async def acompletion_with_retries(*args, **kwargs): kwargs["num_retries"] = 0 retry_strategy: Final[RetryStrategy] = kwargs.pop("retry_strategy", "constant_retry") original_function: Final = kwargs.pop("original_function", completion) - if retry_strategy == "exponential_backoff_retry": - retryer = tenacity.AsyncRetrying( - wait=tenacity.wait_exponential(multiplier=1, max=10), - stop=tenacity.stop_after_attempt(num_retries), - reraise=True, - ) - else: - retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + fallback_wait = ( + tenacity.wait_exponential(multiplier=1, max=10) + if retry_strategy == "exponential_backoff_retry" + else tenacity.wait_none() + ) + retryer = tenacity.AsyncRetrying( + wait=partial(_retry_after_wait, fallback_wait=fallback_wait), + stop=tenacity.stop_after_attempt(num_retries), + reraise=True, + ) return await retryer(original_function, *args, **kwargs) diff --git a/litellm/utils.py b/litellm/utils.py index 72ce14912a9..ca37f94c9b5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7248,16 +7248,20 @@ def _get_retry_after_from_exception_header( # ". See https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After#syntax for # details. if response_headers is not None: - retry_header: Final[str] = response_headers.get("retry-after") - try: - retry_after = int(retry_header) - except Exception: - retry_date_tuple: Final = email.utils.parsedate_tz(retry_header) - if retry_date_tuple is None: - retry_after = -1 - else: - retry_date: Final = email.utils.mktime_tz(retry_date_tuple) - retry_after = int(retry_date - time.time()) + retry_header: Final[str | None] = response_headers.get("retry-after") + retry_after_ms_header: Final[str | None] = response_headers.get("retry-after-ms") + if retry_header is None and retry_after_ms_header is not None: + retry_after = float(retry_after_ms_header) / 1000 + else: + try: + retry_after = int(retry_header) + except Exception: + retry_date_tuple: Final = email.utils.parsedate_tz(retry_header) + if retry_date_tuple is None: + retry_after = -1 + else: + retry_date: Final = email.utils.mktime_tz(retry_date_tuple) + retry_after = int(retry_date - time.time()) else: retry_after = -1 diff --git a/tests/unit/test_completion_retry_after.py b/tests/unit/test_completion_retry_after.py new file mode 100644 index 00000000000..13af6601e07 --- /dev/null +++ b/tests/unit/test_completion_retry_after.py @@ -0,0 +1,52 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +import httpx +import pytest + +import litellm +from litellm.main import _retry_after_wait + + +def _retry_state(exception: Exception) -> SimpleNamespace: + return SimpleNamespace(outcome=SimpleNamespace(exception=lambda: exception)) + + +def _rate_limit_error(headers: dict[str, str]) -> litellm.RateLimitError: + response = httpx.Response( + status_code=429, + headers=headers, + request=httpx.Request("POST", "https://example.invalid/v1"), + ) + return litellm.RateLimitError( + message="rate limited", + model="test-model", + llm_provider="test-provider", + response=response, + ) + + +@pytest.mark.parametrize( + ("headers", "expected_wait"), + [ + ({"retry-after": "20"}, 20), + ({"retry-after-ms": "2500"}, 2.5), + ({"retry-after": "120"}, 60), + ], +) +def test_retry_after_wait_uses_and_caps_provider_hint(headers: dict[str, str], expected_wait: float) -> None: + fallback_wait = Mock(return_value=1.0) + + wait = _retry_after_wait(_retry_state(_rate_limit_error(headers)), fallback_wait) + + assert wait == expected_wait + fallback_wait.assert_not_called() + + +def test_retry_after_wait_preserves_fallback_strategy_without_hint() -> None: + fallback_wait = Mock(return_value=3.0) + + wait = _retry_after_wait(_retry_state(RuntimeError("request failed")), fallback_wait) + + assert wait == 3.0 + fallback_wait.assert_called_once()