diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 16202321709..58164fa0f40 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -338,6 +338,7 @@ class Timeout(openai.APITimeoutError): num_retries: int | None = None, headers: dict | None = None, exception_status_code: int | None = None, + response: httpx.Response | None = None, ): request: Final = httpx.Request( method="POST", @@ -352,6 +353,12 @@ class Timeout(openai.APITimeoutError): self.max_retries = max_retries self.num_retries = num_retries self.headers = headers + # The upstream response, when the timeout came from a provider reply + # rather than a client-side deadline. Retry and cooldown logic reads + # `retry-after` off it, and the proxy prefixes its headers before + # returning them, matching every other mapped provider exception. + if response is not None: + self.response = response # custom function to convert to str def __str__(self): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 03f839c5345..82708d412c9 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -843,15 +843,6 @@ def _get_bedrock_mantle_context_window_message(error_str: str) -> str | None: return f"prompt is too long: {prompt_tokens} tokens > {max_tokens} maximum" -def _bedrock_timeout_headers(original_exception: object) -> dict | None: - from litellm.litellm_core_utils.llm_response_utils.get_headers import get_response_headers - - provider_headers = getattr(getattr(original_exception, "response", None), "headers", None) - if not provider_headers: - return None - return get_response_headers(provider_headers) - - def _map_bedrock_exception( *, model: str, @@ -935,7 +926,7 @@ def _map_bedrock_exception( message=f"BedrockException: Timeout Error - {error_str}", model=model, llm_provider="bedrock", - headers=_bedrock_timeout_headers(original_exception), + response=getattr(original_exception, "response", None), ) elif "Could not process image" in error_str: raise litellm.InternalServerError( @@ -979,7 +970,7 @@ def _map_bedrock_exception( model=model, llm_provider=custom_llm_provider, litellm_debug_info=extra_information, - headers=_bedrock_timeout_headers(original_exception), + response=getattr(original_exception, "response", None), ) elif original_exception.status_code == 422: raise BadRequestError( @@ -1012,7 +1003,7 @@ def _map_bedrock_exception( llm_provider=custom_llm_provider, litellm_debug_info=extra_information, exception_status_code=original_exception.status_code, - headers=_bedrock_timeout_headers(original_exception), + response=getattr(original_exception, "response", None), ) diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 730041c73e2..42d3df76902 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -9,6 +9,7 @@ import litellm from litellm.litellm_core_utils.exception_mapping_utils import ( ExceptionCheckers, _get_body_error_code, + _get_response_headers, exception_type, extract_and_raise_litellm_exception, ) @@ -1349,10 +1350,10 @@ def test_bedrock_classified_errors_preserve_provider_response_headers( ], ) def test_bedrock_timeout_mapping_preserves_provider_headers(status_code, provider_message): - """Timeout takes no response argument, so the headers have to ride on the exception itself. + """A mapped bedrock timeout keeps the upstream response, like every other mapped bedrock error. - The proxy reads e.headers before e.response.headers, so they arrive already - llm_provider-prefixed rather than as raw upstream names. + The proxy prefixes those headers on the way out, while retry and cooldown + logic still reads the raw retry-after off the response. """ provider_response = httpx.Response( status_code=status_code, @@ -1376,6 +1377,35 @@ def test_bedrock_timeout_mapping_preserves_provider_headers(status_code, provide extra_kwargs={}, ) - headers = exc_info.value.headers or {} - assert headers["llm_provider-x-amzn-requestid"] == "req-timeout" - assert "set-cookie" not in headers + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-timeout" + assert exc_info.value.headers is None + + +@pytest.mark.parametrize("status_code", [504, 408]) +def test_bedrock_timeout_mapping_keeps_retry_after_readable(status_code): + """Cooldown and retry timing read retry-after through _get_response_headers.""" + provider_response = httpx.Response( + status_code=status_code, + headers={"x-amzn-RequestId": "req-retry-after", "retry-after": "7"}, + text='{"message":"Bedrock did not answer in time"}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=status_code, + message='{"message":"Bedrock did not answer in time"}', + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(litellm.Timeout) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + exception_headers = _get_response_headers(original_exception=exc_info.value) + assert exception_headers is not None + assert litellm.utils._get_retry_after_from_exception_header(response_headers=exception_headers) == 7