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 42ffd4cbcff..65008fd8add 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -277,6 +277,7 @@ fn core_error_kind(error: &Error) -> &'static str { Error::Http { .. } => "HttpError", Error::InvalidResponse(_) => "InvalidResponse", Error::Network(_) => "NetworkError", + Error::Timeout(_) => "TimeoutError", 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 6343122b9e9..7035a17652f 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -321,6 +321,7 @@ fn core_error_kind(error: &Error) -> &'static str { Error::Http { .. } => "HttpError", Error::InvalidResponse(_) => "InvalidResponse", Error::Network(_) => "NetworkError", + Error::Timeout(_) => "TimeoutError", 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 e4861c3954c..b1d2187166d 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -105,6 +105,7 @@ impl IntoResponse for MessagesRouteError { ), Error::Http { .. } | Error::Network(_) + | Error::Timeout(_) | Error::InvalidResponse(_) | Error::InvalidType { .. } | Error::MissingField(_) => ( diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 30ba0da5e68..8b1e0b9e32b 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,7 +1,7 @@ use serde_json::Value; use crate::error::Error; -use crate::http_utils::truncate_error_body; +use crate::http_utils::{classify_transport_error, truncate_error_body}; use super::client::http_client; use super::types::ProviderAudioTranscriptionRequest; @@ -22,12 +22,9 @@ pub async fn execute_audio_transcription_provider_call( let response = request_builder .send() .await - .map_err(|error| Error::Network(error.to_string()))?; + .map_err(classify_transport_error)?; let status = response.status(); - let text = response - .text() - .await - .map_err(|error| Error::Network(error.to_string()))?; + let text = response.text().await.map_err(classify_transport_error)?; if !status.is_success() { return Err(Error::Http { status: status.as_u16(), diff --git a/litellm-rust/crates/core/src/audio_transcription/tests.rs b/litellm-rust/crates/core/src/audio_transcription/tests.rs index 263d63337b0..651dbfa207d 100644 --- a/litellm-rust/crates/core/src/audio_transcription/tests.rs +++ b/litellm-rust/crates/core/src/audio_transcription/tests.rs @@ -1,11 +1,13 @@ use std::io::{Read, Write}; use std::net::TcpListener; use std::thread; +use std::time::Duration; use serde_json::{Map, json}; use super::audio_transcription; use super::types::AudioTranscriptionRequest; +use crate::error::Error; #[tokio::test] async fn bedrock_request_is_signed_and_contains_audio() { @@ -48,3 +50,47 @@ async fn bedrock_request_is_signed_and_contains_audio() { assert_eq!(response, json!({"text": "hello"})); server.join().expect("server"); } + +#[tokio::test] +async fn an_upstream_that_never_responds_times_out() { + let listener = TcpListener::bind("127.0.0.1:0").expect("listener"); + let address = listener.local_addr().expect("address"); + thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("connection"); + let mut buffer = [0_u8; 16_384]; + loop { + match stream.read(&mut buffer) { + Ok(0) | Err(_) => return, + Ok(_) => {} + } + } + }); + + let optional_params = Map::from_iter([ + ("aws_access_key_id".to_string(), json!("access-key")), + ("aws_secret_access_key".to_string(), json!("secret-key")), + ("aws_region_name".to_string(), json!("us-east-1")), + ]); + let api_base = format!("http://{address}"); + let error = tokio::time::timeout( + Duration::from_secs(2), + audio_transcription(AudioTranscriptionRequest { + model: "mistral.voxtral-mini-3b-2507", + audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}), + api_key: None, + api_base: Some(&api_base), + custom_llm_provider: Some("bedrock"), + extra_headers: None, + optional_params, + timeout: Some(Duration::from_millis(100)), + }), + ) + .await + .expect("client call completes") + .expect_err("upstream never responds"); + + assert!( + matches!(error, Error::Timeout(_)), + "expected a timeout failure, got {error:?}" + ); +} diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index f282e82c4fa..63db38e511b 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, as_response_error}; -use crate::http_utils::{classify_send_error, truncate_error_body}; +use crate::http_utils::{classify_transport_error, truncate_error_body}; use super::client::http_client; use super::transformation::ChatCompletionsAuth; @@ -27,13 +27,13 @@ pub(super) async fn execute_chat_completions_provider_call( request_builder = request_builder.timeout(duration); } - let response = request_builder.send().await.map_err(classify_send_error)?; + let response = request_builder + .send() + .await + .map_err(classify_transport_error)?; let status = response.status(); - let text = response - .text() - .await - .map_err(|err| Error::Network(err.to_string()))?; + let text = response.text().await.map_err(classify_transport_error)?; if !status.is_success() { return Err(Error::Http { diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index bdd413d7cff..5e3a4dfce11 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -787,4 +787,36 @@ mod round_trip { "expected a terminal network failure, got {err:?}" ); } + + #[tokio::test] + async fn an_established_request_that_never_responds_times_out() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("has an address"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _ = read_http_request(&mut socket).await; + std::future::pending::<()>().await; + }); + + let api_base = format!("http://{addr}/v1/messages"); + let mut timed_out_call = call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + ); + timed_out_call.timeout = Some(std::time::Duration::from_millis(100)); + let err = tokio::time::timeout( + std::time::Duration::from_secs(2), + chat_completions(timed_out_call), + ) + .await + .expect("client call completes") + .expect_err("established request times out"); + + server.abort(); + assert!( + matches!(err, Error::Timeout(_)), + "expected a timeout failure, got {err:?}" + ); + } } diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index a5de4b722d2..a7f20387f61 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -21,6 +21,8 @@ pub enum Error { Http { status: u16, body: String }, #[error("upstream network error: {0}")] Network(String), + #[error("upstream request timed out: {0}")] + Timeout(String), #[error("routing error: {0}")] Routing(String), /// The request is outside the surface this route covers in Rust. Hosts that @@ -59,6 +61,8 @@ mod tests { Error::Unsupported("non-text response content block"), Error::InvalidRequest("whatever".to_string()), Error::Auth("whatever".to_string()), + Error::Network("whatever".to_string()), + Error::Timeout("whatever".to_string()), ] { assert!(matches!( as_response_error(original), diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index 759c7ba5652..b79e74ead59 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -5,7 +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 { +pub(crate) fn classify_transport_error(error: reqwest::Error) -> Error { + if error.is_timeout() { + return Error::Timeout(error.to_string()); + } Error::Network(error.to_string()) } diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 441f133b182..89453362d97 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,6 +1,6 @@ use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; use crate::error::{Error, as_response_error}; -use crate::http_utils::classify_send_error; +use crate::http_utils::classify_transport_error; use super::client::http_client; use super::common_utils::truncate_error_body; @@ -17,13 +17,13 @@ pub(super) async fn execute_messages_provider_call( request_builder = request_builder.timeout(duration); } - let response = request_builder.send().await.map_err(classify_send_error)?; + let response = request_builder + .send() + .await + .map_err(classify_transport_error)?; let status = response.status(); - let text = response - .text() - .await - .map_err(|err| Error::Network(err.to_string()))?; + let text = response.text().await.map_err(classify_transport_error)?; if !status.is_success() { return Err(Error::Http { @@ -60,13 +60,10 @@ pub(super) async fn execute_messages_provider_stream( let response = request_builder .send() .await - .map_err(|err| Error::Network(err.to_string()))?; + .map_err(classify_transport_error)?; let status = response.status(); if !status.is_success() { - let text = response - .text() - .await - .map_err(|err| Error::Network(err.to_string()))?; + let text = response.text().await.map_err(classify_transport_error)?; return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), @@ -193,7 +190,7 @@ mod tests { } #[tokio::test] - async fn established_request_timeouts_are_network_errors() { + async fn established_request_timeouts_are_timeout_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(); @@ -223,6 +220,31 @@ mod tests { assert!(received.starts_with("POST / "), "{received}"); release_server_tx.send(()).expect("releases server"); server.await.expect("server task completes"); - assert!(matches!(error, Error::Network(_))); + assert!(matches!(error, Error::Timeout(_))); + } + + #[tokio::test] + async fn stream_request_timeouts_are_timeout_errors() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("has an address"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _ = read_http_request(&mut socket).await; + std::future::pending::<()>().await; + }); + + let error = tokio::time::timeout( + Duration::from_secs(2), + execute_messages_provider_stream(request( + format!("http://{addr}"), + Duration::from_millis(100), + )), + ) + .await + .expect("client call completes") + .expect_err("established stream times out"); + + server.abort(); + assert!(matches!(error, Error::Timeout(_))); } } diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index d3e47937666..84201110b88 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -13,7 +13,7 @@ pyo3::create_exception!( _native, RustUpstreamError, pyo3::exceptions::PyException, - "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." + "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response and 408 when the request timed out before one arrived." ); pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { @@ -37,6 +37,9 @@ pub(crate) fn fallback_route_error_to_pyerr(err: Error) -> PyErr { Error::Http { status, body } => { RustUpstreamError::new_err((status, format!("{status}: {body}"))) } + // 408 is the convention Python reads back to classify the failure as a + // timeout, mirroring an upstream `408 Request Timeout` response. + Error::Timeout(_) => RustUpstreamError::new_err((408u16, err.to_string())), other => RustUpstreamError::new_err((0u16, other.to_string())), } } @@ -73,6 +76,10 @@ mod tests { Error::Network("request timed out".to_string()), (0, "upstream network error: request timed out"), ), + ( + Error::Timeout("request timed out".to_string()), + (408, "upstream request timed out: request timed out"), + ), ( Error::InvalidResponse("bad JSON".to_string()), (0, "invalid response: bad JSON"), diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index c599667ab17..93f5ce74bac 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -21,13 +21,13 @@ import httpx from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger -from litellm.exceptions import APIError from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( convert_to_model_response_object, ) from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge.runtime import BridgeErrorContext, raise_upstream from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.types.utils import ModelResponse @@ -297,8 +297,9 @@ def _reraise_or_decline( A request that never reached the provider is safe to serve on the Python path. One that did is not: the provider has already done the work, so a - second attempt bills for it twice. Those surface as an `APIError` carrying - the upstream status, which LiteLLM's exception mapping already understands. + second attempt bills for it twice. Those surface as the litellm exception + matching the upstream status, so router cooldown and retry logic keyed on + exception type behaves as it does on the Python path. """ exceptions: Final = _rust_bridge_exceptions() if exceptions is None: @@ -309,14 +310,13 @@ def _reraise_or_decline( return declined, upstream_failed = exceptions if 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 chat completions: {message}", - llm_provider=custom_llm_provider or "", - model=model, + raise_upstream( + rust_error, + BridgeErrorContext( + route="chat completions", + provider=custom_llm_provider or "", + model=model, + ), ) if not isinstance(rust_error, declined): raise rust_error diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index 0bf60b1aa0b..1da51825e83 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -5,7 +5,15 @@ from dataclasses import dataclass from enum import Enum from typing import Final, Generic, NoReturn, TypeAlias, TypeVar, cast -from litellm.exceptions import APIError +from litellm.exceptions import ( + APIError, + AuthenticationError, + BadRequestError, + NotFoundError, + RateLimitError, + ServiceUnavailableError, + Timeout, +) from litellm.rust_bridge.bindings import native_exception_types NativeT = TypeVar("NativeT") @@ -91,7 +99,7 @@ def attempt( except declined as error: return RustDeclined(reason=_decline_reason(error)) except upstream as error: - _raise_upstream(error, context) + raise_upstream(error, context) return RustHandled(adapt(value)) @@ -112,7 +120,7 @@ async def aattempt( except declined as error: return RustDeclined(reason=_decline_reason(error)) except upstream as error: - _raise_upstream(error, context) + raise_upstream(error, context) return RustHandled(adapt(value)) @@ -124,7 +132,7 @@ def call(operation: Callable[[], ResultT], context: BridgeErrorContext) -> Resul try: return operation() except upstream as error: - _raise_upstream(error, context) + raise_upstream(error, context) async def acall(operation: Callable[[], Awaitable[ResultT]], context: BridgeErrorContext) -> ResultT: @@ -135,7 +143,7 @@ async def acall(operation: Callable[[], Awaitable[ResultT]], context: BridgeErro try: return await operation() except upstream as error: - _raise_upstream(error, context) + raise_upstream(error, context) def _decline_reason(error: BaseException) -> str: @@ -158,18 +166,57 @@ def _required_reason(result: RustDeclined | RustUnavailable) -> str: return f"declined the request: {reason}" -def _raise_upstream(error: BaseException, context: BridgeErrorContext) -> NoReturn: +def raise_upstream(error: BaseException, context: BridgeErrorContext) -> NoReturn: args: Final = cast(tuple[object, ...], error.args) status_value: Final = args[0] if args else 0 message_value: Final = args[1] if len(args) > 1 else str(error) status: Final = status_value if isinstance(status_value, int) else 0 message: Final = message_value if isinstance(message_value, str) else str(message_value) - raise APIError( - status_code=status or 500, - message=f"litellm rust {context.route}: {message}", - llm_provider=context.provider, - model=context.model, - ) from error + prefixed: Final = f"litellm rust {context.route}: {message}" + match status: + case 401 | 403: + raise AuthenticationError( + message=prefixed, + llm_provider=context.provider, + model=context.model, + ) from error + case 404: + raise NotFoundError( + message=prefixed, + model=context.model, + llm_provider=context.provider, + ) from error + case 408: + raise Timeout( + message=prefixed, + model=context.model, + llm_provider=context.provider, + ) from error + case 429: + raise RateLimitError( + message=prefixed, + llm_provider=context.provider, + model=context.model, + ) from error + case _ if 400 <= status < 500: + raise BadRequestError( + message=prefixed, + model=context.model, + llm_provider=context.provider, + ) from error + case _ if status >= 500: + raise ServiceUnavailableError( + message=prefixed, + llm_provider=context.provider, + model=context.model, + ) from error + case _: + raise APIError( + status_code=status or 500, + message=prefixed, + llm_provider=context.provider, + model=context.model, + ) from error def identity(value: ResultT) -> ResultT: 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 2789def7784..b3089a8b269 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -275,7 +275,7 @@ def assert_packaged_native_loaded(wheel_root: Path) -> None: async def exercise_packaged_messages(api_base: str) -> None: - from litellm.exceptions import APIError + from litellm.exceptions import RateLimitError 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 @@ -328,7 +328,7 @@ async def exercise_packaged_messages(api_base: str) -> None: request_body=dict(rate_limit_kwargs["body"]), timeout=3.0, ) - except APIError as error: + except RateLimitError 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: diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 03921133c77..46d0e96cbbc 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -351,14 +351,35 @@ class TestFailureClassification: bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming"))) assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None - def test_an_upstream_failure_is_surfaced_with_its_status(self): - from litellm.exceptions import APIError + def test_an_upstream_rate_limit_surfaces_as_litellm_rate_limit_error(self): + from litellm.exceptions import RateLimitError bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited"))) - with pytest.raises(APIError) as raised: + with pytest.raises(RateLimitError) as raised: bridge.chat_completions(**_call_kwargs(ModelResponse())) assert raised.value.status_code == 429 assert "rate limited" in str(raised.value) + assert isinstance(raised.value.__cause__, _FakeUpstream) + + def test_an_upstream_auth_failure_surfaces_as_litellm_authentication_error(self): + from litellm.exceptions import AuthenticationError + + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(401, "401: bad key"))) + with pytest.raises(AuthenticationError) as raised: + bridge.chat_completions(**_call_kwargs(ModelResponse())) + assert raised.value.status_code == 401 + assert isinstance(raised.value.__cause__, _FakeUpstream) + + def test_an_upstream_timeout_surfaces_as_litellm_timeout(self): + from litellm.exceptions import Timeout + + bridge.set_rust_chat_completions( + chat_completions=_RecordingCall(error=_FakeUpstream(408, "upstream request timed out: 30s")) + ) + with pytest.raises(Timeout) as raised: + bridge.chat_completions(**_call_kwargs(ModelResponse())) + assert raised.value.status_code == 408 + assert isinstance(raised.value.__cause__, _FakeUpstream) def test_a_transport_failure_with_no_response_surfaces_as_a_500(self): from litellm.exceptions import APIError @@ -367,6 +388,7 @@ class TestFailureClassification: with pytest.raises(APIError) as raised: bridge.chat_completions(**_call_kwargs(ModelResponse())) assert raised.value.status_code == 500 + assert isinstance(raised.value.__cause__, _FakeUpstream) def test_an_unrecognized_error_is_not_swallowed(self): bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=RuntimeError("something else"))) @@ -375,7 +397,7 @@ class TestFailureClassification: @pytest.mark.asyncio async def test_the_async_wrapper_does_not_fall_back_on_an_upstream_failure(self): - from litellm.exceptions import APIError + from litellm.exceptions import ServiceUnavailableError bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom"))) ran = [] @@ -384,7 +406,7 @@ class TestFailureClassification: ran.append(True) return "python" - with pytest.raises(APIError): + with pytest.raises(ServiceUnavailableError): await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert ran == [], "a request the provider already served must not be re-issued" diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index b0fa510069b..bd4fd308415 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -1,10 +1,20 @@ from __future__ import annotations from types import SimpleNamespace +from typing import Final +import openai import pytest -from litellm.exceptions import APIError +from litellm.exceptions import ( + APIError, + AuthenticationError, + BadRequestError, + NotFoundError, + RateLimitError, + ServiceUnavailableError, + Timeout, +) from litellm.rust_bridge import bindings, runtime @@ -52,7 +62,7 @@ def test_invoke_translates_upstream_without_fallback() -> None: def fail() -> object: raise RustUpstreamError(429, "rate limited") - with pytest.raises(APIError, match="rate limited") as caught: + with pytest.raises(RateLimitError, match="rate limited") as caught: runtime.invoke( native_call=fail, fallback=lambda: pytest.fail("fallback must not run"), @@ -93,3 +103,94 @@ def test_required_mode_rejects_unavailable_bridge() -> None: mode=runtime.FallbackMode.RUST_REQUIRED, context=context(), ) + + +class TestUpstreamStatusClassification: + """The router's cooldown, retry, and fallback logic keys on litellm exception + types, so a Rust-surfaced upstream status must raise the same exception the + Python path would""" + + @pytest.mark.parametrize( + ("status", "expected"), + ( + (400, BadRequestError), + (401, AuthenticationError), + (403, AuthenticationError), + (404, NotFoundError), + (408, Timeout), + (422, BadRequestError), + (429, RateLimitError), + (500, ServiceUnavailableError), + (503, ServiceUnavailableError), + (0, APIError), + ), + ) + def test_each_status_raises_its_litellm_exception(self, status: int, expected: type[Exception]) -> None: + def fail() -> object: + raise RustUpstreamError(status, f"{status}: upstream said no") + + with pytest.raises(expected, match="upstream said no") as caught: + runtime.invoke( + native_call=fail, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), + ) + + assert caught.value.__cause__ is not None + assert isinstance(caught.value.__cause__, RustUpstreamError) + + def test_the_timeout_marker_raises_litellm_timeout(self) -> None: + def hang() -> object: + raise RustUpstreamError(408, "upstream request timed out: 30s elapsed") + + with pytest.raises(Timeout) as caught: + runtime.invoke( + native_call=hang, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), + ) + + assert caught.value.status_code == 408 + assert isinstance(caught.value.__cause__, RustUpstreamError) + assert "timed out" in str(caught.value) + + def test_the_cause_chain_preserves_the_rust_error(self) -> None: + cause: Final = RustUpstreamError(500, "500: boom") + + def fail() -> object: + raise cause + + with pytest.raises(ServiceUnavailableError) as caught: + runtime.invoke( + native_call=fail, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), + ) + + assert caught.value.__cause__ is cause + + def test_a_rust_429_is_classified_the_way_router_cooldown_expects(self) -> None: + from litellm.router_utils.cooldown_handlers import _is_cooldown_required + + def fail() -> object: + raise RustUpstreamError(429, "429: slow down") + + with pytest.raises(RateLimitError) as caught: + runtime.invoke( + native_call=fail, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), + ) + + error: Final = caught.value + assert isinstance(error, openai.RateLimitError), "type-based retry and fallback gates must see it" + assert error.status_code == 429 + assert _is_cooldown_required(None, "deployment", error.status_code) is True