Merge ac45a34f2b into f82e4c80ac

litellm/litellm_core_utils/litellm_logging.py

@@ -796,7 +796,8 @@ class Logging(LiteLLMLoggingBaseClass):
         # ids leaking into a different, later request on the same thread. Sync
         # support is deferred to a follow-up PR with its own safe-restore
         # mechanism; async calls (the proxy's only call path) are unaffected.
-        if supports_correlation_logging and litellm.request_correlation_in_logs:
+        self._stamps_correlation_context: bool = supports_correlation_logging and litellm.request_correlation_in_logs
+        if self._stamps_correlation_context:
             set_trace_id(self.litellm_trace_id)
             set_session_id(self.litellm_session_id)
         # set_trace_id()/set_session_id() sanitize (strip control chars, bound
@@ -3578,6 +3579,11 @@ class Logging(LiteLLMLoggingBaseClass):
         trace_id_var.set(self._pre_call_trace_id)
         session_id_var.set(self._pre_call_session_id)
 
+    def _stamp_correlation_context(self) -> None:
+        if self._stamps_correlation_context:
+            trace_id_var.set(self._own_trace_id)
+            session_id_var.set(self._own_session_id)
+
     def _restore_correlation_context_if_unclaimed(self) -> None:
         """Guarded variant for __del__-triggered cleanup only.
 

litellm/main.py

@@ -14,12 +14,13 @@ import contextvars
 import datetime
 import inspect
 import json
+import math
 import os
 import random
 import sys
 import time
 import traceback
-from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Mapping, Sequence
+from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence
 from concurrent import futures
 from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
 from copy import deepcopy
@@ -72,6 +73,7 @@ from litellm.constants import (
 )
 from litellm.exceptions import LiteLLMUnknownProvider
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure
 from litellm.litellm_core_utils.allowed_tools_tool_choice import (
     raise_if_allowed_tools_tool_choice_unsupported,
 )
@@ -86,6 +88,9 @@ from litellm.litellm_core_utils.chat_completion_agentic_loop import (
 )
 from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind
 from litellm.litellm_core_utils.dd_tracing import tracer
+from litellm.litellm_core_utils.exception_mapping_utils import (
+    _get_response_headers,  # pyright: ignore[reportPrivateUsage]  # shared provider header extractor
+)
 from litellm.litellm_core_utils.get_litellm_params import (
     AGENTIC_LOOP_CONTROL_KWARGS_KEYS,
     AWS_CREDENTIAL_KWARGS_KEYS,
@@ -103,6 +108,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import (
 from litellm.litellm_core_utils.get_provider_specific_headers import (
     ProviderSpecificHeaderUtils,
 )
+from litellm.litellm_core_utils.get_retry_from_policy import get_num_retries_from_retry_policy
 from litellm.litellm_core_utils.health_check_utils import (
     _filter_model_params,
     create_health_check_response,
@@ -166,7 +172,7 @@ from litellm.types.completion import (
 )
 from litellm.types.litellm_params import ControlOptions, RetryStrategy
 from litellm.types.llms.openai import AllMessageValues
-from litellm.types.router import GenericLiteLLMParams, validate_litellm_params_strictly
+from litellm.types.router import GenericLiteLLMParams, RetryPolicy, validate_litellm_params_strictly
 from litellm.types.utils import (
     AgenticLoopParams,
     BaseUrlRoutingHint,
@@ -188,8 +194,10 @@ from litellm.utils import (
     TextCompletionResponse,
     TranscriptionResponse,
     Usage,
+    _get_retry_after_from_exception_header,  # pyright: ignore[reportPrivateUsage]  # shared Retry-After parser
     add_provider_specific_params_to_optional_params,
     async_mock_completion_streaming_obj,
+    calculate_retry_after,
     convert_to_model_response_object,
     create_pretrained_tokenizer,
     create_tokenizer,
@@ -207,6 +215,7 @@ from litellm.utils import (
     mock_completion_streaming_obj,
     pre_process_non_default_params,
     read_config_args,
+    should_retry,
     should_run_mock_completion,
     supports_httpx_timeout,
     supports_pdf_input,
@@ -6608,56 +6617,233 @@ def completion(
         )
 
 
-def completion_with_retries(*args, **kwargs):
-    """
-    Executes a litellm.completion() with 3 retries
-    """
+class _RetryOutcome(Protocol):
+    def exception(self) -> BaseException | None: ...
+
+
+class _RetryState(Protocol):
+    @property
+    def attempt_number(self) -> int: ...
+
+    @property
+    def outcome(self) -> _RetryOutcome | None: ...
+
+
+def _compute_retry_wait(
+    exc: BaseException | None,
+    retries_sent: int,
+    num_retries: int,
+    retry_strategy: str,
+) -> float:
+    headers: Final[httpx.Headers | None] = _get_response_headers(exc) if isinstance(exc, Exception) else None
+    if retry_strategy == "exponential_backoff_retry":
+        return float(
+            calculate_retry_after(
+                remaining_retries=max(0, num_retries - retries_sent),
+                max_retries=num_retries,
+                response_headers=headers,
+            )
+        )
+    retry_after: Final[int | None] = _get_retry_after_from_exception_header(headers)
+    return float(retry_after) if retry_after is not None and 0 < retry_after <= 60 else 0.0
+
+
+_NUM_RETRIES_ADAPTER: Final = TypeAdapter(float)
+_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy)
+
+
+def _build_retry_predicate(retry_policy: RetryPolicy | None) -> Callable[[BaseException], bool]:
+    """A RetryPolicy count decides for the errors it covers, and other errors retry only when transient."""
+
+    def _predicate(e: BaseException) -> bool:
+        if is_agentic_followup_failure(e):
+            return False
+        policy_retries: Final[int | None] = (
+            get_num_retries_from_retry_policy(exception=e, retry_policy=retry_policy)
+            if retry_policy is not None and isinstance(e, Exception)
+            else None
+        )
+        if policy_retries is not None:
+            return policy_retries > 0
+        status_code: Final[object] = getattr(e, "status_code", None)
+        if isinstance(status_code, int):
+            return should_retry(status_code)
+        return isinstance(e, (openai.APITimeoutError, openai.APIConnectionError))
+
+    return _predicate
+
+
+@dataclass(frozen=True, slots=True)
+class _RetryPlan:
+    num_retries: int
+    retry_policy: RetryPolicy | None
+    initial_exception: BaseException | None
+
+
+def _retry_plan_after_first_error(
+    first_error: Exception, num_retries: int, retry_policy: RetryPolicy
+) -> _RetryPlan | None:
+    """The retries a direct retry_policy grants after the first error, kept for the whole call like the router's."""
+    if is_agentic_followup_failure(first_error):
+        return None
+    policy_retries: Final[int | None] = get_num_retries_from_retry_policy(
+        exception=first_error, retry_policy=retry_policy
+    )
+    if policy_retries is not None:
+        if policy_retries <= 0:
+            return None
+        return _RetryPlan(num_retries=policy_retries, retry_policy=retry_policy, initial_exception=first_error)
+    if num_retries <= 1 or not _build_retry_predicate(None)(first_error):
+        return None
+    return _RetryPlan(num_retries=num_retries - 1, retry_policy=retry_policy, initial_exception=first_error)
+
+
+def _run_completion_retries(
+    original_function: Callable[..., object],
+    args: tuple[object, ...],
+    kwargs: Mapping[str, object],
+    plan: _RetryPlan,
+    retry_strategy: str,
+) -> object:
     try:
         import tenacity
     except Exception as e:
         raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}")
 
-    num_retries: Final = kwargs.pop("num_retries", 3)
-    # reset retries in .completion()
-    kwargs["max_retries"] = 0
-    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,
+    if plan.initial_exception is not None:
+        time.sleep(
+            _compute_retry_wait(
+                exc=plan.initial_exception,
+                retries_sent=0,
+                num_retries=plan.num_retries,
+                retry_strategy=retry_strategy,
+            )
         )
-    else:
-        retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True)
+    retryer: Final = tenacity.Retrying(
+        wait=_build_retry_wait_fn(
+            plan.num_retries, retry_strategy, first_attempt_is_a_retry=plan.initial_exception is not None
+        ),
+        stop=tenacity.stop_after_attempt(plan.num_retries),
+        retry=tenacity.retry_if_exception(_build_retry_predicate(plan.retry_policy)),
+        reraise=True,
+    )
     return retryer(original_function, *args, **kwargs)
 
 
-async def acompletion_with_retries(*args, **kwargs):
-    """
-    [DEPRECATED]. Use 'acompletion' or router.acompletion instead!
-    Executes a litellm.completion() with 3 retries
-    """
+async def _run_acompletion_retries(
+    original_function: Callable[..., Awaitable[object]],
+    args: tuple[object, ...],
+    kwargs: Mapping[str, object],
+    plan: _RetryPlan,
+    retry_strategy: str,
+) -> object:
     try:
         import tenacity
     except Exception as e:
         raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}")
 
-    num_retries: Final = kwargs.pop("num_retries", 3)
+    if plan.initial_exception is not None:
+        await asyncio.sleep(
+            _compute_retry_wait(
+                exc=plan.initial_exception,
+                retries_sent=0,
+                num_retries=plan.num_retries,
+                retry_strategy=retry_strategy,
+            )
+        )
+    retryer: Final = tenacity.AsyncRetrying(
+        wait=_build_retry_wait_fn(
+            plan.num_retries, retry_strategy, first_attempt_is_a_retry=plan.initial_exception is not None
+        ),
+        stop=tenacity.stop_after_attempt(plan.num_retries),
+        retry=tenacity.retry_if_exception(_build_retry_predicate(plan.retry_policy)),
+        reraise=True,
+    )
+    return await retryer(original_function, *args, **kwargs)
+
+
+def _build_retry_wait_fn(
+    num_retries: int, retry_strategy: str, first_attempt_is_a_retry: bool
+) -> Callable[[_RetryState], float]:
+    def _wait_fn(retry_state: _RetryState) -> float:
+        outcome: Final[_RetryOutcome | None] = retry_state.outcome
+        exc: Final[BaseException | None] = outcome.exception() if outcome is not None else None
+        return _compute_retry_wait(
+            exc=exc,
+            retries_sent=retry_state.attempt_number if first_attempt_is_a_retry else retry_state.attempt_number - 1,
+            num_retries=num_retries,
+            retry_strategy=retry_strategy,
+        )
+
+    return _wait_fn
+
+
+def completion_with_retries(
+    *args: object, original_function: Callable[..., object] = completion, **kwargs: object
+) -> object:
+    """
+    Executes a litellm.completion() with retries
+    """
+    num_retries: Final[int] = math.ceil(_NUM_RETRIES_ADAPTER.validate_python(kwargs.pop("num_retries", 3)))
+    # reset retries in .completion()
     kwargs["max_retries"] = 0
     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,
+    retry_strategy: Final[str] = (
+        "constant_retry" if kwargs.pop("retry_strategy", None) == "constant_retry" else "exponential_backoff_retry"
+    )
+    initial_exc: Final[object] = kwargs.pop("_initial_exception", None)
+    retry_policy: Final[object] = kwargs.pop("retry_policy", None)
+    validated_retry_policy: Final[RetryPolicy | None] = (
+        _RETRY_POLICY_ADAPTER.validate_python(retry_policy) if retry_policy else None
+    )
+    if isinstance(initial_exc, BaseException) or validated_retry_policy is None:
+        initial_plan: Final = _RetryPlan(
+            num_retries=num_retries,
+            retry_policy=validated_retry_policy,
+            initial_exception=initial_exc if isinstance(initial_exc, BaseException) else None,
         )
-    else:
-        retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True)
-    return await retryer(original_function, *args, **kwargs)
+        return _run_completion_retries(original_function, args, kwargs, initial_plan, retry_strategy)
+    try:
+        return original_function(*args, **kwargs)
+    except Exception as first_error:
+        plan: Final = _retry_plan_after_first_error(first_error, num_retries, validated_retry_policy)
+        if plan is None:
+            raise
+        return _run_completion_retries(original_function, args, kwargs, plan, retry_strategy)
+
+
+async def acompletion_with_retries(
+    *args: object, original_function: Callable[..., Awaitable[object]] = acompletion, **kwargs: object
+) -> object:
+    """
+    [DEPRECATED]. Use 'acompletion' or router.acompletion instead!
+    Executes a litellm.acompletion() with retries
+    """
+    num_retries: Final[int] = math.ceil(_NUM_RETRIES_ADAPTER.validate_python(kwargs.pop("num_retries", 3)))
+    kwargs["max_retries"] = 0
+    kwargs["num_retries"] = 0
+    retry_strategy: Final[str] = (
+        "constant_retry" if kwargs.pop("retry_strategy", None) == "constant_retry" else "exponential_backoff_retry"
+    )
+    initial_exc: Final[object] = kwargs.pop("_initial_exception", None)
+    retry_policy: Final[object] = kwargs.pop("retry_policy", None)
+    validated_retry_policy: Final[RetryPolicy | None] = (
+        _RETRY_POLICY_ADAPTER.validate_python(retry_policy) if retry_policy else None
+    )
+    if isinstance(initial_exc, BaseException) or validated_retry_policy is None:
+        initial_plan: Final = _RetryPlan(
+            num_retries=num_retries,
+            retry_policy=validated_retry_policy,
+            initial_exception=initial_exc if isinstance(initial_exc, BaseException) else None,
+        )
+        return await _run_acompletion_retries(original_function, args, kwargs, initial_plan, retry_strategy)
+    try:
+        return await original_function(*args, **kwargs)
+    except Exception as first_error:
+        plan: Final = _retry_plan_after_first_error(first_error, num_retries, validated_retry_policy)
+        if plan is None:
+            raise
+        return await _run_acompletion_retries(original_function, args, kwargs, plan, retry_strategy)
 
 
 def responses_with_retries(*args, **kwargs):

litellm/utils.py

@@ -1090,6 +1090,12 @@ def _restore_correlation_context_if_supported(logging_obj: object) -> None:
         restore()
 
 
+def _stamp_correlation_context_if_supported(logging_obj: object) -> None:
+    stamp: Final = getattr(logging_obj, "_stamp_correlation_context", None)
+    if stamp is not None:
+        stamp()
+
+
 def _is_streaming_response_for_correlation(result: object) -> bool:
     """True if `result` is a lazy stream wrapper rather than an already-complete response.
 
@@ -1377,20 +1383,15 @@ def function_setup(
         dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None
         dynamic_async_failure_callbacks: Final[list[str | Callable[..., object] | CustomLogger] | None] = None
         if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list):
-            removed_async_items = []
-            for index, callback in enumerate(kwargs["success_callback"]):
-                if coroutine_checker.is_async_callable(callback) or callback == "dynamodb" or callback == "s3":
-                    if dynamic_async_success_callbacks is not None and isinstance(
-                        dynamic_async_success_callbacks, list
-                    ):
-                        dynamic_async_success_callbacks.append(callback)
-                    else:
-                        dynamic_async_success_callbacks = [callback]
-                    removed_async_items.append(index)
-            # Pop the async items from success_callback in reverse order to avoid index issues
-            for index in reversed(removed_async_items):
-                kwargs["success_callback"].pop(index)
-            dynamic_success_callbacks = kwargs.pop("success_callback")
+            requested_success_callbacks: Final[list[str | Callable[..., object] | CustomLogger]] = kwargs.pop(
+                "success_callback"
+            )
+            dynamic_async_success_callbacks = [
+                callback for callback in requested_success_callbacks if _runs_as_async_success_callback(callback)
+            ] or None
+            dynamic_success_callbacks = [
+                callback for callback in requested_success_callbacks if not _runs_as_async_success_callback(callback)
+            ]
         if kwargs.get("failure_callback", None) is not None and isinstance(kwargs["failure_callback"], list):
             dynamic_failure_callbacks = kwargs.pop("failure_callback")
 
@@ -1730,32 +1731,39 @@ async def _client_async_logging_helper(
         )
 
 
-def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tuple[int | None, dict[str, Any]]:
+def _get_wrapper_num_retries(
+    kwargs: dict[str, Any], exception: Exception
+) -> tuple[int | None | Exception, dict[str, Any]]:
     """
     Get the number of retries from the kwargs and the retry policy.
-    Used for the wrapper functions.
+    Used for the wrapper functions. A malformed retry policy comes back as its parse error in place of the count.
     """
 
     num_retries = kwargs.get("num_retries", None)
     if num_retries is None:
         num_retries = litellm.num_retries
-    if kwargs.get("retry_policy", None):
-        get_num_retries_from_retry_policy: Final[Callable[..., int | None]] = (
-            litellm_utils.get_num_retries_from_retry_policy
-        )
-        reset_retry_policy: Final = litellm_utils.reset_retry_policy
-        retry_policy_num_retries: Final[int | None] = get_num_retries_from_retry_policy(
-            exception=exception,
-            retry_policy=kwargs.get("retry_policy"),
-        )
-        kwargs["retry_policy"] = reset_retry_policy()
-        if retry_policy_num_retries is not None:
-            num_retries = retry_policy_num_retries
-
-    return num_retries, kwargs
+    retry_policy: Final[object] = kwargs.pop("retry_policy", None)
+    if not retry_policy:
+        return num_retries, kwargs
+    kwargs["retry_policy"] = retry_policy
+    try:
+        retry_policy_num_retries: Final[int | None] = _get_retry_policy_num_retries(kwargs, exception)
+    except (AttributeError, TypeError, ValueError) as retry_policy_error:
+        return retry_policy_error, kwargs
+    if retry_policy_num_retries is None:
+        return num_retries, kwargs
+    return retry_policy_num_retries, kwargs
 
 
-def _get_wrapper_timeout(kwargs: dict[str, object]) -> object:
+def _get_retry_policy_num_retries(kwargs: Mapping[str, object], exception: Exception) -> int | None:
+    get_num_retries_from_retry_policy: Final[Callable[..., int | None]] = (
+        litellm_utils.get_num_retries_from_retry_policy
+    )
+    retry_policy: Final[object] = kwargs.get("retry_policy")
+    return get_num_retries_from_retry_policy(exception=exception, retry_policy=retry_policy) if retry_policy else None
+
+
+def _get_wrapper_timeout(kwargs: Mapping[str, object]) -> object:
     from litellm.litellm_core_utils.litellm_logging import Logging
 
     if kwargs.get("stream") is not True:
@@ -1767,6 +1775,41 @@ def _get_wrapper_timeout(kwargs: dict[str, object]) -> object:
     return kwargs.get("timeout") if stream_timeout is None else stream_timeout
 
 
+def _add_retry_context(error: Exception, kwargs: Mapping[str, object]) -> None:
+    deployment_num_retries: Final = kwargs.get("num_retries")
+    if deployment_num_retries is not None:
+        setattr(error, "num_retries", deployment_num_retries)
+    setattr(error, "timeout", _get_wrapper_timeout(kwargs=kwargs))
+
+
+_RETRY_ATTEMPT_KWARG: Final = "_litellm_wrapper_retry_attempt"
+_RETRY_ATTEMPT_MARKER: Final = object()
+
+
+def _with_retry_policy_spent(kwargs: Mapping[str, object]) -> Mapping[str, object]:
+    if not kwargs.get("retry_policy"):
+        return {**kwargs}
+    return {**kwargs, "retry_policy": litellm_utils.reset_retry_policy()}
+
+
+async def _end_async_wrapper_failure(
+    logging_obj: LiteLLMLoggingObject,
+    error: Exception,
+    traceback_exception: str,
+    start_time: datetime.datetime,
+    end_time: datetime.datetime,
+    log_failure: bool,
+) -> None:
+    unbind_budget_reservation_from_callbacks(logging_obj.litellm_params)
+    if not log_failure:
+        return
+    _stamp_correlation_context_if_supported(logging_obj)
+    logging_obj.failure_handler(
+        error, traceback_exception, start_time, end_time, for_async_request=True
+    )  # DO NOT MAKE THREADED - router retry fallback relies on this!
+    await logging_obj.async_failure_handler(error, traceback_exception, start_time, end_time)
+
+
 def _runs_as_async_success_callback(callback: str | Callable[..., object] | CustomLogger) -> bool:
     get_coroutine_checker: Final = litellm_utils.get_coroutine_checker
     return get_coroutine_checker().is_async_callable(callback) or callback in ("dynamodb", "s3")
@@ -2043,6 +2086,17 @@ def _get_call_type_from_public_name(function_name: str) -> str:
             return function_name
 
 
+def _context_window_fallback_applies(
+    error: Exception, context_window_fallback_dict: object, model: str, *, is_router_call: bool
+) -> bool:
+    return (
+        not is_router_call
+        and isinstance(error, litellm.exceptions.ContextWindowExceededError)
+        and isinstance(context_window_fallback_dict, Mapping)
+        and model in context_window_fallback_dict
+    )
+
+
 def _is_litellm_router_call(kwargs: Mapping[str, object], *, is_async: bool) -> bool:
     """Router completion uses metadata. Async generic calls retry with litellm_metadata; sync generic calls need SDK retries."""
     metadata_buckets: Final = (
@@ -2290,48 +2344,50 @@ def client(original_function):
             return result
         except Exception as e:
             if call_type == CallTypes.completion.value:
-                num_retries = kwargs.get("num_retries", None) or litellm.num_retries or None
-                if kwargs.get("retry_policy", None):
-                    get_num_retries_from_retry_policy: Callable[..., int | None] = (
-                        litellm_utils.get_num_retries_from_retry_policy
-                    )
-                    reset_retry_policy = litellm_utils.reset_retry_policy
-                    num_retries = get_num_retries_from_retry_policy(
-                        exception=e,
-                        retry_policy=kwargs.get("retry_policy"),
-                    )
-                    kwargs["retry_policy"] = reset_retry_policy()  # prevent infinite loops
+                num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e)
+                if isinstance(num_retries, Exception):
+                    raise num_retries
                 litellm.num_retries = None  # set retries to None to prevent infinite loops
                 context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {})
 
                 is_completion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False)
+                context_window_fallback_applies: Final = isinstance(model, str) and _context_window_fallback_applies(
+                    e, context_window_fallback_dict, model, is_router_call=is_completion_litellm_router_call
+                )
                 if (
-                    num_retries and not is_completion_litellm_router_call and not is_agentic_followup_failure(e)
-                ):  # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying
-                    if (
-                        isinstance(e, openai.APIError)
-                        or isinstance(e, openai.Timeout)
-                        or isinstance(e, openai.APIConnectionError)
-                    ):
-                        kwargs["num_retries"] = num_retries
-                        return litellm.completion_with_retries(*args, **kwargs)
-                elif (
-                    isinstance(e, litellm.exceptions.ContextWindowExceededError)
-                    and context_window_fallback_dict
-                    and model in context_window_fallback_dict
+                    num_retries
                     and not is_completion_litellm_router_call
-                ):
-                    if len(args) > 0:
-                        args[0] = context_window_fallback_dict[model]
-                    else:
-                        kwargs["model"] = context_window_fallback_dict[model]
-                    return original_function(*args, **kwargs)
+                    and not context_window_fallback_applies
+                    and not is_agentic_followup_failure(e)
+                ):  # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying
+                    status_code: Final[object] = getattr(e, "status_code", None)
+                    if (
+                        _get_retry_policy_num_retries(kwargs, e) is not None
+                        or not isinstance(status_code, int)
+                        or should_retry(status_code)
+                    ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)):
+                        kwargs["num_retries"] = num_retries
+                        from litellm.chat_completions.dispatch import completion as public_completion
+
+                        return litellm.completion_with_retries(
+                            *args,
+                            original_function=public_completion,
+                            **{**kwargs, "_initial_exception": e},
+                        )
+                if context_window_fallback_applies:
+                    fallback_model: Final = context_window_fallback_dict[model]
+                    fallback_kwargs: Final = _with_retry_policy_spent(kwargs)
+                    if args:
+                        return original_function(fallback_model, *args[1:], **fallback_kwargs)
+                    return original_function(*args, **{**fallback_kwargs, "model": fallback_model})
             elif call_type == CallTypes.responses.value:
                 num_retries = kwargs.get("num_retries", None) or litellm.num_retries or None
                 if kwargs.get("retry_policy", None):
-                    get_num_retries_from_retry_policy = litellm_utils.get_num_retries_from_retry_policy
+                    get_responses_num_retries: Final[Callable[..., int | None]] = (
+                        litellm_utils.get_num_retries_from_retry_policy
+                    )
                     reset_retry_policy = litellm_utils.reset_retry_policy
-                    num_retries = get_num_retries_from_retry_policy(
+                    num_retries = get_responses_num_retries(
                         exception=e,
                         retry_policy=kwargs.get("retry_policy"),
                     )
@@ -2380,8 +2436,10 @@ def client(original_function):
         model: Final[str | None] = args[0] if len(args) > 0 else kwargs.get("model", None)
         is_completion_with_fallbacks: Final = kwargs.get("fallbacks") is not None
         kwargs.pop("_is_litellm_internal_call", None)  # discard if injected
+        is_wrapper_retry_attempt: Final = kwargs.pop(_RETRY_ATTEMPT_KWARG, None) is _RETRY_ATTEMPT_MARKER
         _is_litellm_internal_call: Final = is_internal_call.get()
         _deployment_call_end_time: datetime.datetime | None = None
+        request_kwargs: Final = {**kwargs}
 
         try:
             if logging_obj is None:
@@ -2389,6 +2447,8 @@ def client(original_function):
 
             # Type assertion: logging_obj is guaranteed to be non-None after function_setup
             assert logging_obj is not None, "logging_obj should not be None after function_setup"
+            if is_wrapper_retry_attempt:
+                _stamp_correlation_context_if_supported(logging_obj)
             if not _is_litellm_internal_call:
                 bind_budget_reservation_to_callbacks(logging_obj.litellm_params)
 
@@ -2614,62 +2674,100 @@ def client(original_function):
 
             return _replayed_when_run_unstreamed(result, transcription_runs_unstreamed)
         except Exception as e:
+            wrapper_num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e)
+            num_retries: Final = None if isinstance(wrapper_num_retries, Exception) else wrapper_num_retries
+            is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True)
+            retries_apply: Final = (
+                call_type == CallTypes.acompletion.value
+                and bool(num_retries)
+                and not is_acompletion_litellm_router_call
+                and not isinstance(e, ImportError)
+                and not is_agentic_followup_failure(e)
+            )
+            if retries_apply:
+                litellm.num_retries = None  # set retries to None to prevent infinite loops
+            status_code: Final[object] = getattr(e, "status_code", None)
+            context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {})
+            context_window_fallback_applies: Final = (
+                call_type == CallTypes.acompletion.value
+                and isinstance(model, str)
+                and _context_window_fallback_applies(
+                    e, context_window_fallback_dict, model, is_router_call=is_acompletion_litellm_router_call
+                )
+            )
+            acompletion_retried: Final = (
+                retries_apply
+                and not context_window_fallback_applies
+                and (
+                    _get_retry_policy_num_retries(kwargs, e) is not None
+                    or not isinstance(status_code, int)
+                    or should_retry(status_code)
+                )
+                and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError))
+            )
+            if acompletion_retried:
+                kwargs["num_retries"] = num_retries
+                _restore_correlation_context_if_supported(logging_obj)
+                if logging_obj and not _is_litellm_internal_call:
+                    unbind_budget_reservation_from_callbacks(logging_obj.litellm_params)
+                from litellm.chat_completions.dispatch import acompletion as public_acompletion
+
+                try:
+                    result = await litellm.acompletion_with_retries(
+                        *args,
+                        original_function=public_acompletion,
+                        **{
+                            **request_kwargs,
+                            "litellm_trace_id": (
+                                request_kwargs.get("litellm_trace_id")
+                                if logging_obj is None
+                                else logging_obj.litellm_trace_id
+                            ),
+                            "num_retries": num_retries,
+                            "_initial_exception": e,
+                            _RETRY_ATTEMPT_KWARG: _RETRY_ATTEMPT_MARKER,
+                        },
+                    )
+                except openai.APIError as retry_err:
+                    _add_retry_context(retry_err, kwargs)
+                    if logging_obj and not _is_litellm_internal_call:
+                        await _end_async_wrapper_failure(
+                            logging_obj,
+                            retry_err,
+                            traceback.format_exc(),
+                            start_time,
+                            datetime.datetime.now(),  # noqa: DTZ005  # matches the naive datetimes this whole function already times start_time/end_time with
+                            log_failure=True,
+                        )
+                    raise
+                except Exception as retry_unavailable:
+                    verbose_logger.debug("acompletion retries did not run: %s", retry_unavailable)
+                else:
+                    return result
+
             traceback_exception: Final = traceback.format_exc()
             # Reuse the timestamp taken right when the deployment call itself failed, before
             # the failure hook ran, so a slow callback doesn't inflate the reported duration.
             end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now()  # noqa: DTZ005  # matches the naive datetimes this whole function already times start_time/end_time with
             if logging_obj and not _is_litellm_internal_call:
-                unbind_budget_reservation_from_callbacks(logging_obj.litellm_params)
-                try:
-                    logging_obj.failure_handler(
-                        e, traceback_exception, start_time, end_time, for_async_request=True
-                    )  # DO NOT MAKE THREADED - router retry fallback relies on this!
-                except Exception as e:
-                    raise e
-                try:
-                    await logging_obj.async_failure_handler(e, traceback_exception, start_time, end_time)
-                except Exception as e:
-                    raise e
+                await _end_async_wrapper_failure(
+                    logging_obj,
+                    e,
+                    traceback_exception,
+                    start_time,
+                    end_time,
+                    log_failure=not is_wrapper_retry_attempt,
+                )
+            if isinstance(wrapper_num_retries, Exception):
+                raise wrapper_num_retries
 
-            retry_call_type: Final = _get_call_type_from_public_name(original_function.__name__)
-            num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e)
-            if retry_call_type == CallTypes.acompletion.value:
-                context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {})
-
-                is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True)
-
-                if (
-                    num_retries
-                    and not is_acompletion_litellm_router_call
-                    and not isinstance(e, ImportError)
-                    and not is_agentic_followup_failure(e)
-                ):  # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying
-                    try:
-                        litellm.num_retries = None  # set retries to None to prevent infinite loops
-                        kwargs["num_retries"] = num_retries
-                        kwargs["original_function"] = original_function
-                        if isinstance(e, openai.RateLimitError):  # rate limiting specific error
-                            kwargs["retry_strategy"] = "exponential_backoff_retry"
-                        elif isinstance(e, openai.APIError):  # generic api error
-                            kwargs["retry_strategy"] = "constant_retry"
-                        result = await litellm.acompletion_with_retries(*args, **kwargs)
-                    except Exception:
-                        pass
-                    else:
-                        return result
-                elif (
-                    isinstance(e, litellm.exceptions.ContextWindowExceededError)
-                    and context_window_fallback_dict
-                    and model in context_window_fallback_dict
-                    and not is_acompletion_litellm_router_call
-                ):
-                    if len(args) > 0:
-                        args[0] = context_window_fallback_dict[model]
-                    else:
-                        kwargs["model"] = context_window_fallback_dict[model]
-                    result = await original_function(*args, **kwargs)
-                    return result
-            elif retry_call_type == CallTypes.aresponses.value:
+            if context_window_fallback_applies:
+                fallback_model: Final = context_window_fallback_dict[model]
+                fallback_kwargs: Final = _with_retry_policy_spent(kwargs)
+                if args:
+                    return await original_function(fallback_model, *args[1:], **fallback_kwargs)
+                return await original_function(*args, **{**fallback_kwargs, "model": fallback_model})
+            if call_type == CallTypes.aresponses.value:
                 is_aresponses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True)
 
                 if (
@@ -2686,19 +2784,14 @@ def client(original_function):
                             kwargs["retry_strategy"] = "exponential_backoff_retry"
                         elif isinstance(e, openai.APIError):  # generic api error
                             kwargs["retry_strategy"] = "constant_retry"
-                        result = await litellm.aresponses_with_retries(*args, **kwargs)
+                        result = await litellm.aresponses_with_retries(*args, **_with_retry_policy_spent(kwargs))
                     except Exception:
                         pass
                     else:
                         return result
 
-            deployment_num_retries: Final = kwargs.get("num_retries")
-            if deployment_num_retries is not None:
-                setattr(e, "num_retries", deployment_num_retries)
-
-            timeout: Final = _get_wrapper_timeout(kwargs=kwargs)
-            setattr(e, "timeout", timeout)
-            raise e
+            _add_retry_context(e, kwargs)
+            raise
 
         finally:
             # Restore trace_id/session_id contextvars to their pre-call value once