From c6794ac08b0e14d20760b8f6dd00eae79ed3ebb6 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 9 Nov 2026 16:06:52 -0800 Subject: [PATCH] fix(messages): let proxy /v1/messages requests run on the rust path (#50335) The proxy attaches its shared aiohttp session to every request and a default stream_timeout to every stream, and the messages detectors treated both as Python-only features, so no proxy request reached Rust. A shared session is only a connection-reuse hint and no native client reads it, so the common caller_client detector now fires only for an explicit client. The native messages call now takes the one timeout Python's messages handler resolves from timeout, request_timeout, stream_timeout and an httpx.Timeout, so no timeout shape keeps a messages call on Python --- litellm/messages/dispatch.py | 20 +++++-- litellm/rust_bridge/catalog.py | 4 +- litellm/rust_bridge/detectors.py | 4 +- .../messages/test_dispatch.py | 58 +++++++++++++++++-- .../responses/test_dispatch.py | 1 - tests/unit/rust_bridge/test_detectors.py | 2 +- 6 files changed, 75 insertions(+), 14 deletions(-) diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 4c94f26c6c1..40d21c5bad9 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -6,6 +6,7 @@ from typing import Final, TypeAlias, cast # noqa: TID251 # native binding sele from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.anthropic.messages.native import handler as main +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import detectors from litellm.rust_bridge.catalog import Feature, Route, RouteContext, Target from litellm.rust_bridge.dispatch import PublicDispatch @@ -25,6 +26,7 @@ from litellm.rust_bridge.public_call import ( signature, ) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from litellm.types.router import GenericLiteLLMParams __all__ = ("anthropic_messages", "anthropic_messages_handler") @@ -55,10 +57,20 @@ _AMESSAGES: Final = signature(_PYTHON_AMESSAGES) def _native_kwargs(fields: Mapping[str, object], kwargs: Mapping[str, object]) -> Mapping[str, object]: - stream_timeout: Final = kwargs.get("stream_timeout") - if fields.get("stream") is not True or stream_timeout is None: + """The native call takes one timeout, resolved from timeout, request_timeout and stream_timeout as Python's.""" + _, provider = _resolved_model(str(fields["model"]), optional_str(fields.get("custom_llm_provider"))) + timeout: Final = BaseLLMHTTPHandler._resolve_anthropic_messages_timeout( # pyright: ignore[reportPrivateUsage] # the handler's own timeout resolution + litellm_params=GenericLiteLLMParams( + timeout=kwargs.get("timeout"), + request_timeout=kwargs.get("request_timeout"), + stream_timeout=kwargs.get("stream_timeout"), + ), + stream=fields.get("stream") is True, + custom_llm_provider=provider or "", + ) + if timeout is None: return kwargs - return {**kwargs, "timeout": stream_timeout} + return {**kwargs, "timeout": timeout} def _public_request( @@ -122,7 +134,7 @@ def messages_target(model: str, custom_llm_provider: str | None, model_info: obj _FEATURES: Final[detectors.Detectors] = MappingProxyType( { - **detectors.COMMON, + **{feature: detect for feature, detect in detectors.COMMON.items() if feature is not Feature.TIMEOUT_VARIANTS}, Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, Feature.RESPONSE_CACHE: detectors.response_cache, diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 78040ba7225..c493617003d 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -39,7 +39,7 @@ class Feature(str, Enum): PROXY_MCP_TOOLS = "proxy_mcp_tools" """litellm_proxy MCP tools that Python expands and may auto-execute.""" CALLER_CLIENT = "caller_client" - """A caller-supplied HTTP client or shared session that Python sends through.""" + """A caller-supplied HTTP client that Python sends through.""" TIMEOUT_VARIANTS = "timeout_variants" """request_timeout, stream_timeout or httpx.Timeout phases beyond a plain timeout.""" DROP_PARAMS = "drop_params" @@ -172,7 +172,7 @@ CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { Feature.MODEL_ALIAS_MAP, } -MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { +MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VARIANTS}) | { Feature.PROVIDER_AFFINITY, Feature.FOUNDRY_BETA_ROW, Feature.SAFEGUARDS_CLASSIFIER, diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py index f2591bdbcd1..c38c3ea7135 100644 --- a/litellm/rust_bridge/detectors.py +++ b/litellm/rust_bridge/detectors.py @@ -44,8 +44,8 @@ def latest_alias(request: NativeCall) -> bool: def caller_client(request: NativeCall) -> bool: - """Python sends through the caller's HTTP client or shared session.""" - return present(request, "client") or present(request, "shared_session") + """Python sends through the caller's HTTP client; a shared session only reuses connections, which native pools.""" + return present(request, "client") def extra_body(request: NativeCall) -> bool: diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 6d65f289524..b580eca2e42 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -2,10 +2,12 @@ import inspect from collections.abc import Awaitable, Callable, Mapping from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect +import httpx import pytest import litellm from litellm.caching.caching import Cache +from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger from litellm.llms.anthropic.messages.native import handler as python_messages @@ -172,8 +174,25 @@ def test_native_receives_normalized_request_and_original_call_shape() -> None: assert call_kwargs["litellm_metadata"] is metadata -@pytest.mark.parametrize("stream, expected_timeout", ((True, 0.5), (False, 30.0))) -def test_native_streams_are_bounded_by_stream_timeout(stream: bool, expected_timeout: float) -> None: +@pytest.mark.parametrize( + ("kwargs", "expected_timeout"), + ( + pytest.param({"stream": True, "timeout": 30.0, "stream_timeout": 0.5}, 0.5, id="stream-timeout-bounds-stream"), + pytest.param( + {"stream": False, "timeout": 30.0, "stream_timeout": 0.5}, 30.0, id="stream-timeout-ignored-unary" + ), + pytest.param({"request_timeout": 5}, 5.0, id="request-timeout"), + pytest.param({"timeout": 30.0, "request_timeout": 5}, 30.0, id="timeout-before-request-timeout"), + pytest.param( + {"stream": True, "request_timeout": 5, "stream_timeout": 0.5}, 0.5, id="stream-timeout-over-request" + ), + pytest.param({"timeout": httpx.Timeout(60.0, connect=2.0)}, 60.0, id="httpx-timeout-read"), + pytest.param( + {"timeout": httpx.Timeout(None, connect=2.0)}, COMPLETION_HTTP_FALLBACK_SECONDS, id="httpx-timeout-no-read" + ), + ), +) +def test_native_timeout_resolves_like_python(kwargs: Mapping[str, object], expected_timeout: float) -> None: captured: Final[list[NativeCall]] = [] def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback @@ -185,7 +204,7 @@ def test_native_streams_are_bounded_by_stream_timeout(stream: bool, expected_tim _ = _DISPATCH.run( (16, MESSAGES, "anthropic/claude-sonnet-4-5"), - {"stream": stream, "timeout": 30.0, "stream_timeout": 0.5}, + kwargs, python=python, binding=messages_binding(native), native=lambda hook, request, call_args, call_kwargs: hook(request), @@ -456,7 +475,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": {}, ), (Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}), - (Feature.TIMEOUT_VARIANTS, "anthropic/claude-sonnet-4-5", {"request_timeout": 5}, {}), (Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}), (Feature.RESPONSE_CACHE, "anthropic/claude-sonnet-4-5", {}, {"cache": Cache()}), (Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), @@ -528,6 +546,38 @@ def test_each_detected_feature_reaches_its_rule( assert served(kwargs) is python_result +@pytest.mark.parametrize( + "kwargs", + ( + pytest.param({"shared_session": object()}, id="proxy-shared-session"), + pytest.param({"stream": True, "timeout": 600.0, "stream_timeout": 600.0}, id="proxy-stream-timeout"), + pytest.param({"request_timeout": 5}, id="request-timeout"), + pytest.param({"timeout": httpx.Timeout(60.0, connect=2.0)}, id="httpx-timeout"), + ), +) +def test_proxy_request_shape_stays_native(kwargs: Mapping[str, object]) -> None: + rules: Final[Rules] = ( + RouteRule( + Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.CALLER_CLIENT, Feature.TIMEOUT_VARIANTS}) + ), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ) + native_result: Final = response() + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("A proxy-shaped request must not fall back to Python") + + served: Final = _DISPATCH.run( + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + kwargs, + python=python, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, call_args, native_kwargs: hook(request), + rules=rules, + ) + assert served is native_result + + @pytest.mark.asyncio @pytest.mark.parametrize( "kwargs", diff --git a/tests/test_litellm_rust/responses/test_dispatch.py b/tests/test_litellm_rust/responses/test_dispatch.py index 4498dbdf39d..19255379dcf 100644 --- a/tests/test_litellm_rust/responses/test_dispatch.py +++ b/tests/test_litellm_rust/responses/test_dispatch.py @@ -260,7 +260,6 @@ ENCODED_ID: Final = ResponsesAPIRequestUtils._build_encrypted_item_id("deploymen (Feature.MOCK_RESPONSE, {"mock_response": "hi"}, {}), (Feature.PROXY_MCP_TOOLS, {"tools": [{"type": "mcp", "server_label": "g", "server_url": "litellm_proxy"}]}, {}), (Feature.CALLER_CLIENT, {"client": object()}, {}), - pytest.param(Feature.CALLER_CLIENT, {"shared_session": object()}, {}, id="shared-session"), pytest.param(Feature.EXTRA_BODY, {"extra_body": {"stream": True}}, {}, id="extra-body-stream"), pytest.param( Feature.PROVIDER_AFFINITY, diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index a6e6b2c0dca..fccc01dd8e6 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -111,7 +111,7 @@ def _wrap_nested_level(value: object, depth: int) -> object: (detectors.prompt_management, {"model": "litellm_agent/gpt-5"}, True), (detectors.prompt_management, {"model": "gpt-5"}, False), (detectors.caller_client, {"client": object()}, True), - (detectors.caller_client, {"shared_session": object()}, True), + (detectors.caller_client, {"shared_session": object()}, False), (detectors.caller_client, {"client": None, "shared_session": None}, False), (detectors.caller_client, {}, False), (detectors.extra_body, {"extra_body": {"stream": True}}, True),