From e3163aa98b0ef3c024277053adcd2f2e806b70aa Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 18 Oct 2026 13:10:05 -0700 Subject: [PATCH] fix(rust-bridge): stamp num_retries and timeout on async native failures The async @client wrapper sets num_retries and timeout on every exception it raises, and the Router reads num_retries off the exception to decide how many times to retry. Async native calls raised without them. Both paths now call one stamp_async_failure helper, so a native failure carries the same retry count and timeout a Python one does Co-Authored-By: Claude Opus 5.5 --- litellm/rust_bridge/catalog.py | 6 +++++- litellm/rust_bridge/dispatch.py | 14 ++++++++++++-- litellm/utils.py | 16 +++++++++------- .../test_litellm_rust/test_public_entries.py | 19 +++++++++++++++++++ 4 files changed, 45 insertions(+), 10 deletions(-) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 7f6d560d1db..93d344aa3ea 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -73,7 +73,11 @@ class Feature(str, Enum): MESSAGES_INTERCEPTOR = "messages_interceptor" """A registered MessagesInterceptor that takes over the call.""" AGENTIC_LOOP_HOOK = "agentic_loop_hook" - """A logger overriding async_should_run_agentic_loop, which may run follow-up calls.""" + """A logger overriding async_should_run_agentic_loop, which may run follow-up calls. + + Every interception that turns a stream into a converted stream in its pre-call deployment hook + (websearch, code interpreter, headroom) also overrides this method, so this feature keeps those + calls on Python too.""" TLS_POLICY = "tls_policy" """An OpenSSL security level or ECDH curve that native TLS cannot apply.""" MANAGED_RESOURCE_IDS = "managed_resource_ids" diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py index 42c1a96c03d..aea1519701d 100644 --- a/litellm/rust_bridge/dispatch.py +++ b/litellm/rust_bridge/dispatch.py @@ -25,6 +25,16 @@ def call_hook( return hook(request, args, kwargs) +async def _stamped(start: Callable[[], Awaitable[ResultT]], kwargs: Mapping[str, object]) -> ResultT: + from litellm.utils import stamp_async_failure + + try: + return await start() + except Exception as error: + stamp_async_failure(error, kwargs) + raise + + @dataclass(frozen=True, slots=True) class PublicDispatch(Generic[RequestT]): route: Route @@ -100,7 +110,7 @@ class PublicDispatch(Generic[RequestT]): return await runtime.arun( self.context(native_request), binding=binding, - native=lambda hook: native(hook, native_request, args, kwargs), + native=lambda hook: _stamped(lambda: native(hook, native_request, args, kwargs), kwargs), python=python, rules=selected_rules, ) @@ -115,7 +125,7 @@ class PublicDispatch(Generic[RequestT]): return await runtime.arun( context, binding=binding, - native=lambda hook: native(hook, request, args, kwargs), + native=lambda hook: _stamped(lambda: native(hook, request, args, kwargs), kwargs), python=lambda: python(*args, **kwargs), rules=selected_rules, ) diff --git a/litellm/utils.py b/litellm/utils.py index ab97b69897f..4a1b189c261 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1482,7 +1482,14 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu return num_retries, kwargs -def _get_wrapper_timeout(kwargs: dict[str, object]) -> object: +def stamp_async_failure(error: Exception, kwargs: Mapping[str, object]) -> None: + num_retries: Final = kwargs.get("num_retries") + if num_retries is not None: + setattr(error, "num_retries", num_retries) + setattr(error, "timeout", _get_wrapper_timeout(kwargs=kwargs)) + + +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: @@ -2350,12 +2357,7 @@ def client(original_function): 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) + stamp_async_failure(e, kwargs) raise e finally: diff --git a/tests/test_litellm_rust/test_public_entries.py b/tests/test_litellm_rust/test_public_entries.py index 2f5c33e3ced..47558d4923a 100644 --- a/tests/test_litellm_rust/test_public_entries.py +++ b/tests/test_litellm_rust/test_public_entries.py @@ -285,3 +285,22 @@ async def test_sync_entry_with_flag_true_raises_setup_errors_at_await(entry: Ent assert inspect.iscoroutine(pending) with pytest.raises(RuntimeError, match="is unavailable"): await pending + + +@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id) +async def test_async_native_failures_carry_the_retry_count_and_timeout_the_router_reads(entry: Entry) -> None: + failure: Final = litellm.RateLimitError(message="limited", llm_provider="openai", model="gpt-4o") + + async def failing(_: NativeCall) -> object: + raise failure + + entry.async_binding.override(failing) + try: + with scenario_rules(entry, SCENARIOS["required"]), pytest.raises(litellm.RateLimitError) as caught: + await entry.asynchronous(*entry.args(), **entry.kwargs(num_retries=3, timeout=7.0)) + finally: + entry.async_binding.reset() + + assert caught.value is failure + assert getattr(caught.value, "num_retries", None) == 3 + assert getattr(caught.value, "timeout", None) == 7.0