From bf1cc8fa5b20f2ebc0ede2a05263962c1dc7b519 Mon Sep 17 00:00:00 2001 From: ecoyoung Date: Thu, 8 Oct 2026 09:52:24 +0800 Subject: [PATCH 1/6] 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() From cbac6fc6cb6953679fffb460ab242a5cfe69c760 Mon Sep 17 00:00:00 2001 From: ecoyoung Date: Thu, 8 Oct 2026 10:09:44 +0800 Subject: [PATCH 2/6] fix(retries): address retry-after review feedback --- litellm/main.py | 18 +++-- litellm/utils.py | 4 +- tests/unit/test_completion_retry_after.py | 52 ------------ tests/unit/test_main.py | 97 ++++++++++++++++++++++- 4 files changed, 109 insertions(+), 62 deletions(-) delete mode 100644 tests/unit/test_completion_retry_after.py diff --git a/litellm/main.py b/litellm/main.py index 8c747192ae8..38331e37f5d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -33,6 +33,7 @@ from litellm._uuid import uuid if TYPE_CHECKING: from aiohttp import ClientSession + from tenacity import RetryCallState import dotenv import httpx @@ -6108,14 +6109,15 @@ def completion( ) -def _retry_after_wait(retry_state: Any, fallback_wait: Callable[[Any], float]) -> float: +def _retry_after_wait(retry_state: "RetryCallState", fallback_wait: Callable[["RetryCallState"], 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) + exception: Final = retry_state.outcome.exception() + if isinstance(exception, Exception): + response_headers: Final = _get_response_headers(exception) + retry_after: Final = litellm.utils._get_retry_after_from_exception_header(response_headers) + if retry_after is not None and retry_after > 0: + return min(retry_after, 60) return fallback_wait(retry_state) @@ -6135,7 +6137,7 @@ 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) - fallback_wait = ( + fallback_wait: Final = ( tenacity.wait_exponential(multiplier=1, max=10) if retry_strategy == "exponential_backoff_retry" else tenacity.wait_none() @@ -6163,7 +6165,7 @@ 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) - fallback_wait = ( + fallback_wait: Final = ( tenacity.wait_exponential(multiplier=1, max=10) if retry_strategy == "exponential_backoff_retry" else tenacity.wait_none() diff --git a/litellm/utils.py b/litellm/utils.py index ca37f94c9b5..03ef79b69cf 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7252,7 +7252,7 @@ def _get_retry_after_from_exception_header( 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: + elif retry_header is not None: try: retry_after = int(retry_header) except Exception: @@ -7262,6 +7262,8 @@ def _get_retry_after_from_exception_header( else: retry_date: Final = email.utils.mktime_tz(retry_date_tuple) retry_after = int(retry_date - time.time()) + else: + retry_after = -1 else: retry_after = -1 diff --git a/tests/unit/test_completion_retry_after.py b/tests/unit/test_completion_retry_after.py deleted file mode 100644 index 13af6601e07..00000000000 --- a/tests/unit/test_completion_retry_after.py +++ /dev/null @@ -1,52 +0,0 @@ -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() diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 2d0d8740e12..d13123f5ab2 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -12,8 +12,9 @@ from dataclasses import dataclass from datetime import datetime from importlib import import_module from pathlib import Path +from types import SimpleNamespace from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest @@ -2573,6 +2574,100 @@ def throw_retryable_error(*_, **__): raise RuntimeError("BOOM") +def _retry_state(exception: BaseException) -> 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 = litellm_main._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 = litellm_main._retry_after_wait(_retry_state(RuntimeError("request failed")), fallback_wait) + + assert wait == 3.0 + fallback_wait.assert_called_once() + + +def test_retry_after_wait_falls_back_for_invalid_provider_hint() -> None: + fallback_wait = Mock(return_value=3.0) + + wait = litellm_main._retry_after_wait( + _retry_state(_rate_limit_error({"retry-after-ms": "invalid"})), fallback_wait + ) + + assert wait == 3.0 + fallback_wait.assert_called_once() + + +def test_completion_with_retries_retries_after_provider_error() -> None: + calls: Final[list[int]] = [] + + def original_function(*args: object, **kwargs: object) -> str: + calls.append(1) + if len(calls) == 1: + raise _rate_limit_error({"retry-after": "0"}) + return "ok" + + result = litellm_main.completion_with_retries( + original_function=original_function, + num_retries=2, + retry_strategy="constant_retry", + ) + + assert result == "ok" + assert len(calls) == 2 + + +@pytest.mark.asyncio +async def test_acompletion_with_retries_retries_after_provider_error() -> None: + calls: Final[list[int]] = [] + + async def original_function(*args: object, **kwargs: object) -> str: + calls.append(1) + if len(calls) == 1: + raise _rate_limit_error({"retry-after": "0"}) + return "ok" + + result = await litellm_main.acompletion_with_retries( + original_function=original_function, + num_retries=2, + retry_strategy="constant_retry", + ) + + assert result == "ok" + assert len(calls) == 2 + + @pytest.mark.asyncio async def test_retrying() -> None: litellm.num_retries = 10 From 9b17250660aba4fcddce5d8c92608be2b0a31663 Mon Sep 17 00:00:00 2001 From: ecoyoung Date: Thu, 8 Oct 2026 10:10:13 +0800 Subject: [PATCH 3/6] fix(retries): preserve fallback on malformed headers --- litellm/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 03ef79b69cf..62d89ff979f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7270,7 +7270,7 @@ def _get_retry_after_from_exception_header( return retry_after except Exception: - retry_after = -1 + return -1 def calculate_retry_after( From e8fb8717e526553458c8d5953d2d0eb0ae06268e Mon Sep 17 00:00:00 2001 From: ecoyoung Date: Thu, 8 Oct 2026 10:40:01 +0800 Subject: [PATCH 4/6] test(retries): cover retry-after branches --- litellm/utils.py | 2 +- tests/unit/test_main.py | 59 ++++++++++++++++++++++++++++++++++++++--- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 62d89ff979f..49319b8e0b4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7235,7 +7235,7 @@ _should_retry = should_retry def _get_retry_after_from_exception_header( response_headers: httpx.Headers | None = None, -): +) -> float | int: """ Reimplementation of openai's calculate retry after, since that one can't be imported. https://github.com/openai/openai-python/blob/af67cfab4210d8e497c05390ce14f39105c77519/src/openai/_base_client.py#L631 diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index d13123f5ab2..2d6b2a02a00 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1,11 +1,13 @@ import asyncio import base64 import contextlib +import email.utils import copy import io import json import logging import os +import time import urllib.parse from collections.abc import Mapping from dataclasses import dataclass @@ -2629,7 +2631,55 @@ def test_retry_after_wait_falls_back_for_invalid_provider_hint() -> None: fallback_wait.assert_called_once() -def test_completion_with_retries_retries_after_provider_error() -> None: +def test_retry_after_wait_falls_back_without_outcome_or_for_base_exception() -> None: + fallback_wait = Mock(return_value=3.0) + + no_outcome_wait = litellm_main._retry_after_wait(SimpleNamespace(outcome=None), fallback_wait) + base_exception_wait = litellm_main._retry_after_wait(_retry_state(KeyboardInterrupt()), fallback_wait) + + assert no_outcome_wait == 3.0 + assert base_exception_wait == 3.0 + assert fallback_wait.call_count == 2 + + +def test_retry_after_wait_parses_http_date_header() -> None: + retry_after = email.utils.formatdate(time.time() + 30, usegmt=True) + fallback_wait = Mock(return_value=1.0) + + wait = litellm_main._retry_after_wait(_retry_state(_rate_limit_error({"retry-after": retry_after})), fallback_wait) + + assert 0 < wait <= 30 + fallback_wait.assert_not_called() + + +def test_retry_after_wait_falls_back_for_invalid_retry_after_header() -> None: + fallback_wait = Mock(return_value=3.0) + + wait = litellm_main._retry_after_wait( + _retry_state(_rate_limit_error({"retry-after": "invalid"})), fallback_wait + ) + + assert wait == 3.0 + fallback_wait.assert_called_once() + + +def test_retry_after_wait_falls_back_without_retry_header() -> None: + fallback_wait = Mock(return_value=3.0) + + wait = litellm_main._retry_after_wait(_retry_state(_rate_limit_error({})), fallback_wait) + + assert wait == 3.0 + fallback_wait.assert_called_once() + + +def test_retry_after_parser_handles_empty_headers() -> None: + retry_after = litellm.utils._get_retry_after_from_exception_header(httpx.Headers()) + + assert retry_after == -1 + + +@pytest.mark.parametrize("retry_strategy", ["constant_retry", "exponential_backoff_retry"]) +def test_completion_with_retries_retries_after_provider_error(retry_strategy: str) -> None: calls: Final[list[int]] = [] def original_function(*args: object, **kwargs: object) -> str: @@ -2641,7 +2691,7 @@ def test_completion_with_retries_retries_after_provider_error() -> None: result = litellm_main.completion_with_retries( original_function=original_function, num_retries=2, - retry_strategy="constant_retry", + retry_strategy=retry_strategy, ) assert result == "ok" @@ -2649,7 +2699,8 @@ def test_completion_with_retries_retries_after_provider_error() -> None: @pytest.mark.asyncio -async def test_acompletion_with_retries_retries_after_provider_error() -> None: +@pytest.mark.parametrize("retry_strategy", ["constant_retry", "exponential_backoff_retry"]) +async def test_acompletion_with_retries_retries_after_provider_error(retry_strategy: str) -> None: calls: Final[list[int]] = [] async def original_function(*args: object, **kwargs: object) -> str: @@ -2661,7 +2712,7 @@ async def test_acompletion_with_retries_retries_after_provider_error() -> None: result = await litellm_main.acompletion_with_retries( original_function=original_function, num_retries=2, - retry_strategy="constant_retry", + retry_strategy=retry_strategy, ) assert result == "ok" From af92f563a4fc6b70cf05e78640ccc745e279a2cf Mon Sep 17 00:00:00 2001 From: ecoyoung Date: Thu, 8 Oct 2026 10:40:12 +0800 Subject: [PATCH 5/6] style(tests): order retry coverage imports --- tests/unit/test_main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 2d6b2a02a00..946755f15ee 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1,8 +1,8 @@ import asyncio import base64 import contextlib -import email.utils import copy +import email.utils import io import json import logging From 5c5291ebea4af82a99ffb6f780cd50fd3e3eac64 Mon Sep 17 00:00:00 2001 From: ecoyoung Date: Thu, 8 Oct 2026 10:42:12 +0800 Subject: [PATCH 6/6] test(retries): make date coverage deterministic --- tests/unit/test_main.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 946755f15ee..ebc6a06137b 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -7,7 +7,6 @@ import io import json import logging import os -import time import urllib.parse from collections.abc import Mapping from dataclasses import dataclass @@ -2643,12 +2642,15 @@ def test_retry_after_wait_falls_back_without_outcome_or_for_base_exception() -> def test_retry_after_wait_parses_http_date_header() -> None: - retry_after = email.utils.formatdate(time.time() + 30, usegmt=True) + retry_after = email.utils.formatdate(1_000_030, usegmt=True) fallback_wait = Mock(return_value=1.0) - wait = litellm_main._retry_after_wait(_retry_state(_rate_limit_error({"retry-after": retry_after})), fallback_wait) + with patch.object(litellm.utils.time, "time", return_value=1_000_000): + wait = litellm_main._retry_after_wait( + _retry_state(_rate_limit_error({"retry-after": retry_after})), fallback_wait + ) - assert 0 < wait <= 30 + assert wait == 30 fallback_wait.assert_not_called()