From c548955e9da953e2a46fc07738d09c92ba4e4cdc Mon Sep 17 00:00:00 2001 From: narmstrong-berri Date: Thu, 5 Nov 2026 11:49:31 -0800 Subject: [PATCH] fix(messages): preserve Rust request contracts and stream lifecycle --- litellm-rust/Cargo.lock | 1 + .../gateway-inference/tests/messages.rs | 145 ++++++++- .../gateway-inference/tests/responses.rs | 2 +- litellm-rust/crates/host-http/src/sse.rs | 3 +- litellm-rust/crates/host-http/tests/sse.rs | 26 +- .../crates/inference-messages/Cargo.toml | 2 +- .../inference-messages/src/common_utils.rs | 21 +- .../crates/inference-messages/src/handler.rs | 104 ++++++- .../crates/inference-messages/src/lib.rs | 2 +- .../inference-messages/tests/caching.rs | 282 ++++++++++++++++++ .../inference-messages/tests/messages/host.rs | 167 +++++++++++ .../tests/messages/stream.rs | 38 ++- .../src/formats/messages/request.rs | 108 +++++++ .../src/formats/messages/streaming.rs | 11 + litellm-rust/crates/llms/Cargo.toml | 1 + .../crates/llms/src/anthropic/common_utils.rs | 92 +++++- .../llms/src/anthropic/messages/handler.rs | 1 + .../llms/src/anthropic/messages/thinking.rs | 51 +++- .../src/anthropic/messages/transformation.rs | 29 +- .../messages/mid_conversation_system.rs | 66 ++++ .../src/base_llm/messages/normalization.rs | 89 ++++++ .../messages/transformation.rs | 88 +++++- .../anthropic_claude_transformation.rs | 55 ++++ .../src/edenai/messages/transformation.rs | 128 +++++++- .../src/openrouter/messages/transformation.rs | 61 ++++ .../src/tencent/messages/transformation.rs | 35 ++- .../crates/python-bridge/src/errors.rs | 4 + .../python-bridge/src/routes/messages/host.rs | 12 +- .../python-bridge/src/routes/messages/mod.rs | 15 + litellm/messages/dispatch.py | 3 +- litellm/router_rust/invoke/streaming.py | 16 +- litellm/rust_bridge/host/routes/messages.py | 24 +- .../sdk/messages/test_native_parity.py | 114 ++++++- .../cache/test_python_cache.py | 105 ++++++- tests/test_litellm_rust/conftest.py | 10 +- .../messages/test_callbacks.py | 171 ++++++++++- .../messages/test_dispatch.py | 62 ++++ .../messages/test_lifecycle.py | 93 ++++++ .../messages/test_requests.py | 42 +++ .../messages/test_streaming.py | 30 +- .../support/recording_server.py | 7 + tests/test_litellm_rust/support/served.py | 102 ++++++- .../rust_bridge/host/routes/test_messages.py | 11 + tests/unit/test_router/test_router.py | 103 ++++++- 44 files changed, 2446 insertions(+), 86 deletions(-) create mode 100644 tests/test_litellm_rust/messages/test_lifecycle.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 107b55e89cb..853f05501e2 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4257,6 +4257,7 @@ dependencies = [ "thiserror 2.0.19", "time", "tokio", + "unicode-general-category", "url", ] diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs index 7d6d1f160d1..9ccb4a3231f 100644 --- a/litellm-rust/crates/gateway-inference/tests/messages.rs +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -4,8 +4,10 @@ use axum::{ body::{Body, to_bytes}, http::Request, }; +use futures_util::{StreamExt, TryStreamExt}; use rstest::rstest; use serde_json::json; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; use wiremock::{ Mock, MockServer, ResponseTemplate, @@ -75,25 +77,76 @@ async fn messages_reaches_the_provider_and_preserves_json_or_sse( assert!(!requests[0].headers.contains_key("x-request-id")); } +#[rstest] +#[case::invalid_json("{".into())] +#[case::array("[]".into())] +#[case::null("null".into())] +#[case::missing_model(json!({"messages": [], "max_tokens": 16}).to_string())] +#[case::unknown_model(json!({"model": "missing/model", "messages": [], "max_tokens": 16}).to_string())] +#[case::invalid_messages(json!({"model": "public/model", "messages": "invalid", "max_tokens": 16}).to_string())] +#[case::missing_max_tokens(json!({"model": "public/model", "messages": [{"role": "user", "content": "hello"}]}).to_string())] #[tokio::test] -async fn invalid_messages_stays_an_anthropic_error() { - let response = support::post( - support::app("anthropic/test-model", "http://127.0.0.1:1"), - "/v1/messages", - json!({"model": "public/model", "messages": "invalid", "max_tokens": 16}), - ) - .await; +async fn invalid_messages_stays_an_anthropic_error_without_provider_io(#[case] body: String) { + let upstream = MockServer::start().await; + let request = Request::post("/v1/messages") + .header("content-type", "application/json") + .header("x-request-id", "invalid-messages-request") + .body(Body::from(body)) + .unwrap(); + let response = support::app("anthropic/test-model", &upstream.uri()) + .oneshot(request) + .await + .unwrap(); + assert_eq!(response.status(), 400); let body = support::json(response).await; assert_eq!(body["type"], "error"); assert_eq!(body["error"]["type"], "invalid_request_error"); + assert_eq!(body["request_id"], "invalid-messages-request"); + assert!(!body["error"]["message"].as_str().unwrap().is_empty()); + assert!(upstream.received_requests().await.unwrap().is_empty()); } const FIRST_EVENT: &str = "event: message_start\ndata: {}\n\n"; +async fn interrupted_stream( + prefix: &'static [u8], +) -> ( + String, + tokio::sync::oneshot::Sender<()>, + tokio::task::JoinHandle<()>, +) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let (release, released) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n", + prefix.len() + 1024 + ).as_bytes() + ).await.unwrap(); + socket.write_all(prefix).await.unwrap(); + let _ = released.await; + }); + (base, release, task) +} + +#[rstest] +#[case::complete_lf(FIRST_EVENT.as_bytes())] +#[case::complete_crlf(b"event: message_start\r\ndata: {}\r\n\r\n")] +#[case::partial_event(b"event: content_blo")] +#[case::partial_data(b"event: content_block_delta\ndata: {\"type\":\"content_block_delta\"")] +#[case::partial_string(b"event: content_block_delta\ndata: {\"delta\":{\"text\":\"partial")] +#[case::partial_utf8(b"event: content_block_delta\ndata: {\"delta\":{\"text\":\"\xe9")] #[tokio::test] -async fn a_stream_that_fails_after_opening_ends_with_an_sse_error_frame() { - let base = support::truncating_upstream(FIRST_EVENT).await; +async fn a_stream_that_fails_after_opening_ends_with_an_sse_error_frame( + #[case] prefix: &'static [u8], +) { + let (base, release, upstream) = interrupted_stream(prefix).await; let request = Request::post("/v1/messages") .header("content-type", "application/json") .body(Body::from( @@ -109,12 +162,26 @@ async fn a_stream_that_fails_after_opening_ends_with_an_sse_error_frame() { .unwrap(); assert_eq!(response.status(), 200); - let body = to_bytes(response.into_body(), 4096).await.unwrap(); - let text = std::str::from_utf8(&body).unwrap(); + let mut chunks = response.into_body().into_data_stream(); + let mut delivered = Vec::with_capacity(prefix.len()); + while delivered.len() < prefix.len() { + let chunk = tokio::time::timeout(std::time::Duration::from_secs(5), chunks.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + delivered.extend_from_slice(&chunk); + } + assert_eq!(delivered, prefix); + release.send(()).unwrap(); + let tail = chunks.try_collect::>().await.unwrap().concat(); + upstream.await.unwrap(); + let text = std::str::from_utf8(&tail).unwrap(); let frame = text - .strip_prefix(FIRST_EVENT) - .and_then(|rest| rest.strip_prefix("event: error\ndata: ")) - .unwrap_or_else(|| panic!("the delivered event then one error frame, got {text:?}")); + .strip_prefix("\n\nevent: error\ndata: ") + .unwrap_or_else(|| { + panic!("an independent error frame after the delivered prefix, got {text:?}") + }); let error: serde_json::Value = serde_json::from_str(frame.trim_end()).unwrap(); assert_eq!(error["type"], "error"); assert_eq!(error["error"]["type"], "api_error"); @@ -217,3 +284,53 @@ async fn openrouter_receives_native_messages_headers_and_its_deployment_bearer( assert_eq!(requests.len(), 1); assert!(!requests[0].headers.contains_key("x-api-key")); } + +#[rstest] +#[case::u64_max("18446744073709551615".to_string(), 200)] +#[case::above_u64("18446744073709551617".to_string(), 500)] +#[case::i64_min("-9223372036854775808".to_string(), 200)] +#[case::below_i64("-9223372036854775809".to_string(), 500)] +#[case::long_numeric_string(format!(r#""{}""#, "9".repeat(4096)), 200)] +#[case::decimal("18446744073709551617.0".to_string(), 200)] +#[case::exponent("1e30".to_string(), 200)] +#[tokio::test] +async fn messages_response_integer_overflow_is_a_provider_error_not_rounded_success( + #[case] token: String, + #[case] status: u16, +) { + let upstream = MockServer::start().await; + let body = format!( + r#"{{"id":"m","type":"message","role":"assistant","model":"test-model","content":[],"future_extension":{{"wide_integer":{token}}}}}"#, + ); + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_raw(body, "application/json")) + .expect(1) + .mount(&upstream) + .await; + let request = Request::post("/v1/messages") + .header("content-type", "application/json") + .body(Body::from(json!({"model":"public/model","messages":[{"role":"user","content":"hi"}],"max_tokens":16}).to_string())) + .unwrap(); + let response = support::app("anthropic/test-model", &upstream.uri()) + .oneshot(request) + .await + .unwrap(); + assert_eq!(response.status(), status); + let response = support::json(response).await; + if status == 200 { + assert_eq!( + response["future_extension"]["wide_integer"], + serde_json::from_str::(&token).unwrap() + ); + } else { + assert_eq!(response["type"], "error"); + assert_eq!(response["error"]["type"], "api_error"); + assert!( + response["error"]["message"] + .as_str() + .unwrap() + .contains("response integer outside") + ); + } + assert_eq!(upstream.received_requests().await.unwrap().len(), 1); +} diff --git a/litellm-rust/crates/gateway-inference/tests/responses.rs b/litellm-rust/crates/gateway-inference/tests/responses.rs index ed7c8e6fe76..8ee3c68f3fc 100644 --- a/litellm-rust/crates/gateway-inference/tests/responses.rs +++ b/litellm-rust/crates/gateway-inference/tests/responses.rs @@ -134,7 +134,7 @@ async fn a_stream_that_fails_after_opening_ends_with_the_responses_error_event() let text = std::str::from_utf8(&body).unwrap(); let frame = text .strip_prefix(RESPONSE_CREATED) - .and_then(|rest| rest.strip_prefix("event: error\ndata: ")) + .and_then(|rest| rest.strip_prefix("\n\nevent: error\ndata: ")) .unwrap_or_else(|| panic!("the delivered event then one error frame, got {text:?}")); let error: serde_json::Value = serde_json::from_str(frame.trim_end()).unwrap(); assert_eq!(error["type"], "error"); diff --git a/litellm-rust/crates/host-http/src/sse.rs b/litellm-rust/crates/host-http/src/sse.rs index 8f7c8ffa3b3..48e3837d085 100644 --- a/litellm-rust/crates/host-http/src/sse.rs +++ b/litellm-rust/crates/host-http/src/sse.rs @@ -53,6 +53,7 @@ where } fn encode_stream_error(&self, error: Error) -> Bytes { - (self.stream_error)(error) + let frame = (self.stream_error)(error); + Bytes::from([b"\n\n".as_slice(), frame.as_ref()].concat()) } } diff --git a/litellm-rust/crates/host-http/tests/sse.rs b/litellm-rust/crates/host-http/tests/sse.rs index af2bb404207..6cace273068 100644 --- a/litellm-rust/crates/host-http/tests/sse.rs +++ b/litellm-rust/crates/host-http/tests/sse.rs @@ -40,11 +40,15 @@ impl Protocol for TestProtocol { } #[rstest] -#[case::complete(false)] -#[case::failed(true)] +#[case::complete(false, "event: custom\ndata: first\n\n")] +#[case::failed_after_frame(true, "event: custom\ndata: first\n\n")] +#[case::failed_during_frame(true, "event: custom\ndata: partial")] #[tokio::test] -async fn sse_preserves_encoded_chunks_and_uses_the_supplied_error_format(#[case] fail: bool) { - let first = Bytes::from_static(b"event: custom\ndata: first\n\n"); +async fn sse_preserves_encoded_chunks_and_delimits_the_supplied_error_frame( + #[case] fail: bool, + #[case] prefix: &'static str, +) { + let first = Bytes::from_static(prefix.as_bytes()); let last = Bytes::from_static(b"data: [DONE]\n\n"); let machine = hosted_call::((), None, move |(), _, _, _observations| async move { @@ -70,14 +74,12 @@ async fn sse_preserves_encoded_chunks_and_uses_the_supplied_error_format(#[case] assert_eq!(response.headers()[CONTENT_TYPE], "text/event-stream"); assert_eq!(errors.load(Ordering::SeqCst), 0); let body = to_bytes(response.into_body(), 1024).await.unwrap(); - assert_eq!( - body, - if fail { - "event: custom\ndata: first\n\nevent: custom_error\ndata: Call(Upstream)\n\n" - } else { - "event: custom\ndata: first\n\ndata: [DONE]\n\n" - } - ); + let expected = if fail { + format!("{prefix}\n\nevent: custom_error\ndata: Call(Upstream)\n\n") + } else { + format!("{prefix}data: [DONE]\n\n") + }; + assert_eq!(body.as_ref(), expected.as_bytes()); assert_eq!(errors.load(Ordering::SeqCst), usize::from(fail)); } diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index 0c5be52550d..3874faa2fd3 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -23,7 +23,7 @@ litellm-secrets.workspace = true litellm-tracing.workspace = true reqwest.workspace = true serde.workspace = true -serde_json = { workspace = true, features = ["preserve_order"] } +serde_json = { workspace = true, features = ["preserve_order", "raw_value"] } tokio.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index e5ece2837d8..0643f9dab1a 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -72,6 +72,12 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option) -> bool { + litellm_inference::provider::resolve_llm_provider(model, custom_llm_provider, "messages") + .ok() + .is_some_and(|provider| messages_provider(provider.provider, provider.model).is_some()) +} + pub(super) fn string_headers( extra_headers: Option>, ) -> Result, Error> { @@ -84,7 +90,7 @@ mod tests { use rstest::rstest; - use super::{messages_provider, string_headers}; + use super::{messages_provider, string_headers, supports_messages_provider}; use crate::Error; use litellm_core_utils::get_llm_provider_logic::LlmProviders; @@ -113,6 +119,19 @@ mod tests { assert!(messages_provider(provider, model).is_none()); } + #[rstest] + #[case::anthropic("anthropic/claude-test", None, true)] + #[case::inferred_bare_model("claude-test", Some("anthropic"), true)] + #[case::generic_passthrough("openai/gpt-5", None, false)] + #[case::vertex_non_claude("vertex_ai/gemini", None, false)] + fn admission_requires_a_native_adapter( + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] supported: bool, + ) { + assert_eq!(supports_messages_provider(model, provider), supported); + } + #[test] fn string_headers_rejects_non_string_values() { let headers = json!({"x-count": 3}).as_object().unwrap().clone(); diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index e0a7aa5adb8..ccb4c09bd4c 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -1,4 +1,4 @@ -use std::time::Duration; +use std::{collections::BTreeMap, time::Duration}; use bytes::Bytes; use futures_util::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream}; @@ -77,7 +77,10 @@ pub(super) async fn execute( &request.wire, ); let identity = request.identity.clone(); - let (output, source) = match cache.lookup().await { + let hit = cache.lookup().await.filter(|(output, _)| { + matches!(output, MessagesCallResponse::Stream { .. }) == request.stream + }); + let (output, source) = match hit { Some(hit) => hit, None => ( call_provider(&route.http, request, context).await?, @@ -412,18 +415,63 @@ fn decode_response( model: &str, text: &str, ) -> Result { - let response = serde_json::from_str(text).map_err(|err| { - Error::invalid_response(litellm_llms::ErrorDetail::invalid( - "messages response JSON", - err, - )) - })?; + validate_response_numbers(text)?; + let response = serde_json::from_str(text).map_err(invalid_response_json)?; config .transform_anthropic_messages_response(model, response) .map_err(Error::from) .map_err(Error::after_send) } +fn invalid_response_json(error: serde_json::Error) -> Error { + Error::invalid_response(litellm_llms::ErrorDetail::invalid( + "messages response JSON", + error, + )) +} + +fn validate_response_numbers(text: &str) -> Result<(), Error> { + use serde_json::value::RawValue; + + let root: &RawValue = serde_json::from_str(text).map_err(invalid_response_json)?; + let mut pending = vec![(root, 0usize)]; + while let Some((value, depth)) = pending.pop() { + let token = value.get().trim(); + if depth >= 128 { + return Err(invalid_response_json(serde::de::Error::custom( + "messages response nesting exceeds 128 levels", + ))); + } + match token.as_bytes().first() { + Some(b'{') => { + let fields: BTreeMap = + serde_json::from_str(token).map_err(invalid_response_json)?; + pending.extend(fields.into_values().map(|value| (value, depth + 1))); + } + Some(b'[') => { + let items: Vec<&RawValue> = + serde_json::from_str(token).map_err(invalid_response_json)?; + pending.extend(items.into_iter().map(|value| (value, depth + 1))); + } + Some(b'-' | b'0'..=b'9') if !token.contains(['.', 'e', 'E']) => { + let integer = if token.starts_with('-') { + token.parse::().map(|_| ()) + } else { + token.parse::().map(|_| ()) + }; + integer.map_err(|error| { + Error::invalid_response(litellm_llms::ErrorDetail::invalid( + "messages response integer outside i64/u64 range", + error, + )) + })?; + } + _ => {} + } + } + Ok(()) +} + fn streaming_response( response: reqwest::Response, decoder: Option, @@ -495,6 +543,46 @@ mod tests { use crate::{MessagesSettings, MessagesShaping, common_utils::messages_provider}; use litellm_inference::provider::resolve_llm_provider; + #[rstest] + #[case::u64_max("18446744073709551615", Some(json!(u64::MAX)))] + #[case::above_u64("18446744073709551617", None)] + #[case::i64_min("-9223372036854775808", Some(json!(i64::MIN)))] + #[case::below_i64("-9223372036854775809", None)] + #[case::decimal("18446744073709551617.0", Some(json!(18446744073709551617.0_f64)))] + #[case::exponent("1.8446744073709552e19", Some(json!(1.8446744073709552e19_f64)))] + #[case::quoted_integer(r#""18446744073709551617""#, Some(json!("18446744073709551617")))] + #[case::escaped_digit(r#""\u00318446744073709551617""#, Some(json!("18446744073709551617")))] + fn response_integer_tokens_never_round_into_success( + #[case] token: &str, + #[case] expected: Option, + ) { + let text = format!( + r#"{{"id":"m","type":"message","role":"assistant","model":"claude-test","content":[],"future":{{"values":[{token}]}}}}"#, + ); + let config = litellm_llms::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; + let result = decode_response(&config, "claude-test", &text); + match expected { + Some(expected) => assert_eq!(result.unwrap().extra["future"]["values"][0], expected), + None => { + let error = result.unwrap_err(); + assert!(error.is_response()); + assert!(!error.is_request()); + assert!(error.dispatched()); + } + } + } + + #[rstest] + #[case::long_numeric_string(format!(r#""{}""#, "9".repeat(4096)), true)] + #[case::long_integer("9".repeat(4096), false)] + #[case::deeply_nested(format!("{}0{}", "[".repeat(512), "]".repeat(512)), false)] + fn response_numeric_guard_bounds_nesting_and_ignores_string_contents( + #[case] token: String, + #[case] accepted: bool, + ) { + assert_eq!(validate_response_numbers(&token).is_ok(), accepted); + } + #[fixture] fn shaping() -> MessagesShaping { MessagesShaping::default() diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index 745442d0314..00271af55f1 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -14,7 +14,7 @@ use std::sync::Arc; use crate::common_utils::messages_provider; -pub use common_utils::merge_messages_header_sources; +pub use common_utils::{merge_messages_header_sources, supports_messages_provider}; pub use litellm_inference::Error; pub use types::{ MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 663c87453a9..877724fee62 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -293,3 +293,285 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_ assert!(matches!(cached.source, ResultSource::Cache { .. })); assert_eq!(cached.cost, Cost::CacheHit); } + +#[fixture] +fn streaming_call() -> litellm_inference_messages::MessagesCall { + litellm_inference_messages::MessagesCall { + body: serde_json::from_value(json!({ + "model": "anthropic/cache-test-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 32, + "stream": true + })) + .unwrap(), + api_key: Some("stream-key".into()), + api_base: None, + custom_llm_provider: None, + litellm_params: Default::default(), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(5)), + shaping: Default::default(), + } +} + +const COMPLETE_STREAM: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + +async fn streamed_bytes(output: litellm_inference_messages::MessagesCallResponse) -> Vec { + use futures_util::TryStreamExt; + let litellm_inference_messages::MessagesCallResponse::Stream { chunks, .. } = output else { + panic!("expected a Messages stream"); + }; + chunks.try_collect::>().await.unwrap().concat() +} + +#[rstest] +#[case::complete(COMPLETE_STREAM.to_owned(), true)] +#[case::data_only("data: {\"type\":\"message_stop\"}\n\n".into(), true)] +#[case::crlf("event: message_stop\r\ndata: {\"type\":\"message_stop\"}\r\n\r\n".into(), true)] +#[case::empty(String::new(), false)] +#[case::missing_stop("data: {\"type\":\"message_start\"}\n\n".into(), false)] +#[case::unterminated_stop("data: {\"type\":\"message_stop\"}".into(), false)] +#[case::partial_tail(format!("{COMPLETE_STREAM}data: {{"), false)] +#[case::error_before_stop(format!("data: {{\"type\":\"error\",\"error\":{{\"message\":\"failed\"}}}}\n\n{COMPLETE_STREAM}"), false)] +#[case::error_after_stop(format!("{COMPLETE_STREAM}data: {{\"type\":\"error\"}}\n\n"), false)] +#[case::invalid_json(format!("data: invalid\n\n{COMPLETE_STREAM}"), false)] +#[case::event_name_without_type("event: message_stop\ndata: {}\n\n".into(), false)] +#[case::over_entry_limit(format!("data: {{\"type\":\"content_block_delta\",\"text\":\"{}\"}}\n\n{COMPLETE_STREAM}", "x".repeat(4096)), false)] +#[tokio::test] +async fn stream_cache_eligibility_preserves_every_delivered_byte( + cache: Arc, + streaming_call: litellm_inference_messages::MessagesCall, + #[case] wire: String, + #[case] cache_hit: bool, +) { + use wiremock::ResponseTemplate; + let response = ResponseTemplate::new(200).set_body_raw(wire.clone(), "text/event-stream"); + let upstream = support::upstream([response.clone(), response]).await; + let route = support::messages_route(litellm_inference_testing::no_secrets()) + .with_cache(ScopedCache::new(cache, CacheScope::Shared)); + let request = || litellm_inference_messages::MessagesCall { + body: streaming_call.body.clone(), + api_base: Some(upstream.uri()), + ..self::streaming_call() + }; + let hooks = ChangingHooks::default(); + + let first = streamed_bytes(route.execute(request(), &hooks, None).await.unwrap()).await; + let second = streamed_bytes(route.execute(request(), &hooks, None).await.unwrap()).await; + + assert_eq!(first, wire.as_bytes()); + assert_eq!(second, wire.as_bytes()); + assert_eq!( + support::received(&upstream).await.len(), + if cache_hit { 1 } else { 2 } + ); + let facts = hooks.facts.lock().unwrap(); + let [first, second] = facts.as_slice() else { + panic!("each call reports exactly one result"); + }; + assert_eq!(first.source, ResultSource::Provider); + assert_eq!( + matches!(second.source, ResultSource::Cache { .. }), + cache_hit + ); + assert_eq!( + second.cost, + if cache_hit { + Cost::CacheHit + } else { + Cost::Deferred + } + ); +} + +#[rstest] +#[case::before_first_chunk(false)] +#[case::after_terminal_chunk(true)] +#[tokio::test] +async fn dropping_a_stream_before_eof_cancels_once_and_does_not_cache( + cache: Arc, + streaming_call: litellm_inference_messages::MessagesCall, + #[case] read_chunk: bool, +) { + use futures_util::StreamExt; + use litellm_host::lifecycle::{CallEvent, ExecutionEvent}; + use litellm_inference_messages::MessagesCallResponse; + use wiremock::ResponseTemplate; + + let response = ResponseTemplate::new(200).set_body_raw(COMPLETE_STREAM, "text/event-stream"); + let upstream = support::upstream([response.clone(), response]).await; + let route = support::messages_route(litellm_inference_testing::no_secrets()) + .with_cache(ScopedCache::new(cache, CacheScope::Shared)); + let request = || litellm_inference_messages::MessagesCall { + body: streaming_call.body.clone(), + api_base: Some(upstream.uri()), + ..self::streaming_call() + }; + let hooks = ChangingHooks::default(); + let observations = support::Observations::default(); + let first = route + .execute(request(), &hooks, Some(observations.sender.clone())) + .await + .unwrap(); + let MessagesCallResponse::Stream { mut chunks, .. } = first else { + panic!("expected a Messages stream"); + }; + if read_chunk { + let mut delivered = Vec::new(); + while delivered.len() < COMPLETE_STREAM.len() { + delivered.extend_from_slice(&chunks.next().await.unwrap().unwrap()); + } + assert_eq!(delivered, COMPLETE_STREAM.as_bytes()); + } + assert!(matches!( + observations.lock().unwrap().as_slice(), + [ + CallEvent::Started { .. }, + CallEvent::Execution(ExecutionEvent::ResultReady { .. }) + ] + )); + drop(chunks); + assert!(matches!( + observations.lock().unwrap().as_slice(), + [ + CallEvent::Started { .. }, + CallEvent::Execution(ExecutionEvent::ResultReady { .. }), + CallEvent::Cancelled { .. } + ] + )); + + let second = streamed_bytes(route.execute(request(), &hooks, None).await.unwrap()).await; + let third = streamed_bytes(route.execute(request(), &hooks, None).await.unwrap()).await; + + assert_eq!(second, COMPLETE_STREAM.as_bytes()); + assert_eq!(third, COMPLETE_STREAM.as_bytes()); + assert_eq!(support::received(&upstream).await.len(), 2); + let facts = hooks.facts.lock().unwrap(); + let [cancelled, completed, cached] = facts.as_slice() else { + panic!("each call reports exactly one result"); + }; + assert_eq!(cancelled.source, ResultSource::Provider); + assert_eq!(completed.source, ResultSource::Provider); + assert!(matches!(cached.source, ResultSource::Cache { .. })); + assert_eq!(cached.cost, Cost::CacheHit); +} + +#[rstest] +#[case::platform("bedrock/claude_platform/claude-test")] +#[case::mantle("bedrock_mantle/claude-test")] +#[case::mantle_alias("bedrock/mantle/claude-test")] +#[tokio::test] +async fn signed_messages_bypass_the_response_cache( + cache: Arc, + #[case] model: &str, +) { + use litellm_inference_testing::RecordingSecrets; + use support::ReceivedRequest; + + let response = support::json_response(json!({ + "id": "signed-message", "type": "message", "role": "assistant", "model": "claude-test", + "content": [{"type": "text", "text": "signed reply"}], "stop_reason": "end_turn" + })); + let upstream = support::upstream([response.clone(), response]).await; + let secrets = Arc::new(RecordingSecrets::new([ + ("AWS_ACCESS_KEY_ID", "test-access"), + ("AWS_SECRET_ACCESS_KEY", "test-secret"), + ("AWS_SESSION_TOKEN", "test-session"), + ("AWS_REGION_NAME", "us-west-2"), + ("ANTHROPIC_AWS_WORKSPACE_ID", "test-workspace"), + ])); + let route = + support::messages_route(secrets).with_cache(ScopedCache::new(cache, CacheScope::Shared)); + let request = || litellm_inference_messages::MessagesCall { + body: serde_json::from_value(json!({ + "model": model, "messages": [{"role": "user", "content": "hello"}], "max_tokens": 32 + })) + .unwrap(), + api_key: None, + api_base: Some(upstream.uri()), + ..streaming_call() + }; + let hooks = ChangingHooks::default(); + + let first = route.execute(request(), &hooks, None).await.unwrap(); + let second = route.execute(request(), &hooks, None).await.unwrap(); + + let ( + litellm_inference_messages::MessagesCallResponse::Complete(first), + litellm_inference_messages::MessagesCallResponse::Complete(second), + ) = (first, second) + else { + panic!("expected completed signed messages"); + }; + assert_eq!(first, second); + let requests = support::received(&upstream).await; + let [first, second] = requests.as_slice() else { + panic!("signed requests must each reach the provider"); + }; + assert!( + first + .header("authorization") + .unwrap() + .starts_with("AWS4-HMAC-SHA256 Credential=test-access/") + ); + assert!( + second + .header("authorization") + .unwrap() + .starts_with("AWS4-HMAC-SHA256 Credential=test-access/") + ); + let facts = hooks.facts.lock().unwrap(); + let [first, second] = facts.as_slice() else { + panic!("each signed request reports one result"); + }; + assert_eq!(first.source, ResultSource::Provider); + assert_eq!(second.source, ResultSource::Provider); +} + +#[rstest] +#[tokio::test] +async fn oversized_response_integers_are_rejected_before_cache_storage( + cache: Arc, +) { + use litellm_inference_messages::{MessagesCall, MessagesCallResponse}; + use wiremock::ResponseTemplate; + + let invalid = r#"{"id":"bad","type":"message","role":"assistant","model":"cache-test","content":[],"future_extension":{"wide_integer":18446744073709551617}}"#; + let valid = json!({"id":"good","type":"message","role":"assistant","model":"cache-test","content":[],"future_extension":{"wide_integer":u64::MAX}}); + let upstream = support::upstream([ + ResponseTemplate::new(200).set_body_raw(invalid, "application/json"), + support::json_response(valid.clone()), + ]) + .await; + let route = support::messages_route(litellm_inference_testing::no_secrets()) + .with_cache(ScopedCache::new(cache, CacheScope::Shared)); + let request = || { + MessagesCall { + body: serde_json::from_value(json!({"model":"anthropic/cache-test","messages":[{"role":"user","content":"hi"}],"max_tokens":16})).unwrap(), + api_key: Some("synthetic".into()), api_base: Some(upstream.uri()), custom_llm_provider: None, + litellm_params: Default::default(), extra_headers: None, provider_specific_header: None, + timeout: None, shaping: Default::default(), + } + }; + let rejected = route + .execute(request(), &(), None) + .await + .err() + .expect("oversized provider integer must fail"); + assert!(rejected.is_response()); + assert!(!rejected.is_request()); + let complete = route.execute(request(), &(), None).await.unwrap(); + let cached = route.execute(request(), &(), None).await.unwrap(); + let (MessagesCallResponse::Complete(complete), MessagesCallResponse::Complete(cached)) = + (complete, cached) + else { + panic!("expected complete messages"); + }; + assert_eq!( + complete.extra["future_extension"], + valid["future_extension"] + ); + assert_eq!(cached, complete); + assert_eq!(support::received(&upstream).await.len(), 2); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index c3d06027c32..6ac1539269c 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -450,3 +450,170 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); assert_eq!(optional_params, json!({"max_tokens": 16})); } + +#[derive(Default)] +struct GatedResult { + entered: tokio::sync::Notify, + release: tokio::sync::Notify, +} + +impl litellm_host::interceptors::Interceptors for GatedResult { + async fn before_provider_request( + &self, + wire: WireRequest, + _: RequestContext, + ) -> Result { + Ok(wire) + } + + async fn after_provider_response( + &self, + _: litellm_host::interceptors::RawResponse, + ) -> Result<(), Error> { + Ok(()) + } + + async fn result_ready(&self, _: ExecutionFacts) -> Result<(), Error> { + self.entered.notify_one(); + self.release.notified().await; + Ok(()) + } +} + +#[rstest] +#[case::complete_response(false, false)] +#[case::complete_stream(true, false)] +#[case::cancel_response(false, true)] +#[case::cancel_stream(true, true)] +#[tokio::test] +async fn pending_result_acceptance_blocks_delivery_and_cache_storage( + call: MessagesCall, + #[case] streaming: bool, + #[case] cancel: bool, + #[values(false, true)] cached: bool, +) { + use futures_util::TryStreamExt; + use litellm_cache_memory::InMemoryCache; + use litellm_cache_response::{CachePolicy, CacheScope, ResponseCache, ScopedCache}; + + let response = if streaming { + ResponseTemplate::new(200).set_body_raw( + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + "text/event-stream", + ) + } else { + message_response() + }; + let upstream = upstream([response.clone(), response.clone(), response]).await; + let route = messages_route(no_secrets()).with_cache(ScopedCache::new( + Arc::new(ResponseCache::new(Arc::new(InMemoryCache::default()))), + CacheScope::Shared, + )); + let request = || { + authenticated( + with_fields( + MessagesCall { + body: call.body.clone(), + ..super::call() + }, + json!({"stream": streaming}), + ), + upstream.uri(), + ) + }; + if cached { + let output = route.execute(request(), &(), None).await.unwrap(); + if let MessagesCallResponse::Stream { chunks, .. } = output { + chunks.try_collect::>().await.unwrap(); + } + } + let gate = GatedResult::default(); + let observations = super::support::Observations::default(); + let mut pending = Box::pin(route.execute(request(), &gate, Some(observations.sender.clone()))); + tokio::time::timeout(Duration::from_secs(5), async { + tokio::select! { + _ = gate.entered.notified() => {}, + _ = &mut pending => panic!("a result cannot be delivered before acceptance"), + } + }) + .await + .unwrap(); + assert!(futures_util::poll!(&mut pending).is_pending()); + + let probe = RecordingHost::passthrough(request()); + let output = route + .execute( + request(), + &probe, + litellm_inference::CallOptions { + cache: Some(CachePolicy { + no_store: true, + ..Default::default() + }), + ..Default::default() + }, + ) + .await + .unwrap(); + if let MessagesCallResponse::Stream { chunks, .. } = output { + chunks.try_collect::>().await.unwrap(); + } + assert_eq!( + matches!( + probe.facts.lock().unwrap()[0].source, + ResultSource::Cache { .. } + ), + cached + ); + assert_eq!(received(&upstream).await.len(), if cached { 1 } else { 2 }); + + if cancel { + drop(pending); + let events = observations.lock().unwrap(); + assert!(matches!(events.last(), Some(CallEvent::Cancelled { .. }))); + assert_eq!( + events + .iter() + .filter(|event| matches!(event, CallEvent::Cancelled { .. })) + .count(), + 1 + ); + assert!(!events.iter().any(|event| matches!( + event, + CallEvent::Succeeded { .. } | CallEvent::Failed { .. } + ))); + } else { + gate.release.notify_one(); + let output = pending.await.unwrap(); + if let MessagesCallResponse::Stream { chunks, .. } = output { + chunks.try_collect::>().await.unwrap(); + } + assert!(matches!( + observations.lock().unwrap().last(), + Some(CallEvent::Succeeded { .. }) + )); + } + + let retry = RecordingHost::passthrough(request()); + let output = route.execute(request(), &retry, None).await.unwrap(); + if let MessagesCallResponse::Stream { chunks, .. } = output { + chunks.try_collect::>().await.unwrap(); + } + assert_eq!( + matches!( + retry.facts.lock().unwrap()[0].source, + ResultSource::Cache { .. } + ), + cached || !cancel + ); + assert_eq!( + received(&upstream).await.len(), + if cached { + 1 + } else if cancel { + 3 + } else { + 2 + } + ); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/stream.rs b/litellm-rust/crates/inference-messages/tests/messages/stream.rs index 99ba08495aa..d7c978dce5e 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/stream.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/stream.rs @@ -477,16 +477,36 @@ async fn the_timeout_covers_a_stalled_stream_error_body(call: MessagesCall) { #[rstest] #[tokio::test] -async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesCall) { +async fn the_sdk_yields_a_body_error_once_after_delivered_chunks( + call: MessagesCall, + #[values(false, true)] cached: bool, +) { + use litellm_cache_memory::InMemoryCache; + use litellm_cache_response::{CacheScope, ResponseCache, ScopedCache}; + use litellm_host::lifecycle::{CallEvent, ExecutionEvent, FailureOrigin}; + let (base, connection) = stalling_upstream().await; - let response = messages_route(no_secrets()) + let observations = support::Observations::default(); + let route = messages_route(no_secrets()); + let route = if cached { + route.with_cache(ScopedCache::new( + Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new( + Some(100), + Some(Duration::from_secs(60)), + )))), + CacheScope::Shared, + )) + } else { + route + }; + let response = route .execute( MessagesCall { timeout: Some(Duration::from_millis(300)), ..streaming(call, base) }, &(), - None, + Some(observations.sender.clone()), ) .await .unwrap(); @@ -507,6 +527,18 @@ async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesC assert!(!error.is_request(), "{error:?}"); assert_eq!(error.upstream_response(), None); assert!(chunks.next().await.is_none()); + drop(chunks); + assert!(matches!( + observations.lock().unwrap().as_slice(), + [ + CallEvent::Started { .. }, + CallEvent::Execution(ExecutionEvent::ResultReady { .. }), + CallEvent::Failed { + origin: FailureOrigin::Call, + .. + } + ] + )); tokio::time::timeout(Duration::from_secs(5), connection) .await .expect("the failed stream closes its upstream connection") diff --git a/litellm-rust/crates/llms-types/src/formats/messages/request.rs b/litellm-rust/crates/llms-types/src/formats/messages/request.rs index d30d21e74e7..3740d499f48 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/request.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/request.rs @@ -67,6 +67,7 @@ impl From for String { #[macro_rules_attribute::apply(crate::wire_type)] #[derive(Default)] +#[serde(remote = "Self")] pub struct ContentBlock { #[serde(rename = "type", default, skip_serializing_if = "Option::is_none")] pub block_type: Option, @@ -96,6 +97,33 @@ pub struct ContentBlock { pub extra: Map, } +impl serde::Serialize for ContentBlock { + fn serialize(&self, serializer: S) -> Result { + Self::serialize(self, serializer) + } +} + +impl<'de> serde::Deserialize<'de> for ContentBlock { + fn deserialize>(deserializer: D) -> Result { + let value = Value::deserialize(deserializer)?; + let nulls: Map = value + .as_object() + .into_iter() + .flatten() + .filter(|(_, value)| value.is_null()) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + let block = Self::deserialize(value).map_err(serde::de::Error::custom)?; + if !matches!(block.block_type, Some(ContentBlockType::Other(_))) { + return Ok(block); + } + Ok(Self { + extra: block.extra.into_iter().chain(nulls).collect(), + ..block + }) + } +} + impl ContentBlock { pub fn text(text: impl Into) -> Self { Self { @@ -430,6 +458,12 @@ mod tests { }))] #[case::provider_specific_fields(json!({"type": "tool_use", "id": "t", "name": "f", "input": {}, "provider_specific_fields": {"x": 1}}))] #[case::untyped(json!({"unknown": {"nested": true}}))] + #[case::unknown_block_with_familiar_nulls(json!({ + "type": "future_block", "text": null, "thinking": null, "signature": null, + "data": null, "id": null, "name": null, "tool_use_id": null, + "input": null, "content": null, "provider_specific_fields": null, + "cache_control": null, "source": {"data": "123"}, "future": null + }))] fn content_block_round_trips_unchanged(#[case] block: Value) { assert_eq!(round_trip::(&block), block); } @@ -852,4 +886,78 @@ mod tests { assert_eq!(block.extra.get("source"), wire.get("source")); assert_eq!(serde_json::to_value(block).unwrap(), wire); } + + #[rstest] + #[case::zero(0, 0)] + #[case::above_float_precision(9_007_199_254_740_993, 9_007_199_254_740_993)] + #[case::integer_limits(u64::MAX, i64::MAX)] + #[case::negative_top_k(1, i64::MIN)] + fn request_integer_fields_preserve_exact_wire_values( + #[case] max_tokens: u64, + #[case] top_k: i64, + ) { + let wire = format!( + r#"{{"model":"test","messages":[],"max_tokens":{max_tokens},"top_k":{top_k}}}"# + ); + let parsed: MessagesRequest = serde_json::from_str(&wire).unwrap(); + assert_eq!(parsed.params.max_tokens, Some(max_tokens)); + assert_eq!(parsed.params.top_k, Some(top_k)); + assert_eq!( + serde_json::to_value(parsed).unwrap(), + serde_json::from_str::(&wire).unwrap() + ); + } + + #[rstest] + #[case::negative_max_tokens(r#""max_tokens":-1"#)] + #[case::fractional_max_tokens(r#""max_tokens":1.5"#)] + #[case::overflow_max_tokens(r#""max_tokens":18446744073709551616"#)] + #[case::string_max_tokens(r#""max_tokens":"16""#)] + #[case::boolean_max_tokens(r#""max_tokens":true"#)] + #[case::overflow_top_k(r#""top_k":9223372036854775808"#)] + #[case::underflow_top_k(r#""top_k":-9223372036854775809"#)] + #[case::fractional_top_k(r#""top_k":1.5"#)] + #[case::string_stream(r#""stream":"true""#)] + #[case::numeric_stream(r#""stream":1"#)] + #[case::non_string_stop_sequence(r#""stop_sequences":[null]"#)] + fn request_rejects_invalid_scalar_fields(#[case] field: &str) { + let wire = format!(r#"{{"model":"test","messages":[],{field}}}"#); + assert!(serde_json::from_str::(&wire).is_err()); + } + + #[rstest] + #[case::missing(json!({}), json!({}))] + #[case::null( + json!({"max_tokens":null,"system":null,"tools":null,"stream":null,"metadata":null,"future":null}), + json!({"future":null}) + )] + #[case::empty( + json!({"max_tokens":0,"system":"","tools":[],"stream":false,"metadata":{},"future":null}), + json!({"max_tokens":0,"system":"","tools":[],"stream":false,"metadata":{},"future":null}) + )] + fn request_distinguishes_empty_values_and_unknown_nulls_from_absent_options( + #[case] fields: Value, + #[case] expected: Value, + ) { + let required = json!({"model":"test","messages":[]}); + let wire: Map = required + .as_object() + .unwrap() + .clone() + .into_iter() + .chain(fields.as_object().unwrap().clone()) + .collect(); + let expected: Map = required + .as_object() + .unwrap() + .clone() + .into_iter() + .chain(expected.as_object().unwrap().clone()) + .collect(); + let parsed: MessagesRequest = serde_json::from_value(Value::Object(wire)).unwrap(); + assert_eq!( + serde_json::to_value(parsed).unwrap(), + Value::Object(expected) + ); + } } diff --git a/litellm-rust/crates/llms-types/src/formats/messages/streaming.rs b/litellm-rust/crates/llms-types/src/formats/messages/streaming.rs index b2f70ee3753..f5ecd68d7d7 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/streaming.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/streaming.rs @@ -169,6 +169,9 @@ mod tests { "type": "message_delta", "delta": {"stop_reason": "end_turn", "future_delta": 42}, "usage": {"output_tokens": 7, "future_usage": true} }))] + #[case::large_index(json!({"type":"content_block_stop","index":u64::MAX}))] + #[case::large_usage(json!({"type":"message_stop","usage":{"input_tokens":u64::MAX,"output_tokens":9_007_199_254_740_993_u64,"future_usage":null}}))] + #[case::partial_json(json!({"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"x\":"}}))] #[case::error(json!({ "type": "error", "error": {"type": "future_error", "message": "failed", "future": 42} }))] @@ -197,6 +200,14 @@ mod tests { #[rstest] #[case::invalid_index(json!({"type": "content_block_stop", "index": "zero"}))] #[case::missing_delta(json!({"type": "content_block_delta", "index": 0}))] + #[case::negative_index(json!({"type":"content_block_stop","index":-1}))] + #[case::fractional_index(json!({"type":"content_block_stop","index":0.5}))] + #[case::null_index(json!({"type":"content_block_stop","index":null}))] + #[case::unknown_delta(json!({"type":"content_block_delta","index":0,"delta":{"type":"future_delta","text":"x"}}))] + #[case::tool_fragment_not_text(json!({"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":{"x":1}}}))] + #[case::negative_usage(json!({"type":"message_stop","usage":{"output_tokens":-1}}))] + #[case::fractional_usage(json!({"type":"message_delta","delta":{},"usage":{"input_tokens":0.5}}))] + #[case::error_without_message(json!({"type":"error","error":{"type":"api_error"}}))] #[case::unknown_event(json!({"type": "future_event"}))] fn malformed_or_unrecognized_events_remain_rejected(#[case] wire: Value) { assert!(serde_json::from_value::(wire).is_err()); diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 79ee17e8b97..d04eb9c6405 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -37,6 +37,7 @@ thiserror.workspace = true time.workspace = true tokio = { workspace = true, features = ["sync"] } url.workspace = true +unicode-general-category = "1.1" [dev-dependencies] litellm-http = { workspace = true, features = ["test-support"] } diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 13653f438f7..eabc9075087 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1,4 +1,5 @@ use crate::base_llm::messages::context::MessagesModelCapabilities; +use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_auth::CredentialPlacement; use litellm_secret_value::SecretValue; @@ -317,7 +318,12 @@ fn retain_blocks(messages: Vec, keep: impl Fn(&ContentBlock) -> bool) - pub fn strip_empty_content_blocks(messages: Vec) -> Vec { retain_blocks(messages, |block| { - !is_empty_text_block(block) && !is_empty_thinking_block(block) + !is_empty_text_block(block) + && (!is_empty_thinking_block(block) + || block + .signature + .as_ref() + .is_some_and(|signature| !signature.is_empty())) }) } @@ -397,6 +403,11 @@ pub fn strip_provider_specific_fields(messages: Vec) -> Vec { .into_iter() .map(|block| ContentBlock { provider_specific_fields: None, + extra: block + .extra + .into_iter() + .filter(|(key, _)| key != "provider_specific_fields") + .collect(), ..block }) .collect(), @@ -444,8 +455,45 @@ pub fn filter_billing_headers_from_system(system: SystemPrompt) -> Option Option { + let decoded = STANDARD.decode(data).ok()?; + let native = std::str::from_utf8(&decoded).ok()?; + if !native.is_ascii() || STANDARD.decode(native).is_err() { + return None; + } + Some(native.to_string()) +} + +fn with_native_redacted_data(block: ContentBlock) -> ContentBlock { + if !block.is_type(ContentBlockType::RedactedThinking) { + return block; + } + let Some(data) = block + .data + .as_deref() + .and_then(unwrapped_redacted_thinking_data) + else { + return block; + }; + ContentBlock { + data: Some(data), + ..block + } +} + pub fn strip_encrypted_reasoning_blocks(messages: Vec) -> Vec { retain_blocks(messages, |block| !is_encrypted_reasoning_block(block)) + .into_iter() + .map(|message| match message.content { + MessageContent::Blocks(blocks) => Message { + content: MessageContent::Blocks( + blocks.into_iter().map(with_native_redacted_data).collect(), + ), + ..message + }, + MessageContent::Text(_) => message, + }) + .collect() } fn is_advisor_use(block: &ContentBlock) -> bool { @@ -858,10 +906,18 @@ mod tests { json!([ {"role": "user", "content": "weather?"}, {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "", "signature": "sig_abc"}, {"type": "tool_use", "id": "toolu_01A", "name": "get_weather", "input": {"city": "Paris"}} ]} ]) )] + #[case::empty_thinking_with_empty_signature_is_removed( + json!([{"role": "assistant", "content": [ + {"type": "thinking", "thinking": "", "signature": ""}, + {"type": "text", "text": "answer"} + ]}]), + json!([{"role": "assistant", "content": [{"type": "text", "text": "answer"}]}]) + )] #[case::whitespace_thinking_beside_real_and_redacted_thinking( json!([{"role": "assistant", "content": [ {"type": "thinking", "thinking": " \n "}, @@ -896,6 +952,9 @@ mod tests { #[case::non_empty_text(json!([{"role": "assistant", "content": [{"type": "text", "text": "hi"}]}]))] #[case::padded_text(json!([{"role": "assistant", "content": [{"type": "text", "text": " hi "}]}]))] #[case::empty_string_content(json!([{"role": "user", "content": ""}]))] + #[case::signed_whitespace_thinking(json!([{"role": "assistant", "content": [ + {"type": "thinking", "thinking": " \n ", "signature": "sig_abc"} + ]}]))] #[case::textless_non_text_block(json!([{"role": "user", "content": [ {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AA=="}} ]}]))] @@ -1001,6 +1060,10 @@ mod tests { {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "ok"}]} ]) )] + #[case::unknown_block_null_provider_fields( + json!([{"role": "user", "content": [{"type": "future", "provider_specific_fields": null, "text": null}]}]), + json!([{"role": "user", "content": [{"type": "future", "text": null}]}]) + )] fn strip_provider_specific_fields_rewrites(#[case] input: Value, #[case] expected: Value) { assert_eq!(apply(strip_provider_specific_fields, input), expected); } @@ -1075,6 +1138,33 @@ mod tests { assert_eq!(apply(strip_encrypted_reasoning_blocks, input), expected); } + #[rstest] + #[case::native(false)] + #[case::converse_wrapped(true)] + fn redacted_thinking_history_uses_the_native_data(#[case] wrapped: bool) { + use base64::{Engine, engine::general_purpose::STANDARD}; + + let native = "EmwKAhgBEgy3va3pzix/LafPsn4aDFIT2Xlxh0L5L8rLVyIwxtE3rAFBa8cr3qpP"; + let data = if wrapped { + STANDARD.encode(native) + } else { + native.to_string() + }; + let input = json!([{"role": "assistant", "content": [ + {"type": "redacted_thinking", "data": data}, + {"type": "text", "text": "answer"} + ]}]); + assert_eq!( + apply(strip_encrypted_reasoning_blocks, input), + json!([ + {"role": "assistant", "content": [ + {"type": "redacted_thinking", "data": native}, + {"type": "text", "text": "answer"} + ]} + ]) + ); + } + #[rstest] #[case::anthropic_signed_blocks(json!([ {"role": "user", "content": "Solve it."}, diff --git a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs index 54f924ddec6..eb0146f41f2 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs @@ -184,6 +184,7 @@ mod tests { json!([ {"role": "user", "content": "run it"}, {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "", "signature": "sig"}, {"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}, {"type": "server_tool_use", "id": "srv_1", "name": "web_search", "input": {"query": "q"}}, {"type": "text", "text": "Web search results:\n\nURL: u"} diff --git a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs index ae3488a21f9..7bde055f3bb 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs @@ -11,6 +11,7 @@ use litellm_llms_types::{ use litellm_python_compat::{json::from_json, repr::repr, truthy::truthy}; use serde_json::Value; use strum::VariantArray; +use unicode_general_category::{GeneralCategory, get_general_category}; use crate::base_llm::messages::context::{ MessagesModelCapabilities, ThinkingBudgets, ThinkingContext, @@ -248,9 +249,34 @@ fn numeric_budget_tokens(budget: &Recognized) -> Option { } } +fn decimal_digit(character: char) -> Option { + if character.is_ascii_digit() { + return Some(character); + } + let decimal_run_length = + std::iter::successors(Some(u32::from(character)), |code| code.checked_sub(1)) + .take_while(|code| { + char::from_u32(*code).is_some_and(|character| { + get_general_category(character) == GeneralCategory::DecimalNumber + }) + }) + .count(); + decimal_run_length + .checked_sub(1) + .and_then(|offset| char::from_digit((offset % 10) as u32, 10)) +} + /// Python's `int(str)`: surrounding whitespace, a sign, and single underscores between digits. fn python_int(text: &str) -> Option { - let trimmed = text.trim(); + let normalized: String = text + .trim() + .chars() + .map(|character| match character { + '+' | '-' | '_' => Some(character), + character => decimal_digit(character), + }) + .collect::>()?; + let trimmed = normalized.as_str(); let digits = trimmed.trim_start_matches(['+', '-']); let well_formed = trimmed.len() - digits.len() <= 1 && !digits.is_empty() @@ -1361,6 +1387,19 @@ mod tests { #[case::float_below_a_tier(json!(1023.9), "low")] #[case::string_budget(json!(" 16_000 "), "xhigh")] #[case::negative_budget(json!(-5), "low")] + #[case::fullwidth_budget(json!("8192"), "xhigh")] + #[case::arabic_indic_budget(json!("٨١٩٢"), "xhigh")] + #[case::mixed_unicode_budget(json!(" +8_١92 "), "xhigh")] + #[case::supplementary_decimal_budget(json!("𝟠𝟙𝟡𝟚"), "xhigh")] + #[case::negative_unicode_budget(json!("-8_١92"), "low")] + #[case::signed_underscored_budget(json!(" +8_192 "), "xhigh")] + #[case::negative_underscored_budget(json!("-8_192"), "low")] + #[case::unsigned_integer_limit(json!(u64::MAX), "xhigh")] + #[case::positive_string_overflow(json!("184467440737095516160"), "xhigh")] + #[case::negative_string_overflow(json!("-184467440737095516160"), "low")] + #[case::negative_fraction_truncates_toward_zero(json!(-0.9), "low")] + #[case::below_medium_fraction(json!(2047.9), "low")] + #[case::at_medium_fraction(json!(2048.0), "medium")] fn legacy_budgets_read_like_python_numbers(#[case] budget: Value, #[case] effort: &str) { let translated = translate_thinking( request( @@ -1380,6 +1419,16 @@ mod tests { #[rstest] #[case::text(json!("lots"), "str")] + #[case::superscript_digits(json!("⁸¹⁹²"), "str")] + #[case::roman_numeral(json!("Ⅷ"), "str")] + #[case::leading_separator(json!("_8192"), "str")] + #[case::trailing_separator(json!("8192_"), "str")] + #[case::repeated_separator(json!("8__192"), "str")] + #[case::two_signs(json!("--8192"), "str")] + #[case::sign_only(json!("+"), "str")] + #[case::decimal_string(json!("8192.0"), "str")] + #[case::exponent_string(json!("8e3"), "str")] + #[case::whitespace_only(json!(" "), "str")] #[case::list(json!([1]), "list")] #[case::null(Value::Null, "NoneType")] #[case::boolean(json!(true), "bool")] diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 18789e5b2be..472e93bdd24 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -388,8 +388,13 @@ fn compact_edit_from_openai(entry: &Map) -> Option { } let trigger = entry .get("compact_threshold") - .and_then(Value::as_f64) - .map(|threshold| json!({"type": "input_tokens", "value": threshold as i64})); + .and_then(|threshold| match threshold { + Value::Number(number) if number.is_i64() || number.is_u64() => Some(threshold.clone()), + Value::Number(number) => number.as_f64().map(|number| json!(number as i64)), + Value::Bool(value) => Some(json!(u8::from(*value))), + _ => None, + }) + .map(|threshold| json!({"type": "input_tokens", "value": threshold})); let passthrough = entry .iter() .filter(|(key, _)| !matches!(key.as_str(), "type" | "compact_threshold")) @@ -701,6 +706,26 @@ mod tests { json!([{"type": "compaction", "compact_threshold": 200000}]), Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 200000}}]})) )] + #[case::threshold_above_float_precision( + json!([{"type": "compaction", "compact_threshold": 9_007_199_254_740_993_u64}]), + Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 9_007_199_254_740_993_u64}}]})) + )] + #[case::threshold_unsigned_maximum( + json!([{"type": "compaction", "compact_threshold": u64::MAX}]), + Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": u64::MAX}}]})) + )] + #[case::threshold_signed_minimum( + json!([{"type": "compaction", "compact_threshold": i64::MIN}]), + Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": i64::MIN}}]})) + )] + #[case::true_threshold_matches_python_int( + json!([{"type": "compaction", "compact_threshold": true}]), + Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1}}]})) + )] + #[case::false_threshold_matches_python_int( + json!([{"type": "compaction", "compact_threshold": false}]), + Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 0}}]})) + )] #[case::other_keys_pass_through( json!([{"type": "compaction", "compact_threshold": 150000, "instructions": "Focus on preserving code snippets"}]), Some(json!({"edits": [{ diff --git a/litellm-rust/crates/llms/src/base_llm/messages/mid_conversation_system.rs b/litellm-rust/crates/llms/src/base_llm/messages/mid_conversation_system.rs index ec1a9863786..ce77bafe7bd 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/mid_conversation_system.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/mid_conversation_system.rs @@ -198,4 +198,70 @@ mod tests { json!({"type": "text", "text": "reminder"}) ); } + + #[rstest] + #[case::tool_result_first( + json!({"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":"ok"},{"type":"text","text":"next"}],"future":null}), + 1 + )] + #[case::tool_result_after_text( + json!({"role":"user","content":[{"type":"text","text":"first"},{"type":"tool_result","tool_use_id":"t","content":"ok"}]}), + 2 + )] + #[case::non_user_tool_result( + json!({"role":"assistant","content":[{"type":"tool_result","tool_use_id":"t","content":"ok"}]}), + 2 + )] + #[case::empty_blocks(json!({"role":"user","content":[]}), 2)] + fn only_a_user_turn_starting_with_tool_results_moves_before_the_system_run( + #[case] follower: Value, + #[case] follower_index: usize, + ) { + let first = json!({"role":"assistant","content":[{"type":"tool_use","id":"t","name":"lookup","input":{}}]}); + let converted = convert_mid_conversation_system_turns(messages(json!([ + first, + {"role":"system","content":[{"type":"future","payload":{"keep":null}}]}, + follower + ]))); + let wire = serde_json::to_value(&converted).unwrap(); + assert_eq!(wire[0], first); + assert_eq!(wire[follower_index], follower); + assert_eq!( + wire[3 - follower_index], + json!({"role":"user","content":[ + {"type":"text","text":CONVERTED_SYSTEM_NOTE}, + {"type":"future","payload":{"keep":null}} + ]}) + ); + assert_eq!( + convert_mid_conversation_system_turns(converted.clone()), + converted + ); + } + + #[rstest] + fn consecutive_system_runs_keep_their_order_across_multiple_tool_turns() { + let first = json!({"role":"assistant","content":"first"}); + let result = json!({"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":"ok"}]}); + let next = json!({"role":"assistant","content":"next"}); + let converted = convert_mid_conversation_system_turns(messages(json!([ + first, + {"role":"system","content":"one"}, + {"role":"system","content":"two"}, + result, + next, + {"role":"system","content":"tail"} + ]))); + assert_eq!( + serde_json::to_value(converted).unwrap(), + json!([ + first, + result, + {"role":"user","content":[{"type":"text","text":CONVERTED_SYSTEM_NOTE},{"type":"text","text":"one"}]}, + {"role":"user","content":[{"type":"text","text":CONVERTED_SYSTEM_NOTE},{"type":"text","text":"two"}]}, + next, + {"role":"user","content":[{"type":"text","text":CONVERTED_SYSTEM_NOTE},{"type":"text","text":"tail"}]} + ]) + ); + } } diff --git a/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs b/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs index 7e24da2b5eb..2952e0d532d 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs @@ -160,4 +160,93 @@ mod tests { normalized ); } + + #[rstest] + #[case::blocks( + json!([{"type":"text","text":"context","cache_control":{"type":"ephemeral","scope":"global","ttl":"1h","future":null}}]), + json!([{"type":"text","text":"context","cache_control":{"type":"ephemeral","ttl":"1h","future":null}}]) + )] + #[case::text(json!("context"), json!("context"))] + #[case::absent(Value::Null, Value::Null)] + fn scope_removal_preserves_cache_hints_and_opaque_nested_fields( + #[case] system: Value, + #[case] expected_system: Value, + ) { + let nested = json!([{"type":"text","text":"nested","cache_control":{"scope":"global"}}]); + let tool = json!({"name":"lookup","input_schema":{},"cache_control":{"scope":"global","ttl":"1h"}}); + let request: MessagesRequest = serde_json::from_value(json!({ + "model":"test", "messages":[ + {"role":"user","content":"plain","future_message":null}, + {"role":"assistant","content":[ + {"type":"future","cache_control":{"type":"ephemeral","scope":"global","ttl":"1h","future":null},"payload":7}, + {"type":"text","text":"uncached"} + ]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":nested, + "cache_control":{"scope":"global","ttl":"5m"}}]} + ], + "system":system, "tools":[tool], "future_request":null + })).unwrap(); + let stripped = super::strip_cache_control_scope(request); + let wire = serde_json::to_value(&stripped).unwrap(); + assert_eq!( + wire["messages"], + json!([ + {"role":"user","content":"plain","future_message":null}, + {"role":"assistant","content":[ + {"type":"future","cache_control":{"type":"ephemeral","ttl":"1h","future":null},"payload":7}, + {"type":"text","text":"uncached"} + ]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":nested, + "cache_control":{"ttl":"5m"}}]} + ]) + ); + assert_eq!( + wire.get("system"), + (!expected_system.is_null()).then_some(&expected_system) + ); + assert_eq!(wire["tools"], json!([tool])); + assert_eq!(wire.get("future_request"), Some(&Value::Null)); + assert_eq!(super::strip_cache_control_scope(stripped.clone()), stripped); + } + + #[rstest] + #[case::native_mid_conversation(true, json!({"role":"system","content":"later","name":"operator"}))] + #[case::converted_mid_conversation(false, json!({"role":"user","content":[ + {"type":"text","text":crate::base_llm::messages::mid_conversation_system::CONVERTED_SYSTEM_NOTE}, + {"type":"text","text":"later"} + ]}))] + fn leading_system_hoisting_and_later_system_conversion_compose( + #[case] supports_mid_conversation: bool, + #[case] expected_later: Value, + ) { + let request: MessagesRequest = serde_json::from_value(json!({ + "model":"test", "max_tokens":16, "system":"existing", + "messages":[ + {"role":"system","content":"leading"}, + {"role":"assistant","content":[{"type":"tool_use","id":"t","name":"lookup","input":{}}]}, + {"role":"system","content":"later","name":"operator"}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":"ok"}]} + ] + })).unwrap(); + let assistant = request.messages[1].clone(); + let result = request.messages[3].clone(); + let normalized = normalize_system_role_messages(request, supports_mid_conversation); + let expected_messages = if supports_mid_conversation { + json!([assistant, expected_later, result]) + } else { + json!([assistant, result, expected_later]) + }; + assert_eq!( + serde_json::to_value(&normalized).unwrap(), + json!({ + "model":"test", "max_tokens":16, + "system":[{"type":"text","text":"existing"},{"type":"text","text":"leading"}], + "messages":expected_messages + }) + ); + assert_eq!( + normalize_system_role_messages(normalized.clone(), supports_mid_conversation), + normalized + ); + } } diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 5b89345a1c5..95263392c14 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -225,7 +225,7 @@ fn filter_body(body: Value, strip_connection_fields: bool) -> Value { || name.as_str() == UNSUPPORTED_OVERRIDE)) }) .map(|(name, value)| { - let value = if name == "model" { + let value = if strip_connection_fields && name == "model" { value .as_str() .map(|model| { @@ -385,6 +385,26 @@ mod tests { ); } + #[rstest] + #[case::platform_route("claude_platform/claude-test")] + #[case::bare_model("claude-test")] + fn platform_billing_identity_is_retained_until_wire_serialization(#[case] model: &str) { + let request = serde_json::from_value(json!({ + "model": model, "max_tokens": 16, "messages": [], + })) + .unwrap(); + let transformed = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .transform_anthropic_messages_request(request, &MessagesTransformContext::default()) + .unwrap(); + assert_eq!(transformed.model, model); + let wire = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .wire_body(serde_json::to_value(transformed).unwrap()); + assert_eq!( + wire, + json!({"model":"claude-test", "max_tokens":16, "messages":[]}) + ); + } + #[rstest] #[case::default(json!({}), false)] #[case::override_empty(json!({"claude_platform_unsupported_params": []}), true)] @@ -421,4 +441,70 @@ mod tests { assert_eq!(wire.get("aws_region_name"), None); assert_eq!(wire.get(UNSUPPORTED_OVERRIDE), None); } + + #[rstest] + #[case::request(json!({"workspace_id": "request"}), json!({"workspace_id": "deployment"}), "request")] + #[case::deployment(json!({}), json!({"workspace_id": "deployment"}), "deployment")] + #[case::blank_request(json!({"workspace_id": " "}), json!({"workspace_id": "deployment"}), "deployment")] + #[case::environment(json!({}), json!({}), "environment")] + fn workspace_sources_select_one_header_without_changing_unrelated_headers( + #[case] optional: Value, + #[case] deployment: Value, + #[case] expected: &str, + ) { + let optional_params = serde_json::from_value(optional).unwrap(); + let params = LitellmParams { + extra: serde_json::from_value(deployment).unwrap(), + ..Default::default() + }; + let environment = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .validate_environment( + vec![("x-routing".into(), "keep".into())], + Some("key"), + "claude-test", + &optional_params, + ¶ms, + &|name| (name == WORKSPACE_ENVS[0]).then(|| "environment".into()), + ) + .unwrap(); + assert_eq!( + environment.headers, + vec![ + ("x-routing".into(), "keep".into()), + (WORKSPACE_HEADER.into(), expected.into()), + ] + ); + } + + #[rstest] + fn custom_unsupported_fields_do_not_remove_nested_tool_data_or_extensions() { + let input = json!({ + "model": "claude_platform/claude-test", "max_tokens": 16, "messages": [], + "workspace_id": "workspace", "aws_region_name": "us-west-2", + "claude_platform_unsupported_params": ["temperature"], + "temperature": 0.5, + "context_management": {"edits": [{"type": "clear_thinking_20251015"}]}, + "tools": [{"name": "lookup", "input_schema": { + "type": "object", "properties": {"workspace_id": {"type": "string"}}, + "aws_region_name": "application-data", "temperature": 7 + }}], + "future": null + }); + let transformed = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .transform_anthropic_messages_request( + serde_json::from_value(input.clone()).unwrap(), + &MessagesTransformContext::default(), + ) + .unwrap(); + let wire = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .wire_body(serde_json::to_value(transformed).unwrap()); + assert_eq!( + wire, + json!({ + "model": "claude-test", "max_tokens": 16, "messages": [], + "context_management": input["context_management"], + "tools": input["tools"], "future": null + }) + ); + } } diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs index 1d380c917f6..a9b918e9580 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs @@ -690,4 +690,59 @@ mod tests { (other, _) => panic!("unexpected auth {other:?}"), } } + + #[rstest] + #[case::zero_stop_counts_override_nonzero_delta( + json!({"cache_read_input_tokens": 90, "cache_creation_input_tokens": 80}), + json!({"input_tokens": 7, "output_tokens": 5, "cache_read_input_tokens": 4, "cache_creation_input_tokens": 3}), + json!({"input_tokens": 0, "output_tokens": 99, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0}), + json!({"input_tokens": 0, "output_tokens": 5, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0}) + )] + #[case::delta_extensions_outrank_start_details( + json!({"cache_creation": {"ephemeral_5m_input_tokens": 8}, "start_only": 1}), + json!({"output_tokens": 5, "cache_creation": {"ephemeral_1h_input_tokens": 3}, "future": null}), + json!({"output_tokens": 99, "stop_only": 2}), + json!({"output_tokens": 5, "cache_creation": {"ephemeral_1h_input_tokens": 3}, "future": null}) + )] + fn usage_precedence_preserves_explicit_zero_and_delta_extensions( + #[case] start: Value, + #[case] delta: Value, + #[case] stop: Value, + #[case] expected: Value, + ) { + let output = promote(vec![ + message_start(start.clone()), + message_delta(delta), + message_stop(Some(stop.clone())), + ]); + assert_eq!( + output, + vec![ + message_start(start), + message_delta(expected), + message_stop(Some(stop)), + ] + ); + } + + #[rstest] + fn only_the_last_pending_delta_receives_terminal_usage() { + let start = message_start(json!({"cache_read_input_tokens": 80})); + let first = message_delta(json!({"output_tokens": 1})); + let last = message_delta(json!({"output_tokens": 5})); + let stop = message_stop(Some( + json!({"input_tokens": 3, "cache_read_input_tokens": 100}), + )); + assert_eq!( + promote(vec![start.clone(), first.clone(), last, stop.clone()]), + vec![ + start, + first, + message_delta( + json!({"output_tokens": 5, "input_tokens": 3, "cache_read_input_tokens": 100}) + ), + stop, + ] + ); + } } diff --git a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs index d5944484522..4ea252e253b 100644 --- a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs @@ -13,7 +13,10 @@ use crate::{ Error, anthropic::{ common_utils::{filter_billing_headers_from_system, has_anthropic_credential}, - messages::transformation::update_headers_with_anthropic_beta, + messages::{ + handler::shape_anthropic_messages_request, + transformation::update_headers_with_anthropic_beta, + }, }, base_llm::{ auth::AuthScheme, @@ -33,6 +36,14 @@ pub struct EdenAIAnthropicMessagesConfig; pub const EDENAI_MESSAGES_CONFIG: EdenAIAnthropicMessagesConfig = EdenAIAnthropicMessagesConfig; impl BaseMessagesConfig for EdenAIAnthropicMessagesConfig { + fn shape_request( + &self, + request: MessagesRequest, + reasoning_auto_summary: bool, + ) -> Result { + shape_anthropic_messages_request(request, reasoning_auto_summary) + } + fn validate_environment( &self, headers: Headers, @@ -121,3 +132,118 @@ impl BaseMessagesConfig for EdenAIAnthropicMessagesConfig { serde_json::Number::from_f64(cost) } } + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + use serde_json::json; + + #[rstest] + #[case::bearer("Authorization", "Bearer caller")] + #[case::api_key("X-Api-Key", "caller")] + fn forwarded_credentials_still_require_an_edenai_key(#[case] name: &str, #[case] value: &str) { + let error = EDENAI_MESSAGES_CONFIG + .validate_environment( + vec![(name.into(), value.into())], + None, + "test-model", + &MessagesOptionalParams::default(), + &LitellmParams::default(), + &|_| None, + ) + .unwrap_err(); + assert_eq!( + error, + Error::Auth(litellm_auth::Error::MissingApiKey { + provider: "Eden AI", + environment_variable: API_KEY_ENV, + }) + ); + } + + #[rstest] + #[case::explicit(Some("explicit"), "explicit")] + #[case::missing(None, "environment")] + #[case::blank(Some(" "), "environment")] + fn configured_credentials_use_bearer_auth( + #[case] api_key: Option<&str>, + #[case] expected: &str, + ) { + let headers = vec![("x-routing".into(), "keep".into())]; + let environment = EDENAI_MESSAGES_CONFIG + .validate_environment( + headers.clone(), + api_key, + "test-model", + &MessagesOptionalParams::default(), + &LitellmParams::default(), + &|name| (name == API_KEY_ENV).then(|| "environment".into()), + ) + .unwrap(); + assert_eq!(environment.headers, headers); + assert!(matches!( + environment.auth, + AuthScheme::Credential { placement: CredentialPlacement::Bearer, secret } + if secret.expose() == expected + )); + } + + #[rstest] + #[case::punctuation("functions.lookup:0", "functions_lookup_0")] + #[case::thought_signature("lookup__thought__signature", "lookup")] + fn shaping_normalizes_matching_tool_ids_and_preserves_native_parameters( + #[case] source_id: &str, + #[case] expected_id: &str, + ) { + let input = json!({"id": source_id, "cache_control": {"scope": "application"}}); + let thinking = json!({"type": "enabled", "budget_tokens": 512, "native_field": true}); + let body = json!({ + "model": "test-model", "max_tokens": 64, "thinking": thinking, + "messages": [ + {"role": "assistant", "content": [ + {"type": "text", "text": " \n "}, + {"type": "tool_use", "id": source_id, "name": "lookup", "input": input} + ]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": source_id, "content": "ok"}]} + ] + }); + let shaped = EDENAI_MESSAGES_CONFIG + .shape_request(serde_json::from_value(body).unwrap(), false) + .unwrap(); + assert_eq!( + serde_json::to_value(shaped).unwrap(), + json!({ + "model": "test-model", "max_tokens": 64, "thinking": thinking, + "messages": [ + {"role": "assistant", "content": [{"type": "tool_use", "id": expected_id, "name": "lookup", "input": input}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": expected_id, "content": "ok"}]} + ] + }) + ); + } + + #[rstest] + fn cache_sanitization_preserves_tool_application_data_and_extensions() { + let payload = json!({"cache_control": {"type": "application", "ttl": "keep"}}); + let wire = EDENAI_MESSAGES_CONFIG.wire_body(json!({ + "model": "test-model", + "messages": [{"role": "assistant", "content": [{ + "type": "tool_use", "id": "call-1", "name": "lookup", "input": payload, + "cache_control": {"type": "ephemeral", "ttl": "1h", "scope": "global"} + }]}], + "future": null + })); + assert_eq!( + wire, + json!({ + "model": "test-model", + "messages": [{"role": "assistant", "content": [{ + "type": "tool_use", "id": "call-1", "name": "lookup", "input": payload, + "cache_control": {"type": "ephemeral"} + }]}], + "future": null + }) + ); + } +} diff --git a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs index eeb3d1cb7f8..a0b3ff6e874 100644 --- a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs @@ -198,4 +198,65 @@ mod tests { assert_eq!(body["provider"], input["provider"]); assert_eq!(body["thinking"], input["thinking"]); } + + #[rstest] + #[case::empty_bearer("Bearer ")] + #[case::basic_auth("Basic legacy")] + #[case::unrelated_scheme("Token legacy")] + fn an_invalid_forwarded_credential_is_replaced_without_leaking_an_api_key( + #[case] authorization: &str, + ) { + let environment = OPENROUTER_MESSAGES_CONFIG + .validate_environment( + vec![ + ("AuThOrIzAtIoN".into(), authorization.into()), + ("X-Api-Key".into(), "wrong-provider-key".into()), + ("x-routing".into(), "keep".into()), + ], + Some("deployment-key"), + "test-model", + &MessagesOptionalParams::default(), + &LitellmParams::default(), + &|name| (name == API_KEY_ENV).then(|| "environment-key".into()), + ) + .unwrap(); + assert_eq!( + environment.headers, + vec![("x-routing".into(), "keep".into())] + ); + assert!(matches!( + environment.auth, + AuthScheme::Credential { placement: CredentialPlacement::Bearer, secret } + if secret.expose() == "deployment-key" + )); + } + + #[rstest] + #[case::canonical("Bearer caller")] + #[case::case_and_whitespace(" bEaReR caller ")] + fn a_forwarded_bearer_wins_but_foreign_api_keys_are_removed(#[case] bearer: &str) { + let environment = OPENROUTER_MESSAGES_CONFIG + .validate_environment( + vec![ + ("Authorization".into(), bearer.into()), + ("x-api-key".into(), "foreign-key".into()), + ("X-API-KEY".into(), "another-foreign-key".into()), + ("x-routing".into(), "keep".into()), + ], + Some("deployment-key"), + "test-model", + &MessagesOptionalParams::default(), + &LitellmParams::default(), + &|name| (name == API_KEY_ENV).then(|| "environment-key".into()), + ) + .unwrap(); + assert_eq!( + environment.headers, + vec![ + ("Authorization".into(), bearer.into()), + ("x-routing".into(), "keep".into()), + ] + ); + assert!(matches!(environment.auth, AuthScheme::Forwarded)); + } } diff --git a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs index 621a09f4fb0..bfc81079a64 100644 --- a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs @@ -2,9 +2,10 @@ use litellm_auth::CredentialPlacement; use litellm_secret_value::SecretValue; use litellm_core_utils::settings::resolve_non_empty; +use litellm_http::request::without_headers; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, - providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH}, + providers::anthropic::{BETA_HEADER, DEFAULT_HEADERS, MESSAGES_PATH}, }; use litellm_router_types::LitellmParams; @@ -13,8 +14,7 @@ use crate::{ anthropic::{ common_utils::{filter_billing_headers_from_system, has_anthropic_credential}, messages::{ - handler::shape_anthropic_messages_request, - transformation::{transform_messages_request, update_headers_with_anthropic_beta}, + handler::shape_anthropic_messages_request, transformation::transform_messages_request, }, }, base_llm::{ @@ -117,8 +117,8 @@ impl BaseMessagesConfig for TencentAnthropicMessagesConfig { DEFAULT_HEADERS } - fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { - update_headers_with_anthropic_beta(headers, request) + fn request_headers(&self, headers: Headers, _request: &MessagesRequest) -> Headers { + without_headers(headers, &[BETA_HEADER]) } } @@ -132,6 +132,31 @@ mod tests { None } + #[rstest] + #[case::lowercase("anthropic-beta")] + #[case::mixed_case("Anthropic-Beta")] + fn request_headers_remove_caller_and_generated_anthropic_betas(#[case] name: &str) { + let request = serde_json::from_value(serde_json::json!({ + "model": "test-model", "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + "output_config": {"format": {"type": "json_schema", "schema": {"type": "object"}}}, + "tools": [{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}] + })) + .unwrap(); + let retained = vec![("x-routing".into(), "keep".into())]; + let headers = vec![ + ( + name.into(), + "caller-beta,structured-outputs-2025-11-13".into(), + ), + ("x-routing".into(), "keep".into()), + ]; + assert_eq!( + TENCENT_MESSAGES_CONFIG.request_headers(headers, &request), + retained + ); + } + #[rstest] #[case::root("https://proxy.test", "https://proxy.test/v1/messages")] #[case::root_slash("https://proxy.test/", "https://proxy.test/v1/messages")] diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 426cd861729..3f56725def1 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -23,6 +23,7 @@ pyo3::create_exception!( ); const MESSAGES_REQUEST_ERROR: &str = "messages_request_error"; +const MESSAGES_RESPONSE_ERROR: &str = "messages_response_error"; /// The rejection is a parameter the model does not support, which Python raises as /// `UnsupportedParamsError`. const MESSAGES_UNSUPPORTED_PARAM: &str = "messages_unsupported_param"; @@ -137,6 +138,9 @@ pub(crate) fn route_failure( if error.is_missing_credential() { value.setattr(MESSAGES_MISSING_CREDENTIAL, true)?; } + if error.is_response() { + value.setattr(MESSAGES_RESPONSE_ERROR, true)?; + } native } (RequestMarker::Unmarked, _) => by_fault(error.is_request(), &error), diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index af614f17360..1a18bf7a4c5 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -186,13 +186,11 @@ impl MessagesPythonHost { } fn provider(&self, py: Python<'_>) -> String { - self.view - .request(py) - .and_then(|request| request.get_item("custom_llm_provider")) - .ok() - .flatten() - .and_then(|value| value.extract::>().ok().flatten()) - .unwrap_or_else(|| "anthropic".into()) + py.import(ROUTE_HOST_MODULE) + .and_then(|module| module.getattr("request_provider")) + .and_then(|resolve| resolve.call1((self.view.request(py)?,))) + .and_then(|value| value.extract()) + .unwrap_or_else(|_| "anthropic".into()) } fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index a9ae01bd7ee..e43369e0cd2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -7,6 +7,21 @@ use pyo3::prelude::*; use super::NativeCall; fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult> { + let request = call.resolved()?; + let model: String = request + .get_item("model")? + .ok_or_else(|| pyo3::exceptions::PyTypeError::new_err("messages requires a model"))? + .extract()?; + let provider: String = py + .import("litellm.rust_bridge.host.routes.messages")? + .getattr("request_provider")? + .call1((&request,))? + .extract()?; + if !litellm_inference_messages::supports_messages_provider(&model, Some(&provider)) { + return Err(crate::errors::RustBridgeDeclined::new_err(format!( + "messages provider {provider} has no native adapter" + ))); + } let (arguments, hooks) = crate::routes::call_hooks(py, LoggingOperation::Messages, &call, asynchronous)?; crate::routes::run_public_call( diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 4c94f26c6c1..8ef09676a84 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -69,8 +69,7 @@ def _public_request( return None model: Final = fields.get("model") messages: Final = optional_sequence(fields.get("messages")) - max_tokens: Final = fields.get("max_tokens") - if not isinstance(model, str) or messages is None or not isinstance(max_tokens, int): + if not isinstance(model, str) or messages is None: return None return native_call(legacy, args, _native_kwargs(fields, kwargs)) diff --git a/litellm/router_rust/invoke/streaming.py b/litellm/router_rust/invoke/streaming.py index d8b4f87501a..164708e007a 100644 --- a/litellm/router_rust/invoke/streaming.py +++ b/litellm/router_rust/invoke/streaming.py @@ -1111,6 +1111,8 @@ class FallbackAwareAnthropicMessagesStream: getattr(source_iterator, "_hidden_params", None) or {} ) self._followed_source_params: object = None + self._started: bool = False + self._closed: bool = False @property def has_buffered_provider_output(self) -> bool: @@ -1154,10 +1156,18 @@ class FallbackAwareAnthropicMessagesStream: return self async def __anext__(self) -> bytes: + self._started = True return await self._async_generator.__anext__() async def aclose(self) -> None: + from litellm.llms.anthropic.messages.native.streaming_iterator import aclose_if_supported + + if self._closed: + return await self._async_generator.aclose() + self._closed = True + if not self._started: + await aclose_if_supported(self._source_iterator) def merge_fallback_hidden_params( self, fallback_hidden_params: Mapping[str, object], fallback_headers: Mapping[str, object] @@ -1295,15 +1305,13 @@ def wrap_anthropic_stream( if _anthropic_stream_forwards_ping_live(chunk, has_generated_content): yield chunk continue - if _stream_commits_now(chunk, has_generated_content, len(buffered)): - has_generated_content = True - # A transport can split one SSE data line across byte chunks, so pre-content detection - # parses the held frames plus the current chunk. parse_window = ( b"".join(c for c in (*buffered, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # see above else chunk ) + if _stream_commits_now(parse_window, has_generated_content, len(buffered)): + has_generated_content = True error_event = parse_anthropic_error_event(parse_window) bridged_failure = _bridged_stream_failure(source_iterator) if error_event is not None else None if bridged_failure is not None: diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index c8b53097f46..91610505060 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -23,6 +23,7 @@ from litellm.llms.anthropic.messages.native.handler import ( _with_deployment_output_effort as _deployment_output_effort, # pyright: ignore[reportPrivateUsage] # the handler's own effort rounding ) from litellm.llms.anthropic.messages.utils import is_reasoning_auto_summary_enabled +from litellm.llms.bedrock.claude_platform.common_utils import CLAUDE_PLATFORM_ON_AWS_NON_REQUEST_PARAMS from litellm.rust_bridge import failures from litellm.rust_bridge.host import callbacks_legacy_python from litellm.rust_bridge.public_call import known_params, optional_str @@ -84,8 +85,20 @@ def inferred_provider(model: str, api_base: str | None) -> str | None: return provider +def request_provider(request: Mapping[str, object]) -> str: + return ( + optional_str(request.get("custom_llm_provider")) + or inferred_provider(str(request["model"]), optional_str(request.get("api_base"))) + or "anthropic" + ) + + def body_params(request: Mapping[str, object]) -> Mapping[str, object]: - params: Final = known_params(_BODY_FIELDS, request) + platform: Final = request_provider(request) == "bedrock" and str(request["model"]).removeprefix( + "bedrock/" + ).startswith("claude_platform/") + fields: Final = _BODY_FIELDS | CLAUDE_PLATFORM_ON_AWS_NON_REQUEST_PARAMS if platform else _BODY_FIELDS + params: Final = known_params(fields, request) if deployment_reasoning_effort_levels(request.get("model_info")) is None: return params return _with_deployment_effort(params, request) @@ -236,6 +249,15 @@ def map_failure(error: Exception, request: Mapping[str, object], request_provide _log_resolved_call(request, request_provider) if getattr(error, "messages_request_error", False): return _rejection(error, request, request_provider) + if getattr(error, "messages_response_error", False): + response_error: Final = litellm.APIError( + status_code=500, + message=str(error), + llm_provider=request_provider, + model=str(request["model"]).removeprefix(f"{request_provider}/"), + ) + response_error.__context__ = error + return response_error if getattr(error, "messages_missing_credential", False): return litellm.AuthenticationError( message=_MISSING_KEY, diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py index 41008f0d57f..97c291fa0c0 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py @@ -769,7 +769,119 @@ FALLBACK_CASES: Final = ( Case("stream_timeout_kwarg", _kw(stream=True, stream_timeout=30), upstream=(_stream(TEXT_STREAM),)), ) -CASES: Final = (*REQUEST_CASES, *RESPONSE_CASES, *FAILURE_CASES, *STREAM_CASES, *SYNC_CASES, *FALLBACK_CASES) +CONTRACT_REGRESSION_CASES: Final = ( + Case( + "tencent_filters_caller_and_feature_betas", + _kw( + extra_headers={"Anthropic-Beta": "structured-outputs-2025-11-13,caller-beta", "x-routing": "keep"}, + output_config={"format": {"type": "json_schema", "schema": {"type": "object"}}}, + tools=[{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}], + ), + model="tencent/claude-sonnet-4-5", + ), + Case( + "edenai_normalizes_replayed_tool_ids", + _kw( + messages=[ + {"role": "user", "content": "Look it up"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": " \n "}, + { + "type": "tool_use", + "id": "functions.lookup:0", + "name": "lookup", + "input": {"id": "functions.lookup:0"}, + }, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "functions.lookup:0", "content": "ok"}], + }, + ] + ), + model="edenai/anthropic/claude-sonnet-4-5", + ), + Case( + "signature_only_thinking_replay", + _kw( + messages=[ + {"role": "user", "content": "Look it up"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "", "signature": "sig_abc"}, + {"type": "tool_use", "id": "t", "name": "lookup", "input": {}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t", "content": "ok"}]}, + ] + ), + ), + Case( + "converse_wrapped_redacted_replay", + _kw( + messages=[ + {"role": "user", "content": "Solve it"}, + { + "role": "assistant", + "content": [ + { + "type": "redacted_thinking", + "data": "RW13S0FoZ0JFZ3kzdmEzcHppeC9MYWZQc240YURCSVQyWGx4aDBMNUw4ckxWeUl3eHRFM3JBRkJhOGNyM3FwUA==", + }, + {"type": "text", "text": "answer"}, + ], + }, + {"role": "user", "content": "Continue"}, + ] + ), + ), + Case( + "compaction_threshold_exact_integer", + _kw(context_management=[{"type": "compaction", "compact_threshold": 9007199254740993}]), + ), + Case( + "compaction_boolean_threshold", + _kw(context_management=[{"type": "compaction", "compact_threshold": True}]), + ), + Case( + "unknown_content_familiar_null_fields", + _kw( + messages=[ + { + "role": "user", + "content": [ + { + "type": "future_block", + "text": None, + "input": None, + "source": {"data": "123"}, + "future": None, + } + ], + } + ] + ), + ), + Case( + "unicode_legacy_thinking_budget", + _kw(thinking={"type": "enabled", "budget_tokens": " +8_١92 "}, max_tokens=16000), + model=OPUS_ADAPTIVE, + ), +) + +CASES: Final = ( + *REQUEST_CASES, + *RESPONSE_CASES, + *FAILURE_CASES, + *STREAM_CASES, + *SYNC_CASES, + *FALLBACK_CASES, + *CONTRACT_REGRESSION_CASES, +) @dataclass(frozen=True, slots=True) diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py index c9ac7af1a78..cd88408ee34 100644 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ b/tests/test_litellm_rust/cache/test_python_cache.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import AsyncIterator from typing import Final, Literal import pytest @@ -19,10 +20,10 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelReq from litellm.proxy.utils import InternalUsageCache from litellm.rust_bridge.catalog import Route from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload -from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE +from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE, call from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec -from tests.test_litellm_rust.support.routing import native_route +from tests.test_litellm_rust.support.routing import assert_served_natively, native_route pytestmark = pytest.mark.requires_rust_extension @@ -360,3 +361,103 @@ async def test_legacy_cache_keeps_public_messages_responses_compatible( second_payload: Final = await collect(second) if stream else payload(second) assert second_payload == first_payload assert len(recording_server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream")) +async def test_cancelling_a_cache_write_does_not_store_or_log_success( + recording_server: RecordingServer, stream: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.caching.base_cache import BaseCache + from litellm.caching.in_memory_cache import InMemoryCache + + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + storage: Final = InMemoryCache() + + class GatedCache(BaseCache): + async def async_get_cache(self, key: str, **kwargs: object) -> object: + return storage.get_cache(key, **kwargs) + + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: + entered.set() + await release.wait() + storage.set_cache(key, value, **kwargs) + + async def async_set_cache_pipeline( + self, cache_list: list[tuple[str, object]], ttl: float | None = None + ) -> None: + entered.set() + await release.wait() + await storage.async_set_cache_pipeline(cache_list, ttl=ttl) + + monkeypatch.setattr(litellm, "cache", Cache(_backend=GatedCache())) + recording_server.expected_requests = 2 + recording_server.default_response = ( + ResponseSpec(body=None, events=MESSAGES_EVENTS) if stream else ResponseSpec(body=MESSAGES_RESPONSE) + ) + recorder: Final = RecordingLogger() + + async def complete_call() -> object: + response: Final = await call(recording_server, stream=stream, callbacks=[recorder]) + if stream: + await collect(response) + return response + + with native_route(Route.MESSAGES): + task: Final = asyncio.create_task(complete_call()) + await asyncio.wait_for(entered.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + release.set() + await drain_logging() + assert storage.cache_dict == {} + assert len(recording_server.requests) == 1 + assert not any("success" in name or "failure" in name for name in recorder.names) + + second: Final = await call(recording_server, stream=stream) + assert cache_key(second) is None + if stream: + await collect(second) + third: Final = await call(recording_server, stream=stream) + assert isinstance(cache_key(third), str) + if stream: + await collect(third) + + assert_served_natively(recording_server, count=2) + assert len(recording_server.requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("first_stream", (False, True), ids=("response_first", "stream_first")) +@pytest.mark.parametrize("preset", (False, True), ids=("normal_key", "preset_key")) +async def test_cached_messages_preserve_the_requested_response_mode( + recording_server: RecordingServer, first_stream: bool, preset: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "cache", Cache()) + recording_server.expected_requests = None + + async def request(stream: bool) -> object: + recording_server.default_response = ( + ResponseSpec(body=None, events=MESSAGES_EVENTS) if stream else ResponseSpec(body=MESSAGES_RESPONSE) + ) + options: Final = {"litellm_params": {"preset_cache_key": "messages-mode"}} if preset else {} + response: Final = await call(recording_server, stream=stream, **options) + assert isinstance(response, AsyncIterator) is stream + if stream: + assert await collect(response) == b"".join(recording_server.default_response.payloads()) + else: + assert isinstance(response, dict) + assert response["id"] == MESSAGES_RESPONSE["id"] + return response + + with native_route(Route.MESSAGES): + await request(first_stream) + second: Final = await request(not first_stream) + assert cache_key(second) is None + third: Final = await request(not first_stream) + assert isinstance(cache_key(third), str) + + assert_served_natively(recording_server, count=2) + assert len(recording_server.requests) == 2 diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 8bcb48a09b5..47a31f9cb54 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -7,6 +7,7 @@ import litellm from tests.test_litellm_rust.support.clickhouse import clickhouse_url as clickhouse_url from tests.test_litellm_rust.support.isolation import isolated_test_state from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service +from tests.test_litellm_rust.support.served import NativeEntry, recording_native_entries @pytest_asyncio.fixture(autouse=True, loop_scope="function") @@ -15,9 +16,16 @@ async def isolate_test_state() -> AsyncIterator[None]: yield +@pytest.fixture(autouse=True) +def native_entries() -> Generator[list[NativeEntry]]: + with recording_native_entries() as entries: + yield entries + + @pytest.fixture -def recording_server() -> Generator[RecordingServer]: +def recording_server(native_entries: list[NativeEntry]) -> Generator[RecordingServer]: with recording_service() as server: + server.native_entries = native_entries yield server diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index a3c03898a12..8ccd2e66799 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -1,8 +1,9 @@ import asyncio -from collections.abc import Iterator +import math +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator from contextvars import ContextVar from datetime import datetime -from typing import Final +from typing import Final, Protocol, cast # noqa: TID251 # narrows public exception attributes from the legacy mapper import pytest from pydantic import TypeAdapter @@ -12,7 +13,7 @@ from litellm.integrations.argilla import ArgillaLogger from litellm.integrations.custom_logger import CustomLogger from litellm.types.integrations.argilla import ArgillaItem from litellm.types.utils import ModelResponse, StandardLoggingPayload -from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, call +from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_EVENTS, MESSAGES_RESPONSE, RATE_LIMITED, call from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed @@ -424,3 +425,167 @@ async def test_sync_call_raises_like_python_without_logging(server: RecordingSer assert not server.requests assert not recorder.names + + +@pytest.mark.asyncio +@pytest.mark.parametrize("explicit_provider", (False, True)) +@pytest.mark.parametrize("status", (401, 429)) +async def test_inferred_provider_is_preserved_in_native_public_errors( + server: RecordingServer, explicit_provider: bool, status: int +) -> None: + server.default_response = ResponseSpec( + body={"error": {"type": "provider_error", "message": "synthetic refusal"}}, status=status + ) + error_type: Final = litellm.AuthenticationError if status == 401 else litellm.RateLimitError + with pytest.raises(error_type) as raised: + await call( + server, model="deepseek/deepseek-v4-flash", custom_llm_provider="deepseek" if explicit_provider else None + ) + assert_served_natively(server) + public_error: Final = cast( + _ProviderError, raised.value + ) # cast-ok: the exception mapper supplies these public attributes + assert public_error.llm_provider == "deepseek" + assert public_error.model == "deepseek/deepseek-v4-flash" + assert "synthetic refusal" in str(raised.value) + + +class _ProviderError(Protocol): + @property + def llm_provider(self) -> str: ... + + @property + def model(self) -> str: ... + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", (2**64 - 1, -(2**63), "9" * 4096, 1e30)) +async def test_native_response_preserves_supported_numbers_and_numeric_strings( + server: RecordingServer, value: int | str | float +) -> None: + server.default_response = ResponseSpec(body={**MESSAGES_RESPONSE, "future_extension": {"wide_integer": value}}) + response: Final = await call(server) + assert_served_natively(server, response) + body: Final = _OBJECT.validate_python(response) + actual: Final = _OBJECT.validate_python(body["future_extension"])["wide_integer"] + assert actual == value + assert type(actual) is type(value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", (2**64 + 1, -(2**63) - 1)) +async def test_native_response_rejects_oversized_integers_instead_of_returning_rounded_success( + server: RecordingServer, value: int +) -> None: + server.default_response = ResponseSpec(body={**MESSAGES_RESPONSE, "future_extension": {"wide_integer": value}}) + with pytest.raises(litellm.APIError, match="response integer outside") as raised: + await call(server) + assert raised.value.status_code == 500 + assert not isinstance(raised.value, litellm.APIConnectionError) + assert str(value) not in str(raised.value) + assert_served_natively(server) + + +class _MessagesRouter(Protocol): + def aanthropic_messages(self, **kwargs: object) -> Awaitable[object]: ... + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entry", ("provider_prefix", "explicit_provider", "router_alias")) +@pytest.mark.parametrize("streaming", (False, True), ids=("unary", "stream")) +async def test_platform_billing_identity_survives_wire_model_rewriting( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, entry: str, streaming: bool +) -> None: + model: Final = "claude-billing-identity-test" + backend: Final = f"bedrock/claude_platform/{model}" + input_price: Final = 0.002 + output_price: Final = 0.003 + prices: Final = cast( + dict[str, object], getattr(litellm, "model_cost") + ) # cast-ok: model_cost is the mutable runtime registry + monkeypatch.setitem( + prices, + backend, + { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": input_price, + "output_cost_per_token": output_price, + "max_input_tokens": 200000, + "max_output_tokens": 64000, + }, + ) + provider_response: Final = {**MESSAGES_RESPONSE, "model": model} + events: Final = ( + ( + "message_start", + {"type": "message_start", "message": {**provider_response, "content": [], "stop_reason": None}}, + ), + *MESSAGES_EVENTS[1:], + ) + server.default_response = ResponseSpec(body=provider_response, events=events if streaming else ()) + recorder: Final = RecordingLogger() + metadata: Final = {"caller_marker": "cost-context", "model_group": "billing-alias"} + + async def invoke() -> object: + if entry == "router_alias": + factory: Final = cast( + Callable[..., _MessagesRouter], litellm.Router + ) # cast-ok: Router exposes its public constructor and endpoint dynamically + router: Final = factory( + model_list=[ + { + "model_name": "billing-alias", + "litellm_params": { + "model": backend, + "api_base": server.base_url, + "api_key": "synthetic", + "workspace_id": "billing-workspace", + }, + } + ], + num_retries=0, + disable_cooldowns=True, + ) + return await router.aanthropic_messages( + model="billing-alias", + messages=[{"role": "user", "content": "hi"}], + max_tokens=16, + litellm_metadata=metadata, + stream=streaming, + ) + return await call( + server, + model=f"claude_platform/{model}" if entry == "explicit_provider" else backend, + custom_llm_provider="bedrock" if entry == "explicit_provider" else None, + workspace_id="billing-workspace", + litellm_metadata=metadata, + stream=streaming, + ) + + with rebound(litellm, "callbacks", [recorder]): + response: Final = await invoke() + if streaming: + stream: Final = cast(AsyncIterator[bytes], response) # cast-ok: native Messages streams yield bytes + chunks: Final = [chunk async for chunk in stream] + assert b"".join(chunks) == b"".join(server.default_response.payloads()) + successes: Final = await recorder.wait_for_async("async_log_success_event") + if streaming: + assert_served_natively(server) + else: + assert_served_natively(server, response) + assert _OBJECT.validate_python(response)["model"] == model + usage: Final = _OBJECT.validate_python(provider_response["usage"]) + expected: Final = ( + TypeAdapter(int).validate_python(usage["input_tokens"]) * input_price + + TypeAdapter(int).validate_python(usage["output_tokens"]) * output_price + ) + logged: Final = _OBJECT.validate_python(successes[0].kwargs) + standard: Final = _OBJECT.validate_python(logged["standard_logging_object"]) + assert math.isclose(TypeAdapter(float).validate_python(logged["response_cost"]), expected) + assert math.isclose(TypeAdapter(float).validate_python(standard["response_cost"]), expected) + pre_call: Final = _OBJECT.validate_python(recorder.wait_for("log_pre_api_call")[0].kwargs) + assert pre_call["model"] == f"claude_platform/{model}" + assert _OBJECT.validate_python(server.requests[0].body)["model"] == model + params: Final = _OBJECT.validate_python(logged["litellm_params"]) + assert _OBJECT.validate_python(params["metadata"])["caller_marker"] == "cost-context" diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 6d65f289524..6deaed93ca6 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -22,6 +22,8 @@ from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_ME from litellm.rust_bridge.public_call import NativeCall from litellm.types.llms.anthropic import ANTHROPIC_ADVISOR_TOOL_TYPE from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from tests.test_litellm_rust.messages.data import call +from tests.test_litellm_rust.support.recording_server import NativeEntry, RecordingServer MESSAGES: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final[Rules] = () @@ -659,3 +661,63 @@ async def test_interceptor_feature_follows_the_resolved_provider() -> None: rules=rules, ) assert result is native_result + + +@pytest.mark.asyncio +async def test_generic_messages_admission_declines_before_provider_or_hooks(server: RecordingServer) -> None: + from tests.test_litellm_rust.support.callback_recorder import RecordingLogger + from tests.test_litellm_rust.support.isolation import rebound + from tests.test_litellm_rust.support.routing import routed + + recorder: Final = RecordingLogger() + server.expected_requests = 0 + with routed(*RUST_RULES), rebound(litellm, "callbacks", [recorder]): + with pytest.raises(RuntimeError, match=r"declined.*no native adapter"): + await call( + server, + model="openai/gpt-5", + max_tokens=16, + messages=MESSAGES, + model_info={"supported_endpoints": ["/v1/messages"]}, + api_key="synthetic", + api_base=server.base_url, + ) + assert server.requests == [] + assert recorder.names == () + assert server.native_entries == [NativeEntry(server.base_url, declined=True)] + + +@pytest.mark.asyncio +async def test_python_traffic_cannot_satisfy_native_execution_evidence(server: RecordingServer) -> None: + from tests.test_litellm_rust.support.routing import assert_served_natively, routed + + with routed(*PYTHON_RULES): + await call( + server, + model="anthropic/claude-test", + max_tokens=16, + messages=MESSAGES, + api_key="synthetic", + api_base=server.base_url, + ) + assert server.native_entries == [] + with pytest.raises(AssertionError, match="native execution evidence"): + assert_served_natively(server) + + +@pytest.mark.asyncio +async def test_sync_entry_async_flag_executes_the_real_native_binding(server: RecordingServer) -> None: + from tests.test_litellm_rust.support.routing import assert_served_natively, routed + + with routed(*RUST_RULES): + pending: Final = litellm.anthropic_messages_handler( + model="anthropic/claude-test", + max_tokens=16, + messages=MESSAGES, + api_key="synthetic", + api_base=server.base_url, + is_async=True, + ) + assert inspect.isawaitable(pending) + response: Final = await pending + assert_served_natively(server, response) diff --git a/tests/test_litellm_rust/messages/test_lifecycle.py b/tests/test_litellm_rust/messages/test_lifecycle.py new file mode 100644 index 00000000000..e2bf54257aa --- /dev/null +++ b/tests/test_litellm_rust/messages/test_lifecycle.py @@ -0,0 +1,93 @@ +import asyncio +import json +import threading +from collections.abc import Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final + +import pytest + +from litellm import Router +from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_MODEL +from tests.test_litellm_rust.support.recording_server import RecordedRequest, RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark: Final = pytest.mark.requires_rust_extension + + +@pytest.fixture +def unfinished_messages_peer() -> Iterator[tuple[RecordingServer, threading.Event]]: + requests: Final[list[RecordedRequest]] = [] # mutable-ok: HTTP handler records observed requests + disconnected: Final = threading.Event() + prefix: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS[:1]).payloads()[0] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + body: Final = self.rfile.read(int(self.headers.get("content-length", "0"))) + requests.append( + RecordedRequest( + method="POST", + path=self.path, + headers={name.lower(): value for name, value in self.headers.items()}, + raw_body=body, + body=json.loads(body), + ) + ) + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.send_header("content-length", str(len(prefix) + 1024)) + self.end_headers() + self.wfile.write(prefix) + self.wfile.flush() + self.connection.settimeout(5) + try: + if self.rfile.read(1) == b"": + disconnected.set() + except ConnectionResetError: + disconnected.set() + except TimeoutError: + pass + + def log_message(self, format: str, *args: object) -> None: + pass + + http: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + http.daemon_threads = True + thread: Final = threading.Thread(target=http.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + try: + yield RecordingServer(http, requests, [], ResponseSpec(body={})), disconnected + finally: + http.shutdown() + http.server_close() + thread.join() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("read_first", [False, True], ids=["before-first-read", "after-first-read"]) +async def test_router_closing_a_messages_stream_releases_the_unfinished_provider_connection( + unfinished_messages_peer: tuple[RecordingServer, threading.Event], read_first: bool +) -> None: + server, disconnected = unfinished_messages_peer + router: Final = Router( + model_list=[ + { + "model_name": "messages", + "litellm_params": {"model": MESSAGES_MODEL, "api_base": server.base_url, "api_key": "fake-key"}, + } + ], + num_retries=0, + disable_cooldowns=True, + ) + stream: Final = await router.aanthropic_messages( + model="messages", messages=[{"role": "user", "content": "hello"}], max_tokens=64, stream=True + ) + if read_first: + assert await anext(stream) == ResponseSpec(body=None, events=MESSAGES_EVENTS[:1]).payloads()[0] + await stream.aclose() + await stream.aclose() + + assert await asyncio.to_thread(disconnected.wait, 2), "closing the Router stream must close upstream" + assert_served_natively(server, stream) + with pytest.raises(StopAsyncIteration): + await anext(stream) diff --git a/tests/test_litellm_rust/messages/test_requests.py b/tests/test_litellm_rust/messages/test_requests.py index 92031e72aee..ab8ab02754f 100644 --- a/tests/test_litellm_rust/messages/test_requests.py +++ b/tests/test_litellm_rust/messages/test_requests.py @@ -146,3 +146,45 @@ async def test_optional_positional_body_parameters_reach_the_wire(server: Record assert response["id"] == MESSAGES_RESPONSE["id"] body, _ = sent(server) assert (body["metadata"], body["system"], body["temperature"]) == (args[3], args[6], args[7]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "workspace_field", ("workspace_id", "aws_workspace_id", "anthropic_workspace_id", "anthropic-workspace-id") +) +async def test_platform_workspace_and_unsupported_override_reach_the_native_adapter( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, workspace_field: str +) -> None: + monkeypatch.delenv("ANTHROPIC_AWS_WORKSPACE_ID", raising=False) + monkeypatch.delenv("ANTHROPIC_WORKSPACE_ID", raising=False) + response: Final = await call( + server, + model="bedrock/claude_platform/claude-test", + asynchronous=True, + **{workspace_field: "deployment-workspace"}, + claude_platform_unsupported_params=["temperature"], + temperature=0.3, + tools=[ + {"name": "lookup", "input_schema": {"type": "object", "properties": {"temperature": {"type": "number"}}}} + ], + ) + assert_served_natively(server, response) + body, headers = sent(server) + assert headers["anthropic-workspace-id"] == "deployment-workspace" + assert "temperature" not in body + assert "claude_platform_unsupported_params" not in body + assert workspace_field not in body + assert body["tools"] == [ + {"name": "lookup", "input_schema": {"type": "object", "properties": {"temperature": {"type": "number"}}}} + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("max_tokens", (1.5, True, "16", None, [], {})) +async def test_malformed_max_tokens_is_rejected_natively_before_io(server: RecordingServer, max_tokens: object) -> None: + server.expected_requests = 0 + with pytest.raises( + litellm.BadRequestError, match=r"invalid Anthropic messages request|missing required field: max_tokens" + ): + await call(server, max_tokens=max_tokens) + assert_served_natively(server, count=0) diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py index e57631f48b9..d8ddd169e7f 100644 --- a/tests/test_litellm_rust/messages/test_streaming.py +++ b/tests/test_litellm_rust/messages/test_streaming.py @@ -56,15 +56,37 @@ async def test_sync_stream_raises_like_python_before_sending(server: RecordingSe @pytest.mark.asyncio -async def test_stream_closed_early_logs_success_once(server: RecordingServer) -> None: - server.enqueue(STREAM) +async def test_stream_closed_early_logs_only_delivered_usage_once(server: RecordingServer) -> None: + first_event: Final = MESSAGES_EVENTS[0] + start_message: Final = { + **MESSAGES_RESPONSE, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 5, "output_tokens": 0}, + } + events: Final = ( + (first_event[0], {"type": "message_start", "message": start_message}), + *MESSAGES_EVENTS[1:], + ) + response: Final = ResponseSpec(body=None, events=events, chunk_delay=0.1) + server.enqueue(response) recorder: Final = RecordingLogger() stream: Final = await call(server, stream=True, callbacks=[recorder]) assert isinstance(stream, AsyncIterator) - await anext(stream) + first: Final = await anext(stream) + assert first == response.payloads()[0] + assert b"message_stop" not in first + await stream.aclose() await stream.aclose() - assert len(await recorder.wait_for_async("async_log_success_event")) == 1 + assert_served_natively(server) + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert success[0].response.usage.prompt_tokens == start_message["usage"]["input_tokens"] + assert success[0].response.usage.completion_tokens == 0 + assert not any("failure" in name for name in recorder.names) with pytest.raises(StopAsyncIteration): await anext(stream) + await drain_logging() + assert len(tuple(event for event in recorder.events if event.name == "async_log_success_event")) == 1 diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py index 5743658abdf..b6454997805 100644 --- a/tests/test_litellm_rust/support/recording_server.py +++ b/tests/test_litellm_rust/support/recording_server.py @@ -10,6 +10,12 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Final +@dataclass(frozen=True, slots=True) +class NativeEntry: + api_base: str | None + declined: bool + + @dataclass class RecordedRequest: method: str @@ -45,6 +51,7 @@ class RecordingServer: responses: list[ResponseSpec] default_response: ResponseSpec expected_requests: int | None = 1 + native_entries: list[NativeEntry] = field(default_factory=list) @property def base_url(self) -> str: diff --git a/tests/test_litellm_rust/support/served.py b/tests/test_litellm_rust/support/served.py index 0890077f745..76ac771a814 100644 --- a/tests/test_litellm_rust/support/served.py +++ b/tests/test_litellm_rust/support/served.py @@ -1,16 +1,104 @@ -from typing import Final +from collections.abc import Awaitable, Callable, Generator +from contextlib import contextmanager +from typing import Final, cast # noqa: TID251 # observation adapters preserve the native callable protocol + +from pydantic import TypeAdapter from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict -from tests.test_litellm_rust.support.recording_server import RecordingServer +from litellm.rust_bridge import _native +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.chat_completions.entrypoints import NATIVE_ACOMPLETION, NATIVE_COMPLETION +from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES +from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR +from litellm.rust_bridge.public_call import NativeCall, optional_str +from litellm.rust_bridge.responses.entrypoints import NATIVE_ARESPONSES, NATIVE_RESPONSES +from tests.test_litellm_rust.support.recording_server import NativeEntry, RecordingServer + +SyncEntry = Callable[[NativeCall], object] +AsyncEntry = Callable[[NativeCall], Awaitable[object]] + + +def _sync_observer(native: SyncEntry, entries: list[NativeEntry]) -> SyncEntry: + def observed(request: NativeCall) -> object: + try: + result: Final = native(request) + except _native.RustBridgeDeclined: + entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), True)) + raise + except BaseException: + entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), False)) + raise + entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), False)) + return result + + return observed + + +def _async_observer(native: AsyncEntry, entries: list[NativeEntry]) -> AsyncEntry: + async def observed(request: NativeCall) -> object: + try: + result: Final = await native(request) + except _native.RustBridgeDeclined: + entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), True)) + raise + except BaseException: + entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), False)) + raise + entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), False)) + return result + + return observed + + +def _observe_sync_binding(binding: NativeBinding[SyncEntry], entries: list[NativeEntry]) -> None: + native: Final = binding.load() + if native is not None: + binding.override(_sync_observer(native, entries)) + + +def _observe_async_binding(binding: NativeBinding[AsyncEntry], entries: list[NativeEntry]) -> None: + native: Final = binding.load() + if native is not None: + binding.override(_async_observer(native, entries)) + + +@contextmanager +def recording_native_entries() -> Generator[list[NativeEntry]]: + entries: Final[list[NativeEntry]] = [] + synchronous: Final = cast( # cast-ok: all sync route bindings accept NativeCall and return a public result + tuple[NativeBinding[SyncEntry], ...], + (NATIVE_MESSAGES, NATIVE_RESPONSES, NATIVE_COMPLETION, NATIVE_OCR), + ) + asynchronous: Final = cast( # cast-ok: all sync route bindings accept NativeCall and return a public result + tuple[NativeBinding[AsyncEntry], ...], + (NATIVE_AMESSAGES, NATIVE_ARESPONSES, NATIVE_ACOMPLETION, NATIVE_AOCR), + ) + try: + for binding in synchronous: + _observe_sync_binding(binding, entries) + for async_binding in asynchronous: + _observe_async_binding(async_binding, entries) + yield entries + finally: + for binding in (*synchronous, *asynchronous): + binding.reset() def assert_served_natively(server: RecordingServer, *responses: object, count: int = 1) -> None: assert len(server.requests) == (len(responses) or count) if not responses: - assert not any( - request.headers.get("user-agent", "").startswith("python-httpx") for request in server.requests + accepted: Final = tuple( + entry + for entry in server.native_entries + if not entry.declined and (entry.api_base is None or entry.api_base.startswith(server.base_url)) ) + assert len(accepted) >= max(count, 1), "no matching native execution evidence" for response in responses: - additional_headers: Final = get_hidden_params_dict(response).get("additional_headers") - assert isinstance(additional_headers, dict) - assert additional_headers.get("x-litellm-rust") == "true" + _assert_response_marker(response) + + +def _assert_response_marker(response: object) -> None: + headers: Final = TypeAdapter(dict[str, object]).validate_python( + get_hidden_params_dict(response).get("additional_headers") + ) + assert headers.get("x-litellm-rust") == "true" diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index 9d8b09dd1c4..8def0fe4b4c 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -245,3 +245,14 @@ def test_a_late_rejection_tells_the_logger_the_resolved_call_first() -> None: assert isinstance(mapped, litellm.BadRequestError) assert logger.model_call_details["custom_llm_provider"] == "anthropic" assert logger.model_call_details["optional_params"] == {"max_tokens": 8} + + +def test_invalid_provider_response_is_an_api_error_with_its_original_cause() -> None: + error: Final = RuntimeError("invalid messages response integer outside i64/u64 range") + error.messages_response_error = True # pyright: ignore[reportAttributeAccessIssue] # marker supplied by the native boundary + mapped: Final = messages_host.map_failure(error, {"model": "deepseek/model"}, "deepseek") + assert isinstance(mapped, litellm.APIError) + assert not isinstance(mapped, litellm.APIConnectionError) + assert mapped.status_code == 500 + assert mapped.__context__ is error + assert "response integer outside" in str(mapped) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index dbc47fd6962..9eb98c46432 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -13,7 +13,7 @@ import os import sys import threading import warnings -from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta, timezone from pathlib import Path from types import CodeType, FrameType, MappingProxyType, SimpleNamespace @@ -52,6 +52,7 @@ from litellm.router_rust.invoke.streaming import ( MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS, FallbackAwareAnthropicMessagesStream, + StreamAttempt, _stream_commits_now, _anthropic_stream_error_is_gateway_verdict, _stream_fallback_error_for_raised, @@ -19879,3 +19880,103 @@ def test_router_warns_when_a_deployment_enforces_endpoints_without_a_list(caplog assert "Deployment 'embed-from-config' sets enforce_supported_endpoints without" in caplog.text assert "Deployment 'embed-from-api' sets enforce_supported_endpoints without" in caplog.text + + +class _PartitionedMessagesSSE(httpx.AsyncByteStream): + def __init__(self, chunks: tuple[bytes, ...]) -> None: + self._chunks: Final = chunks + self.reads = 0 + self.closes = 0 + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self._chunks: + self.reads += 1 + yield chunk + + async def aclose(self) -> None: + self.closes += 1 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("isolate_litellm_state", "_router_httpx_transport") +@pytest.mark.parametrize("split_at", [0, 12, 24, 50], ids=["whole", "event-name", "event-end", "data"]) +@pytest.mark.parametrize("newline", [b"\n", b"\r\n"], ids=["lf", "crlf"]) +async def test_messages_stream_preserves_generated_content_across_partitions_before_a_retryable_error( + respx_mock: respx.MockRouter, split_at: int, newline: bytes +) -> None: + content: Final = _anthropic_messages_content_chunk("original").replace(b"\n", newline) + chunks: Final = ( + _anthropic_messages_message_start_chunk().replace(b"\n", newline), + *((content[:split_at], content[split_at:]) if split_at else (content,)), + _anthropic_messages_overloaded_error_chunk().replace(b"\n", newline), + ) + route: Final = respx_mock.post("https://messages-partitions.test/v1/messages").mock( + side_effect=lambda _: httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_PartitionedMessagesSSE(chunks) + ) + ) + router: Final = Router( + model_list=[{ + "model_name": "messages", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", "api_key": "fake-key", + "api_base": "https://messages-partitions.test", + }, + }], + num_retries=1, + retry_after=0, + disable_cooldowns=True, + ) + stream: Final = await router.aanthropic_messages( + model="messages", messages=[{"role": "user", "content": "hello"}], max_tokens=64, stream=True + ) + received: Final = b"".join([chunk async for chunk in stream]) + + assert received == b"".join(chunks) + assert route.call_count == 1 + + +@pytest.fixture +def _messages_stream_attempt() -> Iterator[StreamAttempt]: + from litellm.router_config import build_router_config + from litellm.router_python_host import python_host + from litellm.router_rust.core import RustCore + + core: Final = RustCore.build(build_router_config({"num_retries": 0}), python_host) + try: + yield StreamAttempt( + core=core, op="aanthropic_messages", request={}, kwargs={}, group="messages", + order=None, depth=0, level=0, deployment_id=None, attempted_retries=0, + ) + finally: + core.retire() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("read_first", [False, True], ids=["unread", "started"]) +async def test_messages_stream_close_releases_its_source_once_without_reading_more( + _messages_stream_attempt: StreamAttempt, read_first: bool +) -> None: + first: Final = _anthropic_messages_message_start_chunk() + source: Final = _PartitionedMessagesSSE((first, _anthropic_messages_content_chunk("unread"))) + + async def frames() -> AsyncGenerator[bytes, None]: + try: + async for chunk in source: + yield chunk + finally: + await source.aclose() + + stream: Final = FallbackAwareAnthropicMessagesStream(frames(), source, _messages_stream_attempt) + assert source.reads == source.closes == 0 + if read_first: + assert await anext(stream) == first + await stream.aclose() + await stream.aclose() + + assert source.reads == int(read_first) + assert source.closes == 1 + with pytest.raises(StopAsyncIteration): + await anext(stream) + assert source.reads == int(read_first) + assert source.closes == 1