diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index dbe2d3a325b..42ffd4cbcff 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -277,7 +277,6 @@ fn core_error_kind(error: &Error) -> &'static str { Error::Http { .. } => "HttpError", Error::InvalidResponse(_) => "InvalidResponse", Error::Network(_) => "NetworkError", - Error::Connect(_) => "ConnectError", Error::Routing(_) => "RoutingError", Error::Unsupported(_) => "UnsupportedRequest", } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 401e26d3b29..6343122b9e9 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -321,7 +321,6 @@ fn core_error_kind(error: &Error) -> &'static str { Error::Http { .. } => "HttpError", Error::InvalidResponse(_) => "InvalidResponse", Error::Network(_) => "NetworkError", - Error::Connect(_) => "ConnectError", Error::Routing(_) => "RoutingError", Error::Unsupported(_) => "UnsupportedRequest", } diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index e9f8c477f36..e4861c3954c 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -105,7 +105,6 @@ impl IntoResponse for MessagesRouteError { ), Error::Http { .. } | Error::Network(_) - | Error::Connect(_) | Error::InvalidResponse(_) | Error::InvalidType { .. } | Error::MissingField(_) => ( diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 7e2731442cc..f282e82c4fa 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::Error; -use crate::http_utils::truncate_error_body; +use crate::error::{Error, 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() { - Error::Connect(err.to_string()) - } else { - Error::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: Error) -> Error { - match err { - already @ (Error::InvalidResponse(_) | Error::Http { .. }) => already, - other => Error::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 2858d180e27..bdd413d7cff 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -769,11 +769,7 @@ mod round_trip { } #[tokio::test] - async fn a_connection_that_is_never_established_declines_instead_of_failing() { - // Nothing was sent, so nothing was billed and the host can still serve - // the request. Classing this with the post-send failures would turn a - // recoverable fallback into a user-facing error on exactly the - // deployments whose transport is configured only on the Python client. + async fn a_connection_that_is_never_established_is_still_terminal() { let port = { let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); listener.local_addr().expect("has an address").port() @@ -787,34 +783,8 @@ mod round_trip { .await .expect_err("nothing is listening"); assert!( - matches!(err, Error::Connect(_)), - "expected a pre-send connect failure, got {err:?}" + matches!(err, Error::Network(_)), + "expected a terminal network 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 [ - Error::MissingField("usage"), - Error::Unsupported("non-text response content block"), - Error::InvalidRequest("whatever".to_string()), - Error::Auth("whatever".to_string()), - ] { - let label = format!("{original:?}"); - assert!( - matches!(as_response_error(original), Error::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(Error::Http { - status: 500, - body: "boom".to_string() - }), - Error::Http { status: 500, .. } - )); - } } diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index db3fa2ec704..a5de4b722d2 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -21,13 +21,6 @@ pub enum Error { Http { status: u16, body: String }, #[error("upstream network error: {0}")] Network(String), - /// The provider was never reached: DNS, TCP, TLS or proxy setup failed - /// before any byte of the request went out. Nothing was billed, so a host - /// that keeps a reference implementation can serve the request itself. - /// A timeout is deliberately not this, since the provider may have received - /// and answered the request already. - #[error("could not reach the provider: {0}")] - Connect(String), #[error("routing error: {0}")] Routing(String), /// The request is outside the surface this route covers in Rust. Hosts that @@ -36,6 +29,14 @@ pub enum Error { Unsupported(&'static str), } +/// Re-tag an error raised after the provider has already returned a response. +pub(crate) fn as_response_error(err: Error) -> Error { + match err { + already @ (Error::InvalidResponse(_) | Error::Http { .. }) => already, + other => Error::InvalidResponse(other.to_string()), + } +} + pub fn json_type_name(value: &serde_json::Value) -> &'static str { match value { serde_json::Value::Null => "null", @@ -46,3 +47,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 [ + Error::MissingField("usage"), + Error::Unsupported("non-text response content block"), + Error::InvalidRequest("whatever".to_string()), + Error::Auth("whatever".to_string()), + ] { + assert!(matches!( + as_response_error(original), + Error::InvalidResponse(_) + )); + } + } + + #[test] + fn response_errors_preserve_an_upstream_status() { + assert!(matches!( + as_response_error(Error::Http { + status: 500, + body: "boom".to_string() + }), + Error::Http { status: 500, .. } + )); + } +} diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index 10661fadf96..759c7ba5652 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -5,6 +5,10 @@ use serde_json::{Map, Value}; use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; use crate::error::{Error, json_type_name}; +pub(crate) fn classify_send_error(error: reqwest::Error) -> Error { + Error::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 13a65d86131..441f133b182 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::Error; +use crate::error::{Error, 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| Error::Network(err.to_string()))?; + let response = request_builder.send().await.map_err(classify_send_error)?; let status = response.status(); let text = response @@ -36,7 +34,10 @@ pub(super) async fn execute_messages_provider_call( let response = serde_json::from_str(&text) .map_err(|err| Error::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( @@ -73,3 +74,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, + ) -> Result { + unreachable!() + } + + fn resolve_api_key( + &self, + _api_key: Option<&str>, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + unreachable!() + } + + fn transform_response( + &self, + _model: &str, + _response: AnthropicMessagesResponse, + ) -> Result { + Err(Error::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, Error::InvalidResponse(message) if message.contains("normalized_content")) + ); + } + + #[tokio::test] + async fn refused_connections_are_terminal_network_errors() { + 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, Error::Network(_))); + } + + #[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, Error::Network(_))); + } +} diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 3b253ac3766..3d10a9d824f 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -25,8 +25,9 @@ pub(super) fn prepare_messages_call( let model = provider_info.model.to_string(); let provider = provider_info.custom_llm_provider; - let config = messages_provider_config(provider) - .ok_or_else(|| Error::InvalidProvider(provider.to_string()))?; + let config = messages_provider_config(provider).ok_or(Error::Unsupported( + "messages provider is not registered in the Rust bridge", + ))?; let env_lookup = |key: &str| std::env::var(key).ok(); let mut headers = string_headers(request.extra_headers)?; diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index df9f7051011..27dc2281d64 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -437,5 +437,8 @@ async fn messages_rejects_unsupported_provider() { .await .expect_err("unsupported provider errors"); - assert!(matches!(err, Error::InvalidProvider(provider) if provider == "openai")); + assert!(matches!( + err, + Error::Unsupported("messages provider is not registered in the Rust bridge") + )); } diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 914e2e1e033..d3e47937666 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -29,28 +29,15 @@ pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { /// Map a core error for a route whose host keeps a Python implementation. /// -/// The distinction the host needs is whether the provider was already called. -/// 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: Error) -> PyErr { +/// Only an explicit capability decline permits the host to try Python. Every +/// other error may have happened after provider dispatch and must be terminal. +pub(crate) fn fallback_route_error_to_pyerr(err: Error) -> PyErr { match err { - Error::Unsupported(_) - | Error::Auth(_) - | Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::InvalidType { .. } - | Error::MissingField(_) - | Error::Routing(_) - // Nothing reached the provider, so serving it on Python cannot double - // bill and is the only way the caller gets an answer at all. - | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), + Error::Unsupported(_) => RustBridgeDeclined::new_err(err.to_string()), Error::Http { status, body } => { RustUpstreamError::new_err((status, format!("{status}: {body}"))) } - Error::Network(message) | Error::InvalidResponse(message) => { - RustUpstreamError::new_err((0u16, message)) - } + other => RustUpstreamError::new_err((0u16, other.to_string())), } } @@ -59,3 +46,53 @@ 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 = [Error::Unsupported("unsupported")]; + for error in declines { + let mapped = fallback_route_error_to_pyerr(error); + assert!(mapped.is_instance_of::(py)); + } + + let upstream_failures = [ + ( + Error::Http { + status: 429, + body: "rate limited".to_string(), + }, + (429, "429: rate limited"), + ), + ( + Error::Network("request timed out".to_string()), + (0, "upstream network error: request timed out"), + ), + ( + Error::InvalidResponse("bad JSON".to_string()), + (0, "invalid response: bad JSON"), + ), + (Error::Auth("missing key".to_string()), (0, "missing key")), + ( + Error::InvalidRequest("invalid".to_string()), + (0, "invalid request: invalid"), + ), + ]; + 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 68b7762cb10..a0257a4ade2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -8,7 +8,7 @@ use litellm_core::chat_completions::{ use pyo3::prelude::*; use serde_json::Value; -use crate::errors::chat_completions_error_to_pyerr; +use crate::errors::fallback_route_error_to_pyerr; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value}; fn prepare_chat_completions( @@ -86,6 +86,6 @@ bridge_route! { timeout_seconds: Option, }, prepare = prepare_chat_completions, - errors = chat_completions_error_to_pyerr, + errors = fallback_route_error_to_pyerr, extra = [chat_completions_decline], } diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 5ba1cc6f1ad..d2674a80a10 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -401,6 +401,35 @@ asyncio.run(exercise()) }); } + #[test] + fn messages_routes_only_decline_unsupported_requests() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "routes").expect("module should be created"); + crate::routes::register(&module).expect("routes should register"); + let body = PyDict::new(py); + let kwargs = PyDict::new(py); + kwargs + .set_item("custom_llm_provider", "openai") + .expect("kwargs should accept provider"); + + let error = module + .getattr("messages") + .and_then(|function| function.call(("model", &body), Some(&kwargs))) + .expect_err("unsupported provider should decline"); + assert!(error.is_instance_of::(py)); + + kwargs + .set_item("custom_llm_provider", "anthropic") + .expect("kwargs should accept provider"); + let error = module + .getattr("messages") + .and_then(|function| function.call(("model", &body), Some(&kwargs))) + .expect_err("missing credentials should fail"); + assert!(error.is_instance_of::(py)); + }); + } + #[test] fn route_registration_rejects_duplicate_python_names() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index 2bb64a7a763..da028486e2f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -5,7 +5,7 @@ use pyo3::prelude::*; use serde_json::Value; use std::future::Future; -use crate::errors::core_error_to_pyerr; +use crate::errors::fallback_route_error_to_pyerr; use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value}; fn prepare_messages( @@ -61,5 +61,5 @@ bridge_route! { timeout_seconds: Option, }, prepare = prepare_messages, - errors = core_error_to_pyerr, + errors = fallback_route_error_to_pyerr, } diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3cd6ee54069..9aaf92f62e4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -19,6 +19,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, @@ -2394,27 +2395,19 @@ class BaseLLMHTTPHandler: from litellm.rust_bridge import messages as rust_messages_bridge upstream_body: Final = {key: value for key, value in request_body.items() if key != "stream"} - try: - rust_response: Final = await rust_messages_bridge.amessages( - model=model, - body=upstream_body, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - 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 - verbose_logger.debug( - "Rust Anthropic messages bridge raised %s; falling back to Python path", - type(rust_error).__name__, - ) - return None + rust_response: Final = await rust_messages_bridge.amessages( + model=model, + body=upstream_body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=headers, + timeout=timeout, + ) if rust_response is None: return None response_obj: Final = cast(AnthropicMessagesResponse, dict(rust_response)) - response_obj["_hidden_params"] = {"additional_headers": {"x-litellm-rust": "true"}} return response_obj @staticmethod @@ -2430,7 +2423,7 @@ class BaseLLMHTTPHandler: ) completion_stream = cast(AsyncIterator[bytes], FakeAnthropicMessagesStreamIterator(response=rust_response)) - hidden_params: Final = AnthropicMessagesStreamHiddenParams(additional_headers={"x-litellm-rust": "true"}) + hidden_params: Final = AnthropicMessagesStreamHiddenParams(additional_headers={}) return AnthropicMessagesStreamingResponse( completion_stream=completion_stream, hidden_params=hidden_params, diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 40d0ddf622b..a4890a626d8 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -3,11 +3,19 @@ from __future__ import annotations from collections.abc import Awaitable -from dataclasses import dataclass -from typing import Final, Protocol, cast +from typing import Final, Protocol import httpx +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + FallbackMode, + ainvoke, + async_none, + identity, + invoke, +) from litellm.rust_bridge.timeouts import timeout_to_seconds @@ -39,20 +47,15 @@ class RustAmessages(Protocol): raise NotImplementedError +_MESSAGES: Final = NativeBinding[RustMessages]("messages") +_AMESSAGES: Final = NativeBinding[RustAmessages]("amessages") + + class _Unset: pass -_UNSET: Final[_Unset] = _Unset() - - -@dataclass(slots=True) -class _RustMessagesState: - messages: RustMessages | None = None - amessages: RustAmessages | None = None - - -_STATE: Final[_RustMessagesState] = _RustMessagesState() +_UNSET: Final = _Unset() def set_rust_messages( @@ -61,31 +64,25 @@ def set_rust_messages( amessages: RustAmessages | None | _Unset = _UNSET, ) -> None: if not isinstance(messages, _Unset): - _STATE.messages = messages + _MESSAGES.reset() if messages is None else _MESSAGES.override(messages) if not isinstance(amessages, _Unset): - _STATE.amessages = amessages + _AMESSAGES.reset() if amessages is None else _AMESSAGES.override(amessages) def load_rust_messages() -> RustMessages | None: - if _STATE.messages is not None: - return _STATE.messages - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustMessages, getattr(native_bridge, "messages", None)) + return _MESSAGES.load() def load_rust_amessages() -> RustAmessages | None: - if _STATE.amessages is not None: - return _STATE.amessages - from litellm.rust_bridge import get_native_bridge + return _AMESSAGES.load() - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustAmessages, getattr(native_bridge, "amessages", None)) + +def _context(model: str, custom_llm_provider: str | None) -> BridgeErrorContext: + return BridgeErrorContext( + route="messages", + provider=custom_llm_provider or "anthropic", + model=model, + ) def messages( @@ -98,17 +95,25 @@ def messages( extra_headers: dict[str, object] | None, timeout: float | httpx.Timeout | None, ) -> dict[str, object] | None: - rust_messages: Final = load_rust_messages() - if rust_messages is None: - return None - return rust_messages( - model=model, - body=body, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=timeout_to_seconds(timeout), + native: Final = load_rust_messages() + return invoke( + native_call=( + None + if native is None + else lambda: native( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ) + ), + fallback=lambda: None, + adapt=identity, + mode=FallbackMode.PYTHON, + context=_context(model, custom_llm_provider), ) @@ -122,15 +127,23 @@ async def amessages( extra_headers: dict[str, object] | None, timeout: float | httpx.Timeout | None, ) -> dict[str, object] | None: - rust_amessages: Final = load_rust_amessages() - if rust_amessages is None: - return None - return await rust_amessages( - model=model, - body=body, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout_seconds=timeout_to_seconds(timeout), + native: Final = load_rust_amessages() + return await ainvoke( + native_call=( + None + if native is None + else lambda: native( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ) + ), + fallback=async_none, + adapt=identity, + mode=FallbackMode.PYTHON, + context=_context(model, custom_llm_provider), ) 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 293f75b7592..90a5c8c35f3 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.rust_bridge import configuration from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -100,12 +102,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) @@ -224,7 +242,7 @@ def _gate(**overrides): @pytest.mark.asyncio -async def test_gate_invokes_rust_and_marks_response_header(): +async def test_gate_invokes_rust(): bridge = RecordingAsyncMessages() litellm.use_litellm_rust(True, amessages=bridge) @@ -232,7 +250,6 @@ async def test_gate_invokes_rust_and_marks_response_header(): assert response is not None assert response["id"] == "msg_123" - assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} call = bridge.calls[0] assert call["model"] == "claude-sonnet-4-5" assert call["body"] == REQUEST_BODY @@ -243,8 +260,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() @@ -253,6 +271,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() @@ -301,7 +358,6 @@ async def test_gate_invokes_rust_for_native_anthropic_provider(): ) assert response is not None - assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} assert bridge.calls[0]["custom_llm_provider"] == "anthropic" assert bridge.calls[0]["api_key"] == "sk-ant" @@ -370,7 +426,6 @@ async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag(): ) assert response is not None - assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} assert "stream" not in bridge.calls[0]["body"] assert bridge.calls[0]["body"] == REQUEST_BODY @@ -380,7 +435,7 @@ async def test_fake_stream_wraps_rust_response_as_anthropic_sse(): response = cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE)) stream = BaseLLMHTTPHandler._rust_anthropic_messages_fake_stream(response) - assert stream._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + assert stream._hidden_params["additional_headers"] == {} chunks = [chunk async for chunk in stream] joined = b"".join(chunks) diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 378194191e5..2789def7784 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -4,6 +4,7 @@ import asyncio import importlib.util import json import os +import shutil import signal import subprocess import sys @@ -170,6 +171,19 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: raise AssertionError(f"unknown route: {route}") +def messages_wrapper_kwargs(api_base: str, outcome: str) -> dict[str, object]: + native_kwargs: Final = route_kwargs("messages", api_base, outcome) + return { + "model": native_kwargs["model"], + "body": native_kwargs["body"], + "api_key": native_kwargs["api_key"], + "api_base": native_kwargs["api_base"], + "custom_llm_provider": native_kwargs["custom_llm_provider"], + "extra_headers": native_kwargs["extra_headers"], + "timeout": native_kwargs["timeout_seconds"], + } + + def assert_success(route: str, response: object) -> None: if not isinstance(response, dict): raise TypeError(f"{route} returned {type(response).__name__}, expected dict") @@ -192,7 +206,7 @@ def success_value(route: str, response: dict[object, object]) -> object: def assert_rate_limit(native: object, route: str, error: BaseException) -> None: - if route == "chat_completions": + if route in ("chat_completions", "messages"): upstream_error: Final = native.RustUpstreamError if not isinstance(error, upstream_error) or error.args[0] != 429: raise AssertionError(f"{route} returned the wrong 429 error: {error!r}") @@ -247,6 +261,86 @@ def exercise_routes(native_path: Path, api_base: str) -> object: return native +def assert_packaged_native_loaded(wheel_root: Path) -> None: + from litellm.rust_bridge import get_native_bridge + + native: Final = get_native_bridge() + if native is None: + raise AssertionError("packaged native bridge was not loaded") + native_file: Final = getattr(native, "__file__", None) + if not isinstance(native_file, str): + raise AssertionError("packaged native bridge has no module path") + if wheel_root.resolve() not in Path(native_file).resolve().parents: + raise AssertionError(f"native bridge loaded outside the wheel: {native_file}") + + +async def exercise_packaged_messages(api_base: str) -> None: + from litellm.exceptions import APIError + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.rust_bridge import messages as messages_bridge + from litellm.types.router import GenericLiteLLMParams + + response: Final = messages_bridge.messages(**messages_wrapper_kwargs(api_base, "success")) + assert_success("messages", response) + + responses: Final = await asyncio.wait_for( + asyncio.gather( + *( + messages_bridge.amessages(**messages_wrapper_kwargs(api_base, "success")) + for _ in range(32) + ) + ), + timeout=15, + ) + for concurrent_response in responses: + assert_success("messages", concurrent_response) + + declined_kwargs: Final = messages_wrapper_kwargs(api_base, "success") + declined_kwargs["custom_llm_provider"] = "openai" + if await messages_bridge.amessages(**declined_kwargs) is not None: + raise AssertionError("unsupported Messages request did not decline") + + params: Final = GenericLiteLLMParams(api_key="sk-native", rust=True) + success_kwargs: Final = route_kwargs("messages", api_base, "success") + gate_response: Final = await BaseLLMHTTPHandler._maybe_rust_anthropic_messages( + custom_llm_provider="anthropic", + litellm_params=params, + has_agentic_hook=False, + model=str(success_kwargs["model"]), + api_key="sk-native", + api_base=api_base, + headers=dict(success_kwargs["extra_headers"]), + request_body=dict(success_kwargs["body"]), + timeout=3.0, + ) + assert_success("messages", gate_response) + + rate_limit_kwargs: Final = route_kwargs("messages", api_base, "429") + try: + await BaseLLMHTTPHandler._maybe_rust_anthropic_messages( + custom_llm_provider="anthropic", + litellm_params=params, + has_agentic_hook=False, + model=str(rate_limit_kwargs["model"]), + api_key="sk-native", + api_base=api_base, + headers=dict(rate_limit_kwargs["extra_headers"]), + request_body=dict(rate_limit_kwargs["body"]), + timeout=3.0, + ) + except APIError as error: + if error.status_code != 429 or "native-rate-limit" not in str(error): + raise AssertionError(f"Messages gate returned the wrong upstream error: {error!r}") from error + else: + raise AssertionError("Messages gate retried or swallowed the native upstream error") + + +def exercise_packaged_python_bridge(wheel_root: Path, api_base: str) -> int: + assert_packaged_native_loaded(wheel_root) + asyncio.run(exercise_packaged_messages(api_base)) + return 0 + + def exercise_signal(native: object, api_base: str) -> int: try: native.messages( @@ -293,6 +387,40 @@ def verify_sigint(native_path: Path, api_base: str) -> None: process.wait(timeout=5) +def verify_packaged_python_bridge(wheel: Path, wheel_root: Path, api_base: str) -> None: + uv: Final = shutil.which("uv") + if uv is None: + raise AssertionError("uv is required to test the packaged Python bridge") + environment: Final = {key: value for key, value in os.environ.items() if key != "ANTHROPIC_API_KEY"} | { + "PYTHONPATH": str(wheel_root) + } + result: Final = subprocess.run( + ( + uv, + "run", + "--isolated", + "--with", + str(wheel.resolve()), + "python", + __file__, + "bridge-child", + str(wheel_root), + api_base, + ), + cwd=wheel_root, + env=environment, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + if result.returncode != 0: + raise AssertionError( + f"packaged Python bridge failed with status {result.returncode}" + f"\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + + def verify_wheel(wheel: Path) -> int: with tempfile.TemporaryDirectory() as temporary_directory, zipfile.ZipFile(wheel) as archive: wheel_root: Final = Path(temporary_directory) @@ -318,6 +446,7 @@ def verify_wheel(wheel: Path) -> int: api_base: Final = f"http://127.0.0.1:{server.server_address[1]}" try: verify_sigint(native_path, api_base) + verify_packaged_python_bridge(wheel, wheel_root, api_base) finally: server.shutdown() server.server_close() @@ -331,6 +460,8 @@ def main() -> int: if len(sys.argv) == 4 and sys.argv[1] == "child": native: Final = exercise_routes(Path(sys.argv[2]), sys.argv[3]) return exercise_signal(native, sys.argv[3]) + if len(sys.argv) == 4 and sys.argv[1] == "bridge-child": + return exercise_packaged_python_bridge(Path(sys.argv[2]), sys.argv[3]) sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") return 2