Merge e3163aa98b into eb93762816

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"

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,
         )

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:

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