fix(rust): run native async calls through python's async wrapper steps (#50423)

Python's async client wrapper stamps the effective timeout and num_retries on every failure, which the proxy reads for x-litellm-timeout and the router reads for retries. The native async paths skipped that wrapper, so error responses from Rust-served async calls had no x-litellm-timeout header

Python's async chat completions, responses, embeddings and messages entries hand the call to run_in_executor before they build the provider body, so a during_call guardrail started next to the call masks the request first. The native paths built their request without yielding, so masked text never reached the provider

The shared native_acall_hook now yields once and stamps the failure, and every async dispatch passes it as its native hook. OCR passes native_acall_hook_on_loop, which only stamps the failure, since its Python entry went through the client wrapper but built its request on the loop

litellm/chat_completions/dispatch.py

@@ -16,6 +16,7 @@ from litellm.rust_bridge.public_call import (
     AsyncFlag,
     NativeCall,
     bind,
+    native_acall_hook,
     native_call,
     native_call_hook,
     optional_sequence,
@@ -131,7 +132,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult:  # kwargs-
         _FLAG.dropped(kwargs),
         python=python,
         binding=NATIVE_ACOMPLETION,
-        native=native_call_hook,
+        native=native_acall_hook,
     )
 
 

litellm/embeddings/dispatch.py

@@ -13,6 +13,7 @@ from litellm.rust_bridge.public_call import (
     AsyncFlag,
     NativeCall,
     bind,
+    native_acall_hook,
     native_call,
     native_call_hook,
     optional_str,
@@ -89,7 +90,7 @@ async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse:  # k
         _FLAG.dropped(kwargs),
         python=_PYTHON_AEMBEDDING,
         binding=NATIVE_AEMBEDDING,
-        native=native_call_hook,
+        native=native_acall_hook,
     )
 
 

litellm/messages/dispatch.py

@@ -19,6 +19,7 @@ from litellm.rust_bridge.public_call import (
     AsyncFlag,
     NativeCall,
     bind,
+    native_acall_hook,
     native_call,
     native_call_hook,
     optional_sequence,
@@ -207,7 +208,7 @@ async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult:
         _FLAG.dropped(kwargs),
         python=python,
         binding=NATIVE_AMESSAGES,
-        native=native_call_hook,
+        native=native_acall_hook,
     )
 
 

litellm/ocr/dispatch.py

@@ -12,6 +12,7 @@ from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
 from litellm.rust_bridge.public_call import (
     AsyncFlag,
     NativeCall,
+    native_acall_hook_on_loop,
     native_call,
     native_call_hook,
     optional_str,
@@ -97,5 +98,5 @@ async def aocr(*args: object, **kwargs: object) -> OCRResponse:  # kwargs-ok: pr
         _FLAG.dropped(kwargs),
         python=runtime.NO_PYTHON,
         binding=NATIVE_AOCR,
-        native=native_call_hook,
+        native=native_acall_hook_on_loop,
     )

litellm/responses/dispatch.py

@@ -14,6 +14,7 @@ from litellm.rust_bridge.public_call import (
     AsyncFlag,
     NativeCall,
     bind,
+    native_acall_hook,
     native_call,
     native_call_hook,
     optional_str,
@@ -150,7 +151,7 @@ async def aresponses(*args: object, **kwargs: object) -> ResponsesResult:  # kwa
         _FLAG.dropped(kwargs),
         python=python,
         binding=NATIVE_ARESPONSES,
-        native=native_call_hook,
+        native=native_acall_hook,
     )
 
 

litellm/rust_bridge/public_call.py

@@ -2,12 +2,17 @@
 
 from __future__ import annotations
 
+import asyncio
 import inspect
-from collections.abc import Callable, Collection, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence
 from dataclasses import dataclass
 from types import MappingProxyType
 from typing import Final, TypeVar, cast  # noqa: TID251  # narrows caller-owned containers without copying them
 
+from litellm.utils import (
+    _get_wrapper_timeout,  # pyright: ignore[reportPrivateUsage]  # the @client wrapper's own failure stamp
+)
+
 
 def signature(legacy: Callable[..., object]) -> inspect.Signature:
     return inspect.signature(legacy)
@@ -115,3 +120,32 @@ def native_call_hook(
     _kwargs: Mapping[str, object],
 ) -> NativeResultT:
     return hook(call)
+
+
+async def native_acall_hook(
+    hook: Callable[[NativeCall], Awaitable[NativeResultT]],
+    call: NativeCall,
+    args: tuple[object, ...],
+    kwargs: Mapping[str, object],
+) -> NativeResultT:
+    # Python's async entries leave the event loop (run_in_executor) before they build the provider body, so
+    # tasks scheduled next to the call, such as the proxy's during_call guardrails, edit the request first
+    await asyncio.sleep(0)
+    return await native_acall_hook_on_loop(hook, call, args, kwargs)
+
+
+async def native_acall_hook_on_loop(
+    hook: Callable[[NativeCall], Awaitable[NativeResultT]],
+    call: NativeCall,
+    _args: tuple[object, ...],
+    kwargs: Mapping[str, object],
+) -> NativeResultT:
+    """``native_acall_hook`` for an async entry whose Python path built its request on the loop, as OCR's did."""
+    try:
+        return await hook(call)
+    except Exception as error:
+        # utils.client stamps every async failure, which the proxy reads for x-litellm-timeout and the router for retries
+        if kwargs.get("num_retries") is not None:
+            setattr(error, "num_retries", kwargs.get("num_retries"))  # noqa: B010  # Exception declares no such attribute, as in utils.client
+        setattr(error, "timeout", _get_wrapper_timeout(kwargs=dict(kwargs)))  # noqa: B010  # Exception declares no such attribute, as in utils.client
+        raise

tests/test_litellm_rust/test_public_entries.py

@@ -1,3 +1,4 @@
+import asyncio
 import inspect
 from collections.abc import Awaitable, Callable, Iterator, Mapping
 from contextlib import contextmanager
@@ -29,7 +30,10 @@ Served: TypeAlias = Callable[[object], bool]
 
 @dataclass(frozen=True, slots=True)
 class Entry:
-    """``python`` is ``None`` for a route with no Python implementation, which only RUST_REQUIRED rules serve."""
+    """``python`` is ``None`` for a route with no Python implementation, which only RUST_REQUIRED rules serve.
+
+    ``executor_hop`` is whether the route's Python async entry left the loop through run_in_executor before it built the request.
+    """
 
     route: Route
     sync: Callable[..., object]
@@ -40,6 +44,7 @@ class Entry:
     arguments: Mapping[str, object]
     positional: int
     python: Served | None
+    executor_hop: bool
 
     def args(self) -> tuple[object, ...]:
         return tuple(self.arguments.values())[: self.positional]
@@ -79,6 +84,7 @@ ENTRIES: Final = (
         arguments={"model": "openai/gpt-4o", "messages": MESSAGES, "mock_response": "python"},
         positional=2,
         python=_chat,
+        executor_hop=True,
     ),
     Entry(
         route=Route.EMBEDDINGS,
@@ -90,6 +96,7 @@ ENTRIES: Final = (
         arguments={"model": "openai/text-embedding-3-small", "input": ["hi"], "mock_response": [0.5]},
         positional=2,
         python=_embedding,
+        executor_hop=True,
     ),
     Entry(
         route=Route.RESPONSES,
@@ -101,6 +108,7 @@ ENTRIES: Final = (
         arguments={"input": "hi", "model": "openai/gpt-4o", "mock_response": "python"},
         positional=2,
         python=_responses,
+        executor_hop=True,
     ),
     Entry(
         route=Route.MESSAGES,
@@ -117,6 +125,7 @@ ENTRIES: Final = (
         },
         positional=0,
         python=_messages,
+        executor_hop=True,
     ),
     Entry(
         route=Route.OCR,
@@ -128,6 +137,7 @@ ENTRIES: Final = (
         arguments={"model": "mistral/mistral-ocr-latest", "document": {"type": "file", "file": b"pdf"}},
         positional=2,
         python=None,
+        executor_hop=False,
     ),
 )
 
@@ -285,3 +295,27 @@ 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_entry_runs_native_through_the_client_wrapper_steps(entry: Entry) -> None:
+    scheduled: Final[list[str]] = []
+
+    async def failing(native_call: NativeCall) -> ModelResponse:
+        scheduled.append("native")
+        raise litellm.Timeout(message="timed out", model="model", llm_provider="openai")
+
+    async def started_before() -> None:
+        scheduled.append("task")
+
+    entry.async_binding.override(failing)
+    task: Final = asyncio.create_task(started_before())
+    try:
+        with scenario_rules(entry, SCENARIOS["required"]), pytest.raises(litellm.Timeout) as raised:
+            await entry.asynchronous(*entry.args(), **entry.kwargs(timeout=30.0, num_retries=0))
+    finally:
+        entry.async_binding.reset()
+        await task
+
+    assert (getattr(raised.value, "timeout", None), getattr(raised.value, "num_retries", None)) == (30.0, 0)
+    assert scheduled == (["task", "native"] if entry.executor_hop else ["native", "task"])

tests/unit/rust_bridge/test_public_call.py

@@ -1,9 +1,22 @@
-from collections.abc import Mapping, Sequence
-from typing import Final
+import asyncio
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from typing import Final, TypeAlias
 
 import pytest
 
-from litellm.rust_bridge.public_call import AsyncFlag, native_call, signature
+import litellm
+from litellm.rust_bridge.public_call import (
+    AsyncFlag,
+    NativeCall,
+    native_acall_hook,
+    native_acall_hook_on_loop,
+    native_call,
+    signature,
+)
+
+NativeAcallHook: TypeAlias = Callable[
+    [Callable[[NativeCall], Awaitable[None]], NativeCall, tuple[object, ...], Mapping[str, object]], Awaitable[None]
+]
 
 
 def _messages(
@@ -76,3 +89,54 @@ def test_flag_keywords_name_positionals_and_drop_only_the_flag() -> None:
 def test_flag_keywords_reject_a_parameter_given_twice() -> None:
     with pytest.raises(TypeError):
         _FLAG.keywords((128,), {"max_tokens": 64})
+
+
+_CALL: Final = native_call(_SIGNATURE, (128, [], "model"), {})
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("kwargs", "expected_timeout", "expected_num_retries"),
+    (
+        ({"timeout": 30.0, "stream_timeout": 0.5, "num_retries": 0}, 30.0, 0),
+        ({"stream": True, "timeout": 30.0, "stream_timeout": 0.5}, 0.5, None),
+        ({"stream": True, "timeout": 30.0}, 30.0, None),
+        ({}, None, None),
+    ),
+)
+@pytest.mark.parametrize("hook", (native_acall_hook, native_acall_hook_on_loop))
+async def test_async_hook_failures_carry_the_client_wrapper_stamp(
+    hook: NativeAcallHook,
+    kwargs: Mapping[str, object],
+    expected_timeout: float | None,
+    expected_num_retries: int | None,
+) -> None:
+    async def native(call: NativeCall) -> None:
+        raise litellm.Timeout(message="timed out", model="model", llm_provider="anthropic")
+
+    with pytest.raises(litellm.Timeout) as raised:
+        await hook(native, _CALL, (), kwargs)
+
+    assert getattr(raised.value, "timeout", "unset") == expected_timeout
+    assert getattr(raised.value, "num_retries", None) == expected_num_retries
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("hook", "yields"), ((native_acall_hook, True), (native_acall_hook_on_loop, False)))
+async def test_only_the_executor_hook_lets_a_task_started_before_the_call_run_first(
+    hook: NativeAcallHook, yields: bool
+) -> None:
+    messages: Final = [{"role": "user", "content": "during about secret"}]
+    sent: Final[list[object]] = []
+
+    async def mask() -> None:
+        messages[0]["content"] = "during about [REDACTED]"
+
+    async def native(call: NativeCall) -> None:
+        sent.append(messages[0]["content"])
+
+    guardrail: Final = asyncio.create_task(mask())
+    await hook(native, _CALL, (), {})
+    await guardrail
+
+    assert sent == ["during about [REDACTED]" if yields else "during about secret"]