From b42e3a35f8ce03d7310fbb195ea5b8e9293d18ae Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 20:55:21 -0700 Subject: [PATCH] fix(python-bridge): prevent duplicate messages fallback requests --- .../core/src/chat_completions/handler.rs | 31 +--- .../crates/core/src/chat_completions/tests.rs | 26 --- litellm-rust/crates/core/src/error.rs | 39 +++++ litellm-rust/crates/core/src/http_utils.rs | 8 + .../crates/core/src/messages/handler.rs | 165 +++++++++++++++++- .../crates/python-bridge/src/errors.rs | 59 ++++++- .../src/routes/chat_completions.rs | 6 +- .../python-bridge/src/routes/messages.rs | 6 +- litellm/llms/custom_httpx/llm_http_handler.py | 24 ++- .../test_rust_bridge_messages.py | 66 ++++++- 10 files changed, 356 insertions(+), 74 deletions(-) diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index afc4529fd26..713938fad51 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,7 +1,7 @@ use serde_json::Value; -use crate::error::{CoreError, CoreResult}; -use crate::http_utils::truncate_error_body; +use crate::error::{CoreError, CoreResult, as_response_error}; +use crate::http_utils::{classify_send_error, truncate_error_body}; use super::client::http_client; use super::transformation::ChatCompletionsAuth; @@ -27,16 +27,7 @@ pub(super) async fn execute_chat_completions_provider_call( request_builder = request_builder.timeout(duration); } - let response = request_builder.send().await.map_err(|err| { - // Failing to establish the connection means the request never went out, - // so the host can still serve it. Everything else here, a timeout - // above all, may have reached the provider and been answered. - if err.is_connect() || err.is_builder() { - CoreError::Connect(err.to_string()) - } else { - CoreError::Network(err.to_string()) - } - })?; + let response = request_builder.send().await.map_err(classify_send_error)?; let status = response.status(); let text = response @@ -60,22 +51,6 @@ pub(super) async fn execute_chat_completions_provider_call( .map_err(as_response_error) } -/// Re-tag an error raised while normalizing a response the provider already -/// returned. -/// -/// A config reports the same variants on either side of the call: a missing -/// field or an unsupported block can mean "this request cannot be translated" -/// during prepare and "this response cannot be normalized" here. Only the -/// second kind has already been billed, and a host that keeps a reference -/// implementation must not retry those, so collapse them to one variant that -/// can only mean the provider was already called. -pub(super) fn as_response_error(err: CoreError) -> CoreError { - match err { - already @ (CoreError::InvalidResponse(_) | CoreError::Http { .. }) => already, - other => CoreError::InvalidResponse(other.to_string()), - } -} - #[cfg(feature = "bedrock-auth")] pub(super) async fn signed_headers( request: &ProviderChatCompletionsRequest, diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index e2383723cb0..ad5e2cc6251 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -791,30 +791,4 @@ mod round_trip { "expected a pre-send connect failure, got {err:?}" ); } - - #[test] - fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { - use crate::chat_completions::handler::as_response_error; - - for original in [ - CoreError::MissingField("usage"), - CoreError::Unsupported("non-text response content block"), - CoreError::InvalidRequest("whatever".to_string()), - CoreError::Auth("whatever".to_string()), - ] { - let label = format!("{original:?}"); - assert!( - matches!(as_response_error(original), CoreError::InvalidResponse(_)), - "{label} must not stay retryable once the provider has answered" - ); - } - // An upstream status is already unambiguous, so it survives intact. - assert!(matches!( - as_response_error(CoreError::Http { - status: 500, - body: "boom".to_string() - }), - CoreError::Http { status: 500, .. } - )); - } } diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 739532f8cb5..15a5975c6a2 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -38,6 +38,14 @@ pub enum CoreError { Unsupported(&'static str), } +/// Re-tag an error raised after the provider has already returned a response. +pub(crate) fn as_response_error(err: CoreError) -> CoreError { + match err { + already @ (CoreError::InvalidResponse(_) | CoreError::Http { .. }) => already, + other => CoreError::InvalidResponse(other.to_string()), + } +} + pub fn json_type_name(value: &serde_json::Value) -> &'static str { match value { serde_json::Value::Null => "null", @@ -48,3 +56,34 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str { serde_json::Value::Object(_) => "object", } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn response_errors_collapse_to_one_non_retryable_variant() { + for original in [ + CoreError::MissingField("usage"), + CoreError::Unsupported("non-text response content block"), + CoreError::InvalidRequest("whatever".to_string()), + CoreError::Auth("whatever".to_string()), + ] { + assert!(matches!( + as_response_error(original), + CoreError::InvalidResponse(_) + )); + } + } + + #[test] + fn response_errors_preserve_an_upstream_status() { + assert!(matches!( + as_response_error(CoreError::Http { + status: 500, + body: "boom".to_string() + }), + CoreError::Http { status: 500, .. } + )); + } +} diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index c541f50275b..d79f178a124 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -5,6 +5,14 @@ use serde_json::{Map, Value}; use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; use crate::error::{CoreError, CoreResult, json_type_name}; +pub(crate) fn classify_send_error(error: reqwest::Error) -> CoreError { + if error.is_connect() || error.is_builder() { + CoreError::Connect(error.to_string()) + } else { + CoreError::Network(error.to_string()) + } +} + /// Bound an upstream error body before it crosses a host boundary, so provider /// bodies stay data-minimized. pub fn truncate_error_body(body: &str) -> String { diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 1c895f66eba..f44240eabe5 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,5 +1,6 @@ use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; -use crate::error::{CoreError, CoreResult}; +use crate::error::{CoreError, CoreResult, as_response_error}; +use crate::http_utils::classify_send_error; use super::client::http_client; use super::common_utils::truncate_error_body; @@ -16,10 +17,7 @@ pub(super) async fn execute_messages_provider_call( request_builder = request_builder.timeout(duration); } - let response = request_builder - .send() - .await - .map_err(|err| CoreError::Network(err.to_string()))?; + let response = request_builder.send().await.map_err(classify_send_error)?; let status = response.status(); let text = response @@ -37,7 +35,10 @@ pub(super) async fn execute_messages_provider_call( let response = serde_json::from_str(&text).map_err(|err| { CoreError::InvalidResponse(format!("invalid messages response JSON: {err}")) })?; - request.config.transform_response(&request.model, response) + request + .config + .transform_response(&request.model, response) + .map_err(as_response_error) } pub(super) async fn execute_messages_provider_stream( @@ -74,3 +75,155 @@ pub(super) async fn execute_messages_provider_stream( } Ok(response) } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use serde_json::json; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + + use super::*; + use crate::messages::transformation::AnthropicMessagesProviderConfig; + + struct RejectingResponseConfig; + + impl AnthropicMessagesProviderConfig for RejectingResponseConfig { + fn complete_url( + &self, + _api_base: Option<&str>, + _model: &str, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + unreachable!() + } + + fn resolve_api_key( + &self, + _api_key: Option<&str>, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + unreachable!() + } + + fn transform_response( + &self, + _model: &str, + _response: AnthropicMessagesResponse, + ) -> CoreResult { + Err(CoreError::MissingField("normalized_content")) + } + } + + static REJECTING_RESPONSE_CONFIG: RejectingResponseConfig = RejectingResponseConfig; + + fn request(url: String, timeout: Duration) -> ProviderMessagesRequest { + ProviderMessagesRequest { + provider: "anthropic".to_string(), + model: "claude-test".to_string(), + config: &REJECTING_RESPONSE_CONFIG, + url, + body: json!({}), + upstream_headers: Vec::new(), + timeout: Some(timeout), + } + } + + async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + loop { + let read = socket.read(&mut buffer).await.expect("reads request"); + if read == 0 { + break; + } + request.extend_from_slice(&buffer[..read]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + String::from_utf8(request).expect("request is utf8") + } + + #[tokio::test] + async fn post_response_transform_errors_are_non_retryable() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _ = read_http_request(&mut socket).await; + let body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-test"}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + }); + + let error = execute_messages_provider_call(request( + format!("http://{addr}/v1/messages"), + Duration::from_secs(5), + )) + .await + .expect_err("response transform should fail"); + + server.await.expect("server task completes"); + assert!( + matches!(error, CoreError::InvalidResponse(message) if message.contains("normalized_content")) + ); + } + + #[tokio::test] + async fn refused_connections_are_safe_to_fallback() { + let port = { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + listener.local_addr().expect("has an address").port() + }; + let error = execute_messages_provider_call(request( + format!("http://127.0.0.1:{port}"), + Duration::from_secs(1), + )) + .await + .expect_err("nothing is listening"); + + assert!(matches!(error, CoreError::Connect(_))); + } + + #[tokio::test] + async fn established_request_timeouts_are_network_errors() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("has an address"); + let (request_received_tx, request_received_rx) = tokio::sync::oneshot::channel(); + let (release_server_tx, release_server_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let received = read_http_request(&mut socket).await; + request_received_tx.send(received).expect("reports request"); + release_server_rx.await.expect("server is released"); + }); + + let error = tokio::time::timeout( + Duration::from_secs(2), + execute_messages_provider_call(request( + format!("http://{addr}"), + Duration::from_millis(100), + )), + ) + .await + .expect("client call completes") + .expect_err("established request times out"); + + let received = tokio::time::timeout(Duration::from_secs(2), request_received_rx) + .await + .expect("server observes request") + .expect("server reports request"); + assert!(received.starts_with("POST / "), "{received}"); + release_server_tx.send(()).expect("releases server"); + server.await.expect("server task completes"); + assert!(matches!(error, CoreError::Network(_))); + } +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index ec68cf7bfa3..b1601930d2e 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -33,7 +33,7 @@ pub(crate) fn core_error_to_pyerr(err: CoreError) -> PyErr { /// Everything raised before the request goes out is safe for the host to retry /// on its own path; anything after it is not, because the provider has already /// done the work and billed for it. -pub(crate) fn chat_completions_error_to_pyerr(err: CoreError) -> PyErr { +pub(crate) fn fallback_route_error_to_pyerr(err: CoreError) -> PyErr { match err { CoreError::Unsupported(_) | CoreError::Auth(_) @@ -59,3 +59,60 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("RustBridgeDeclined", py.get_type::())?; module.add("RustUpstreamError", py.get_type::()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fallback_routes_distinguish_declines_from_upstream_failures() { + Python::initialize(); + Python::attach(|py| { + let declines = [ + CoreError::Unsupported("unsupported"), + CoreError::Auth("missing key".to_string()), + CoreError::InvalidProvider("unsupported".to_string()), + CoreError::InvalidRequest("invalid".to_string()), + CoreError::InvalidType { + expected: "string", + actual: "number", + }, + CoreError::MissingField("model"), + CoreError::Routing("no route".to_string()), + CoreError::Connect("connection refused".to_string()), + ]; + for error in declines { + let mapped = fallback_route_error_to_pyerr(error); + assert!(mapped.is_instance_of::(py)); + } + + let upstream_failures = [ + ( + CoreError::Http { + status: 429, + body: "rate limited".to_string(), + }, + (429, "429: rate limited"), + ), + ( + CoreError::Network("request timed out".to_string()), + (0, "request timed out"), + ), + ( + CoreError::InvalidResponse("bad JSON".to_string()), + (0, "bad JSON"), + ), + ]; + for (error, expected) in upstream_failures { + let mapped = fallback_route_error_to_pyerr(error); + assert!(mapped.is_instance_of::(py)); + let args: (u16, String) = mapped + .value(py) + .getattr("args") + .and_then(|args| args.extract()) + .expect("upstream error should carry status and message"); + assert_eq!(args, (expected.0, expected.1.to_string())); + } + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index b353765e7cc..e1f5445f4bc 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -9,7 +9,7 @@ use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use serde_json::{Map, Value}; -use crate::errors::chat_completions_error_to_pyerr; +use crate::errors::fallback_route_error_to_pyerr; use crate::marshal::{optional_object_to_map, optional_timeout}; fn chat_completions_response_to_py( @@ -112,7 +112,7 @@ fn chat_completions( match result { Ok(response) => chat_completions_response_to_py(py, response), - Err(err) => Err(chat_completions_error_to_pyerr(err)), + Err(err) => Err(fallback_route_error_to_pyerr(err)), } } @@ -150,7 +150,7 @@ fn achat_completions( timeout, }) .await - .map_err(chat_completions_error_to_pyerr)?; + .map_err(fallback_route_error_to_pyerr)?; Python::attach(|py| chat_completions_response_to_py(py, response)) }) diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index ac361bb0e59..c337054e54d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -7,7 +7,7 @@ use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use serde_json::{Map, Value}; -use crate::errors::core_error_to_pyerr; +use crate::errors::fallback_route_error_to_pyerr; use crate::marshal::{optional_object_to_map, optional_timeout}; fn messages_response_to_py( @@ -66,7 +66,7 @@ fn messages( match result { Ok(response) => messages_response_to_py(py, response), - Err(err) => Err(core_error_to_pyerr(err)), + Err(err) => Err(fallback_route_error_to_pyerr(err)), } } @@ -97,7 +97,7 @@ fn amessages( timeout, }) .await - .map_err(core_error_to_pyerr)?; + .map_err(fallback_route_error_to_pyerr)?; Python::attach(|py| messages_response_to_py(py, response)) }) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 834f7d564a2..78f88a0b305 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -20,6 +20,7 @@ import litellm.types.utils from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES +from litellm.exceptions import APIError from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, @@ -2400,10 +2401,27 @@ class BaseLLMHTTPHandler: extra_headers=headers, timeout=timeout, ) - except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path + except Exception as rust_error: # noqa: BLE001 + from litellm.rust_bridge import get_native_bridge + + native_bridge: Final = get_native_bridge() + declined: Final = getattr(native_bridge, "RustBridgeDeclined", None) + upstream_failed: Final = getattr(native_bridge, "RustUpstreamError", None) + if isinstance(upstream_failed, type) and isinstance(rust_error, upstream_failed): + args: Final = rust_error.args + status: Final = args[0] if args else 0 + message: Final = args[1] if len(args) > 1 else "" + raise APIError( + status_code=int(status) or 500, + message=f"litellm rust messages: {message}", + llm_provider=custom_llm_provider, + model=model, + ) from rust_error + if not isinstance(declined, type) or not isinstance(rust_error, declined): + raise verbose_logger.debug( - "Rust Anthropic messages bridge raised %s; falling back to Python path", - type(rust_error).__name__, + "Rust Anthropic messages bridge declined before calling the provider (%s); falling back to Python path", + rust_error, ) return None if rust_response is None: diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index fbd7e36e298..ae1f6331b49 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -1,12 +1,14 @@ """Tests for the optional Rust-backed Anthropic Messages path.""" import importlib +from types import ModuleType from typing import cast import httpx import pytest import litellm +from litellm.exceptions import APIError from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -99,12 +101,28 @@ class ExplodingAsyncMessages: class RaisingAsyncMessages: - def __init__(self) -> None: + def __init__(self, error: Exception) -> None: self.calls = 0 + self.error = error async def __call__(self, **kwargs: object) -> dict[str, object]: self.calls += 1 - raise RuntimeError("upstream request failed with status 400: bad request") + raise self.error + + +class FakeBridgeDeclined(Exception): + pass + + +class FakeUpstreamError(Exception): + pass + + +def _install_fake_bridge_exceptions(monkeypatch) -> None: + native_bridge = ModuleType("_native") + native_bridge.RustBridgeDeclined = FakeBridgeDeclined + native_bridge.RustUpstreamError = FakeUpstreamError + monkeypatch.setattr(rust_bridge_loader, "_cached_bridge", native_bridge) @pytest.fixture(autouse=True) @@ -251,8 +269,9 @@ async def test_gate_invokes_rust_and_marks_response_header(): @pytest.mark.asyncio -async def test_gate_falls_back_to_python_when_bridge_raises(): - bridge = RaisingAsyncMessages() +async def test_gate_falls_back_only_when_bridge_declines(monkeypatch): + _install_fake_bridge_exceptions(monkeypatch) + bridge = RaisingAsyncMessages(FakeBridgeDeclined("unsupported request")) litellm.use_litellm_rust(True, amessages=bridge) response = await _gate() @@ -261,6 +280,45 @@ async def test_gate_falls_back_to_python_when_bridge_raises(): assert bridge.calls == 1 +@pytest.mark.asyncio +async def test_gate_surfaces_an_upstream_failure_without_fallback(monkeypatch): + _install_fake_bridge_exceptions(monkeypatch) + bridge = RaisingAsyncMessages(FakeUpstreamError(429, "429: rate limited")) + litellm.use_litellm_rust(True, amessages=bridge) + + with pytest.raises(APIError) as exc_info: + await _gate() + + assert exc_info.value.status_code == 429 + assert "429: rate limited" in str(exc_info.value) + assert bridge.calls == 1 + + +@pytest.mark.asyncio +async def test_gate_maps_statusless_upstream_failure_to_500_without_fallback(monkeypatch): + _install_fake_bridge_exceptions(monkeypatch) + bridge = RaisingAsyncMessages(FakeUpstreamError(0, "request timed out")) + litellm.use_litellm_rust(True, amessages=bridge) + + with pytest.raises(APIError) as exc_info: + await _gate() + + assert exc_info.value.status_code == 500 + assert "request timed out" in str(exc_info.value) + assert bridge.calls == 1 + + +@pytest.mark.asyncio +async def test_gate_reraises_an_unknown_bridge_failure(): + bridge = RaisingAsyncMessages(RuntimeError("unknown bridge failure")) + litellm.use_litellm_rust(True, amessages=bridge) + + with pytest.raises(RuntimeError, match="unknown bridge failure"): + await _gate() + + assert bridge.calls == 1 + + @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_absent(): bridge = ExplodingAsyncMessages()