From 1467b13fd2a777b0824fe3eb1140f117d733e5cf Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:24:20 -0800 Subject: [PATCH] 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 | 3 +- litellm/embeddings/dispatch.py | 3 +- litellm/messages/dispatch.py | 3 +- litellm/ocr/dispatch.py | 3 +- litellm/responses/dispatch.py | 3 +- litellm/rust_bridge/public_call.py | 36 +++++++++- .../test_litellm_rust/test_public_entries.py | 36 +++++++++- tests/unit/rust_bridge/test_public_call.py | 70 ++++++++++++++++++- 8 files changed, 147 insertions(+), 10 deletions(-) diff --git a/litellm/chat_completions/dispatch.py b/litellm/chat_completions/dispatch.py index 06ba8e9ad76..05572ee8837 100644 --- a/litellm/chat_completions/dispatch.py +++ b/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, ) diff --git a/litellm/embeddings/dispatch.py b/litellm/embeddings/dispatch.py index 772ac68737d..9547757ff8f 100644 --- a/litellm/embeddings/dispatch.py +++ b/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, ) diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 40d21c5bad9..faa604a25ce 100644 --- a/litellm/messages/dispatch.py +++ b/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, ) diff --git a/litellm/ocr/dispatch.py b/litellm/ocr/dispatch.py index 529c3df2bb9..f90a42bb1c9 100644 --- a/litellm/ocr/dispatch.py +++ b/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, ) diff --git a/litellm/responses/dispatch.py b/litellm/responses/dispatch.py index 0dda3f93660..acf6610d06e 100644 --- a/litellm/responses/dispatch.py +++ b/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, ) diff --git a/litellm/rust_bridge/public_call.py b/litellm/rust_bridge/public_call.py index 8625ca71bb2..1ab3d5e10c9 100644 --- a/litellm/rust_bridge/public_call.py +++ b/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 diff --git a/tests/test_litellm_rust/test_public_entries.py b/tests/test_litellm_rust/test_public_entries.py index 2f5c33e3ced..24a3419c5a0 100644 --- a/tests/test_litellm_rust/test_public_entries.py +++ b/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"]) diff --git a/tests/unit/rust_bridge/test_public_call.py b/tests/unit/rust_bridge/test_public_call.py index 5f50db6ea5a..27c3c119c1b 100644 --- a/tests/unit/rust_bridge/test_public_call.py +++ b/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"]