scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

feat(messages): retry rejected thinking signatures on rust

Browse source
This commit is contained in:
narmstrong-berry 2026-11-10 20:43:24 +00:00
parent f0d234c4c6
commit 45dcc9cc80
No known key found for this signature in database
30 changed files with 976 additions and 96 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm-rust/crates/callbacks-legacy-python/python_contract.json
View file

@ -43,6 +43,11 @@
"headers",
"body"
],
"retried_request": [
"logger",
"logged_body",
"body"
],
"record_api_window": [
"logger",
"start",

83
litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
View file

@ -664,6 +664,38 @@ impl LegacyLogging {
Ok(HookStep::Ready(()))
}
pub(crate) fn request_retried(
&mut self,
py: Python<'_>,
wire: &WireRequest,
) -> PyResult<HookStep<Self, ()>> {
let Some(request) = &self.request else {
return Ok(HookStep::Ready(()));
};
let retried = Logging::RetriedRequest.call(
py,
(
self.logger()?.object(py),
request.body.bind(py),
to_py(py, &wire.body)?,
),
)?;
let now = epoch_seconds();
if self.api_started.is_some() {
self.close_api_window(py, now)?;
self.api_started = Some(now);
}
Ok(HookStep::Await(retried.unbind(), Self::resume_retried))
}
fn resume_retried(
&mut self,
_py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, ()>> {
result.map(|_| HookStep::Ready(()))
}
pub(crate) fn succeeded(
&mut self,
py: Python<'_>,
@ -1296,6 +1328,57 @@ assert ('custom_llm_provider' in logger.update['litellm_params']) is in_params,
json!({"type": "document_url", "document_url": source})
}
#[test]
fn a_retried_request_hands_the_logged_body_and_the_retried_body_to_logging() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
let mut logging = LegacyLogging {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
..legacy_call(py, &locals, false)
};
let context = RequestContext {
model: "model".into(),
custom_llm_provider: "provider".into(),
optional_params: json!({}),
secret_fields: Vec::new(),
api_key: None,
};
let wire = WireRequest {
url: "https://provider.invalid/v1/messages".into(),
headers: Vec::new(),
body: json!({"messages": [{"role": "user", "content": "hi"}], "thinking": {"type": "enabled"}}),
};
logging
.before_provider_request(py, Box::new(wire.clone()), &context)
.unwrap();
let retried = WireRequest {
body: json!({"messages": [{"role": "user", "content": "hi"}]}),
..wire
};
let step = logging
.on_event(
py,
PythonCallEvent::Execution(ExecutionEvent::ProviderRequestRetried {
wire: retried,
}),
)
.unwrap();
assert!(matches!(step, HookStep::Await(_, _)));
run(
py,
&locals,
c"
[(logged, body)] = [value for name, value in logger.calls if name == 'retried_request']
assert logged is logger.pre['complete_input_dict'], logged
assert body == {'messages': [{'role': 'user', 'content': 'hi'}]}, body
",
);
});
}
fn before_provider_request(script: &CStr, body: Value) -> WireRequest {
before_send_with_secrets(script, json!({}), body, &[])
}

3
litellm-rust/crates/callbacks-legacy-python/src/mapping.rs
View file

@ -179,6 +179,9 @@ impl CallHooks<PythonRuntime> for LegacyLogging {
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
(AFTER.invoke)(self, py, raw)
}
PythonCallEvent::Execution(ExecutionEvent::ProviderRequestRetried { wire }) => {
self.request_retried(py, &wire)
}
PythonCallEvent::Succeeded { timing, response } => {
(SUCCESS.invoke)(self, py, timing, response)
}

2
litellm-rust/crates/callbacks-legacy-python/src/python.rs
View file

@ -50,6 +50,8 @@ pub(crate) enum Logging {
RecordResponse,
#[strum(serialize = "record_api_window")]
RecordApiWindow,
#[strum(serialize = "retried_request")]
RetriedRequest,
#[strum(serialize = "defers_async_logging")]
DefersAsync,
#[strum(serialize = "claim_async_success_log")]

4
litellm-rust/crates/callbacks-legacy-python/src/test_support.rs
View file

@ -69,6 +69,10 @@ FAKES = {
'record_response': lambda logger, status, headers, body: logger.model_call_details.update(
response_headers=dict(headers)
),
'retried_request': lambda logger, logged_body, body: (
logger.record('retried_request', (logged_body, body)),
'awaitable',
)[1],
'record_api_window': lambda logger, start, end: logger.model_call_details.update(
llm_api_duration_ms=(end - start) * 1000
),

7
litellm-rust/crates/host-native/src/driver.rs
View file

@ -84,6 +84,13 @@ where
reply.send(self.interceptors.after_provider_response(raw).await);
Ok(())
}
HostRequest::Intercept(InterceptRequest::ProviderRequestRetried {
wire,
reply,
}) => {
reply.send(self.interceptors.provider_request_retried(*wire).await);
Ok(())
}
HostRequest::Stream(StreamDelivery::Open(head, reply)) => {
self.demand = Some(reply);
return Ok(Boundary::Open(head));

8
litellm-rust/crates/host-python/src/driver.rs
View file

@ -310,6 +310,14 @@ where
reply.send(Ok(()));
Ok(())
}
HostRequest::Intercept(InterceptRequest::ProviderRequestRetried { wire, reply }) => {
self.notify(PythonCallEvent::Execution(
ExecutionEvent::ProviderRequestRetried { wire: *wire },
))
.await?;
reply.send(Ok(()));
Ok(())
}
HostRequest::Stream(StreamDelivery::Open(head, reply)) => {
self.streaming = true;
let head = {

7
litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs
View file

@ -241,6 +241,9 @@ fn retain_event(py: Python<'_>, event: PythonCallEvent<'_>) -> OwnedEvent {
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw: raw.clone() })
}
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { wire }) => {
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { wire })
}
CallEvent::Succeeded { timing, response } => CallEvent::Succeeded {
timing,
response: response.clone_ref(py),
@ -280,6 +283,10 @@ fn dispatch(
py,
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
),
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { wire }) => hooks.on_event(
py,
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { wire: wire.clone() }),
),
CallEvent::Succeeded { timing, response } => hooks.on_event(
py,
CallEvent::Succeeded {

3
litellm-rust/crates/host-python/tests/support/mod.rs
View file

@ -422,6 +422,9 @@ impl CallHooks<PythonRuntime> for SyntheticHooks {
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
format!("response:{}", raw.body)
}
PythonCallEvent::Execution(ExecutionEvent::ProviderRequestRetried { wire }) => {
format!("retried:{}", wire.body)
}
PythonCallEvent::Succeeded { response, .. } => {
format!("succeeded:{}", response.bind(py))
}

14
litellm-rust/crates/host/src/interceptors.rs
View file

@ -107,6 +107,13 @@ pub trait Interceptors<E>: Send + Sync {
&self,
raw: RawResponse,
) -> impl Future<Output = Result<(), E>> + Send;
fn provider_request_retried(
&self,
_wire: WireRequest,
) -> impl Future<Output = Result<(), E>> + Send {
async { Ok(()) }
}
}
impl<E, T: Interceptors<E> + ?Sized> Interceptors<E> for &T {
@ -128,6 +135,13 @@ impl<E, T: Interceptors<E> + ?Sized> Interceptors<E> for &T {
) -> impl Future<Output = Result<(), E>> + Send {
(**self).after_provider_response(raw)
}
fn provider_request_retried(
&self,
wire: WireRequest,
) -> impl Future<Output = Result<(), E>> + Send {
(**self).provider_request_retried(wire)
}
}
impl<E> Interceptors<E> for () {

13
litellm-rust/crates/host/src/lifecycle.rs
View file

@ -6,7 +6,10 @@ use std::{
use futures_util::TryStreamExt;
use crate::{call::CallOutput, interceptors::RawResponse};
use crate::{
call::CallOutput,
interceptors::{RawResponse, WireRequest},
};
/// Seconds since the Unix epoch, on one clock for every host.
pub fn epoch_seconds() -> f64 {
@ -50,7 +53,7 @@ pub enum CallEvent<Response = (), Error = (), Raw = RawResponse> {
},
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq)]
pub enum ExecutionEvent<Raw = RawResponse> {
ResultReady {
facts: crate::interceptors::ExecutionFacts,
@ -58,6 +61,9 @@ pub enum ExecutionEvent<Raw = RawResponse> {
ProviderResponseReceived {
raw: Raw,
},
ProviderRequestRetried {
wire: WireRequest,
},
}
impl<Response, Error, Raw: std::borrow::Borrow<RawResponse>> CallEvent<Response, Error, Raw> {
@ -76,6 +82,9 @@ impl<Response, Error, Raw: std::borrow::Borrow<RawResponse>> CallEvent<Response,
facts: facts.clone(),
})
}
Self::Execution(ExecutionEvent::ProviderRequestRetried { wire }) => {
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { wire: wire.clone() })
}
Self::Succeeded { timing, .. } => CallEvent::Succeeded {
timing: *timing,
response: (),

11
litellm-rust/crates/host/src/machine/context.rs
View file

@ -119,6 +119,17 @@ where
})
.await?
}
async fn provider_request_retried(&self, wire: WireRequest) -> Result<(), P::Error> {
self.0
.request_reply(|reply| {
HostRequest::Intercept(InterceptRequest::ProviderRequestRetried {
wire: Box::new(wire),
reply,
})
})
.await?
}
}
pub struct StreamSender<P: Protocol>(Channel<P>);

4
litellm-rust/crates/host/src/protocol.rs
View file

@ -33,6 +33,10 @@ pub enum InterceptRequest<E> {
raw: RawResponse,
reply: Reply<Result<(), E>>,
},
ProviderRequestRetried {
wire: Box<WireRequest>,
reply: Reply<Result<(), E>>,
},
}
pub enum StreamDelivery<P: Protocol> {

2
litellm-rust/crates/inference-messages/src/constants.rs
View file

@ -2,3 +2,5 @@
/// seconds. Mirrors the Python Anthropic Messages default. The per-request
/// timeout from the caller still overrides this on the request builder.
pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600;
pub(crate) const MESSAGES_HTTP_ERROR_ATTEMPTS: usize = 2;

114
litellm-rust/crates/inference-messages/src/handler.rs
View file

@ -31,7 +31,7 @@ use serde_json::{Map, Value};
use super::{
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute,
common_utils::{MessagesProvider, merge_messages_header_sources, string_headers},
constants::MESSAGES_TIMEOUT_SECS,
constants::{MESSAGES_HTTP_ERROR_ATTEMPTS, MESSAGES_TIMEOUT_SECS},
route::Messages,
types::invalid_request,
};
@ -44,6 +44,7 @@ struct ProviderCall {
signer: Option<litellm_auth_aws::SigV4Signer>,
timeout: Option<Duration>,
stream: bool,
transform_context: MessagesTransformContext,
}
impl ProviderCall {
@ -217,6 +218,7 @@ async fn provider_call(
signer: authenticated.signer,
timeout: timeout.or_else(|| shaping.settings.request_timeout()),
stream,
transform_context,
})
}
@ -300,32 +302,17 @@ async fn call_provider(
request: ProviderCall,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
let timeout = request
.timeout
.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS));
let response = send_recovering(http, &request, timeout, context).await?;
let ProviderCall {
identity,
wire,
provider,
signer,
timeout,
stream,
..
} = request;
let provider_name = provider.as_str();
log_request_body(provider_name, stream, &wire.body);
let timeout = timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS));
let response = send(
http,
Authenticated {
headers: wire.headers,
signer,
},
&wire.url,
&wire.body,
timeout,
stream,
)
.await?;
if !response.status().is_success() {
return Err(provider_error(response, timeout).await);
}
let config = provider.config();
if stream {
return Ok(streaming_response(
@ -353,6 +340,63 @@ async fn call_provider(
decode_response(config, &identity.model, &text).map(MessagesCallResponse::Complete)
}
async fn send_recovering(
http: &litellm_http::Client,
request: &ProviderCall,
timeout: Duration,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<reqwest::Response, Error> {
let config = request.provider.config();
let mut body = std::borrow::Cow::Borrowed(&request.wire.body);
let mut attempt = 1;
loop {
log_request_body(request.provider.as_str(), request.stream, &body);
let response = send(
http,
Authenticated {
headers: request.wire.headers.clone(),
signer: request.signer.clone(),
},
&request.wire.url,
&body,
timeout,
request.stream,
)
.await?;
if response.status().is_success() {
return Ok(response);
}
let failure = provider_failure(response, timeout).await?;
if attempt == MESSAGES_HTTP_ERROR_ATTEMPTS
|| !config.should_retry_anthropic_messages_on_http_error(
failure.status,
&failure.text,
&body,
&request.transform_context,
)
{
return Err(failure.into());
}
attempt += 1;
tracing::debug!(
attempt,
max_attempts = MESSAGES_HTTP_ERROR_ATTEMPTS,
"retrying after a recoverable HTTP error"
);
let retried = WireRequest {
url: request.wire.url.clone(),
headers: request.wire.headers.clone(),
body: config.transform_anthropic_messages_request_on_http_error(
failure.status,
&failure.text,
body.into_owned(),
),
};
context.provider_request_retried(&retried).await?;
body = std::borrow::Cow::Owned(retried.body);
}
}
fn serialize_failure(err: serde_json::Error) -> Error {
Error::invalid_request(litellm_llms::ErrorDetail::failed(
"Anthropic messages request serialization",
@ -402,18 +446,38 @@ fn timed_chunks(
)
}
struct ProviderFailure {
status: u16,
text: String,
headers: Vec<(String, String)>,
}
impl From<ProviderFailure> for Error {
fn from(failure: ProviderFailure) -> Self {
litellm_http::Error::upstream_with_headers(failure.status, failure.text, failure.headers)
.into()
}
}
/// The provider's error, its body read under `timeout`: a stream has no overall deadline once
/// its head arrives, so an error body that stalls would otherwise hold the call forever.
async fn provider_error(response: reqwest::Response, timeout: Duration) -> Error {
async fn provider_failure(
response: reqwest::Response,
timeout: Duration,
) -> Result<ProviderFailure, Error> {
let status = response.status().as_u16();
let headers = litellm_http::request::header_pairs(response.headers());
match tokio::time::timeout(timeout, response.text()).await {
Ok(Ok(text)) => {
log_error_body(status, &text);
litellm_http::Error::upstream_with_headers(status, text, headers).into()
Ok(ProviderFailure {
status,
text,
headers,
})
}
Ok(Err(error)) => litellm_http::Error::from(error).into(),
Err(elapsed) => litellm_http::Error::read_timeout(elapsed).into(),
Ok(Err(error)) => Err(litellm_http::Error::from(error).into()),
Err(elapsed) => Err(litellm_http::Error::read_timeout(elapsed).into()),
}
}

1
litellm-rust/crates/inference-messages/tests/messages/main.rs
View file

@ -22,6 +22,7 @@ use support::*;
mod host;
mod request;
mod response;
mod retry;
mod secrets;
mod stream;
mod tencent;

209
litellm-rust/crates/inference-messages/tests/messages/retry.rs Normal file
View file

@ -0,0 +1,209 @@
use futures_util::TryStreamExt;
use litellm_host::lifecycle::{CallEvent, ExecutionEvent};
use litellm_inference_messages::MessagesCallResponse;
use litellm_inference_testing::RecordingSecrets;
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities;
use rstest::rstest;
use super::*;
const INVALID_SIGNATURE: &str = "messages.1.content.0: Invalid `signature` in `thinking` block";
const BOUND_ELSEWHERE: &str = "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.";
fn rejection(status: u16, message: &str) -> ResponseTemplate {
status_response(
status,
json!({"type": "error", "error": {"type": "invalid_request_error", "message": message}}),
)
}
fn sse_response() -> ResponseTemplate {
ResponseTemplate::new(200).set_body_raw(
"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
"text/event-stream",
)
}
fn replay(call: MessagesCall, api_base: String, stream: bool, thinking: Value) -> MessagesCall {
MessagesCall {
api_key: Some("sk-ant".into()),
api_base: Some(api_base),
..with_fields(
call,
json!({
"max_tokens": 4096,
"stream": stream,
"thinking": thinking,
"messages": [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": [
{"type": "thinking", "thinking": "plan", "signature": "sig"},
{"type": "text", "text": "hello"}
]},
{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]},
{"role": "user", "content": "again"}
]
}),
)
}
}
fn enabled() -> Value {
json!({"type": "enabled", "budget_tokens": 2048})
}
async fn execute(
host: &RecordingCall<Messages>,
) -> Result<(MessagesCallResponse, Vec<bytes::Bytes>), Error> {
let route = messages_route(Arc::new(RecordingSecrets::empty()));
let response = route
.execute(host.request()?, host, Some(host.events.0.sender.clone()))
.await?;
match response {
MessagesCallResponse::Stream { head, chunks } => Ok((
MessagesCallResponse::Stream {
head,
chunks: Box::pin(futures_util::stream::empty()),
},
chunks.try_collect().await?,
)),
complete => Ok((complete, Vec::new())),
}
}
fn observed_retries(host: &RecordingCall<Messages>) -> Vec<Value> {
host.events
.0
.lock()
.unwrap()
.iter()
.filter_map(|event| match event {
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { wire }) => {
Some(wire.body.clone())
}
_ => None,
})
.collect()
}
#[rstest]
#[case::invalid_signature(INVALID_SIGNATURE)]
#[case::missing_signature(
"messages.1.content.0.thinking.signature.str: Input should be a valid string"
)]
#[case::empty_thinking("messages.1.content.0.thinking: each thinking block must contain thinking")]
#[case::unpinned_binding_mismatch(BOUND_ELSEWHERE)]
#[tokio::test]
async fn a_rejected_thinking_replay_is_sent_once_more_without_thinking(
call: MessagesCall,
#[case] message: &str,
#[values(false, true)] stream: bool,
) {
let success = if stream {
sse_response()
} else {
message_response()
};
let upstream = upstream([rejection(400, message), success]).await;
let host = RecordingCall::<Messages>::new(replay(call, upstream.uri(), stream, enabled()));
let (response, chunks) = execute(&host).await.expect("the retry succeeds");
let [first, second] = <[_; 2]>::try_from(received(&upstream).await).unwrap();
assert_eq!(first.json()["thinking"], enabled());
assert_eq!(
second.json()["messages"],
json!([
{"role": "user", "content": "hi"},
{"role": "assistant", "content": [{"type": "text", "text": "hello"}]},
{"role": "user", "content": "again"}
])
);
let mut expected = first.json();
expected.as_object_mut().unwrap().shift_remove("thinking");
expected["messages"] = second.json()["messages"].clone();
assert_eq!(second.json(), expected);
assert_eq!(first.header("x-hook"), second.header("x-hook"));
let retried: Vec<Value> = host
.retried
.lock()
.unwrap()
.iter()
.map(|wire| wire.body.clone())
.collect();
assert_eq!(retried, vec![second.json()]);
assert_eq!(observed_retries(&host), vec![second.json()]);
assert_eq!(
matches!(response, MessagesCallResponse::Stream { .. }),
stream
);
assert_eq!(chunks.is_empty(), !stream);
}
#[rstest]
#[tokio::test]
async fn a_second_rejection_is_the_error_the_caller_sees(
call: MessagesCall,
#[values(false, true)] stream: bool,
) {
let upstream = upstream([
rejection(400, INVALID_SIGNATURE),
rejection(
400,
"messages.0.content.0: Invalid `signature` in `thinking` block, again",
),
])
.await;
let host = RecordingCall::<Messages>::new(replay(call, upstream.uri(), stream, enabled()));
let error = execute(&host).await.err().expect("the retry fails");
assert_eq!(received(&upstream).await.len(), 2);
let (status, body) = error.upstream_response().unwrap();
assert_eq!(status, 400);
assert!(body.contains("again"), "{body}");
assert_eq!(host.retried.lock().unwrap().len(), 1);
}
#[rstest]
#[case::not_a_400(429, INVALID_SIGNATURE, json!({"type": "adaptive"}))]
#[case::unrelated_400(400, "max_tokens: 4096 > 64000", json!({"type": "adaptive"}))]
#[case::signature_without_thinking(400, "Invalid `signature` in `tool_use` block", json!({"type": "adaptive"}))]
#[case::caller_pinned_binding(
400,
BOUND_ELSEWHERE,
json!({"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}})
)]
#[tokio::test]
async fn other_failures_reach_the_caller_without_a_retry(
call: MessagesCall,
#[case] status: u16,
#[case] message: &str,
#[case] thinking: Value,
) {
let upstream = upstream([rejection(status, message), message_response()]).await;
let adaptive = MessagesShaping {
capabilities: MessagesModelCapabilities {
supports_adaptive_thinking: true,
..MessagesModelCapabilities::default()
},
..MessagesShaping::default()
};
let host = RecordingCall::<Messages>::new(replay(
MessagesCall {
shaping: adaptive,
..call
},
upstream.uri(),
false,
thinking.clone(),
));
let error = execute(&host).await.err().expect("the call fails");
let sent = only_request(&upstream).await;
assert_eq!(sent.json()["thinking"], thinking);
assert_eq!(error.upstream_response().unwrap().0, status);
assert!(host.retried.lock().unwrap().is_empty());
assert!(observed_retries(&host).is_empty());
}

10
litellm-rust/crates/inference-messages/tests/support/mod.rs
View file

@ -117,6 +117,7 @@ pub struct RecordingCall<P: litellm_host::protocol::Protocol> {
pub events: Arc<CallEvents>,
pub chunks: Mutex<Vec<P::Chunk>>,
pub head: Mutex<Option<P::StreamHead>>,
pub retried: Mutex<Vec<litellm_host::interceptors::WireRequest>>,
}
#[derive(Default)]
@ -167,6 +168,7 @@ impl<P: litellm_host::protocol::Protocol> RecordingCall<P> {
events: Arc::new(CallEvents::default()),
chunks: Mutex::new(Vec::new()),
head: Mutex::new(None),
retried: Mutex::new(Vec::new()),
}
}
}
@ -195,6 +197,14 @@ impl<P: litellm_host::protocol::Protocol> litellm_host::interceptors::Intercepto
) -> Result<(), P::Error> {
Ok(())
}
async fn provider_request_retried(
&self,
wire: litellm_host::interceptors::WireRequest,
) -> Result<(), P::Error> {
self.retried.lock().unwrap().push(wire);
Ok(())
}
}
impl<P> RecordingCall<P>

13
litellm-rust/crates/inference/src/context.rs
View file

@ -1,6 +1,6 @@
use litellm_cache_response::CachePolicy;
use litellm_host::{
interceptors::{ExecutionFacts, Interceptors, RawResponse, ResultSource},
interceptors::{ExecutionFacts, Interceptors, RawResponse, ResultSource, WireRequest},
lifecycle::{CallEvent, ExecutionEvent},
observation::ObservationSender,
};
@ -57,4 +57,15 @@ impl<'a, I: Interceptors<Error>> CallContext<'a, I> {
.await
.map_err(Error::post_call)
}
pub async fn provider_request_retried(&self, wire: &WireRequest) -> Result<(), Error> {
if let Some(observers) = &self.observers {
observers.emit(CallEvent::Execution(
ExecutionEvent::ProviderRequestRetried { wire: wire.clone() },
));
}
self.interceptors
.provider_request_retried(wire.clone())
.await
}
}

162
litellm-rust/crates/llms/src/anthropic/common_utils.rs
View file

@ -650,6 +650,62 @@ pub fn flatten_unencrypted_web_search_results(messages: Vec<Message>) -> Vec<Mes
.collect()
}
pub fn is_anthropic_invalid_thinking_block_error(error_text: &str) -> bool {
let lower = error_text.to_lowercase();
if !lower.contains("thinking") {
return false;
}
(lower.contains("signature") && (lower.contains("invalid") || lower.contains("valid string")))
|| lower.contains("must contain thinking")
}
pub fn is_anthropic_thinking_binding_mismatch_error(error_text: &str) -> bool {
error_text
.to_lowercase()
.contains("bound to a different conversation")
}
pub fn anthropic_request_errors_on_thinking_binding_mismatch(request_body: &Value) -> bool {
request_body
.pointer("/thinking/block_binding/prefix_mismatch_behavior")
.and_then(Value::as_str)
== Some("error")
}
fn is_thinking_value_block(block: &Value) -> bool {
matches!(
block.get("type").and_then(Value::as_str),
Some("thinking" | "redacted_thinking")
)
}
fn without_thinking_value_blocks(message: Value) -> Option<Value> {
let Value::Object(mut message) = message else {
return Some(message);
};
if let Some(Value::Array(content)) = message.get_mut("content") {
content.retain(|block| !is_thinking_value_block(block));
if content.is_empty() {
return None;
}
}
Some(Value::Object(message))
}
pub fn strip_thinking_blocks_from_anthropic_messages_request(request_body: Value) -> Value {
let Value::Object(mut body) = request_body else {
return request_body;
};
if let Some(Value::Array(messages)) = body.get_mut("messages") {
*messages = std::mem::take(messages)
.into_iter()
.filter_map(without_thinking_value_blocks)
.collect();
}
body.shift_remove("thinking");
Value::Object(body)
}
#[cfg(test)]
mod tests {
use crate::base_llm::messages::context::SupportedEffortTiers;
@ -1974,4 +2030,110 @@ mod tests {
fn the_messages_url_follows_python_ensure_api_route(#[case] api_base: &str, #[case] url: &str) {
assert_eq!(complete_anthropic_url(Some(api_base), &|_| None), url);
}
#[rstest]
#[case::invalid_signature(
"messages.1.content.0: Invalid `signature` in `thinking` block",
true
)]
#[case::signature_not_a_string(
r#"{"message":"messages.2.content.0.thinking.signature.str: Input should be a valid string"}"#,
true
)]
#[case::empty_thinking(
"messages.1.content.0.thinking: each thinking block must contain thinking",
true
)]
#[case::uppercase("MESSAGES.1: INVALID SIGNATURE IN THINKING BLOCK", true)]
#[case::signature_without_a_verdict("thinking signature mismatch", false)]
#[case::invalid_without_thinking("Invalid `signature` in `tool_use` block", false)]
#[case::unrelated("max_tokens: 4096 > 64000", false)]
#[case::empty("", false)]
fn invalid_thinking_block_errors_are_recognized(
#[case] error_text: &str,
#[case] expected: bool,
) {
assert_eq!(
is_anthropic_invalid_thinking_block_error(error_text),
expected
);
}
#[rstest]
#[case::bound_elsewhere("The block is Bound To A Different Conversation.", true)]
#[case::plain_signature("Invalid `signature` in `thinking` block", false)]
fn binding_mismatch_errors_are_recognized(#[case] error_text: &str, #[case] expected: bool) {
assert_eq!(
is_anthropic_thinking_binding_mismatch_error(error_text),
expected
);
}
#[rstest]
#[case::pinned(json!({"thinking": {"block_binding": {"prefix_mismatch_behavior": "error"}}}), true)]
#[case::other_behavior(json!({"thinking": {"block_binding": {"prefix_mismatch_behavior": "drop"}}}), false)]
#[case::string_binding(json!({"thinking": {"block_binding": "strict"}}), false)]
#[case::no_thinking(json!({"messages": []}), false)]
#[case::not_an_object(json!(["thinking"]), false)]
fn a_pinned_binding_is_the_caller_asking_for_the_error(
#[case] body: Value,
#[case] expected: bool,
) {
assert_eq!(
anthropic_request_errors_on_thinking_binding_mismatch(&body),
expected
);
}
#[rstest]
#[case::mixed_turn(
json!({"model": "m", "thinking": {"type": "enabled"}, "messages": [
{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "s"},
{"type": "text", "text": "a"},
{"type": "redacted_thinking", "data": "d"}
]}
]}),
json!({"model": "m", "messages": [{"role": "assistant", "content": [{"type": "text", "text": "a"}]}]})
)]
#[case::thinking_only_and_empty_turns_are_dropped(
json!({"messages": [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "d"}]},
{"role": "assistant", "content": []},
{"role": "user", "content": "again"}
]}),
json!({"messages": [{"role": "user", "content": "hi"}, {"role": "user", "content": "again"}]})
)]
#[case::non_object_entries_are_kept(
json!({"messages": ["raw", {"role": "user", "content": ["raw", {"type": "thinking"}]}]}),
json!({"messages": ["raw", {"role": "user", "content": ["raw"]}]})
)]
#[case::messages_not_a_list(
json!({"messages": "hi", "thinking": {"type": "adaptive"}}),
json!({"messages": "hi"})
)]
fn retrying_strips_thinking_from_the_replay_and_the_request(
#[case] body: Value,
#[case] expected: Value,
) {
assert_eq!(
strip_thinking_blocks_from_anthropic_messages_request(body),
expected
);
}
#[test]
fn stripping_keeps_the_order_of_the_remaining_fields() {
let stripped = strip_thinking_blocks_from_anthropic_messages_request(
json!({"model": "m", "thinking": {}, "messages": [], "max_tokens": 1}),
);
let keys: Vec<&str> = stripped
.as_object()
.unwrap()
.keys()
.map(String::as_str)
.collect();
assert_eq!(keys, ["model", "messages", "max_tokens"]);
}
}

113
litellm-rust/crates/llms/src/anthropic/messages/transformation.rs
View file

@ -16,10 +16,13 @@ use crate::{
Error,
anthropic::common_utils::{
ANTHROPIC_API_BASE_ENV, ANTHROPIC_API_KEY_ENV, ANTHROPIC_AUTH_TOKEN_ENV,
ANTHROPIC_BASE_URL_ENV, OauthHandling, complete_anthropic_url, get_auth_header,
has_advisor_tool, has_anthropic_credential, is_tool_search_used, merge_beta_headers,
optionally_handle_anthropic_oauth, requires_native_compaction_beta, strip_advisor_blocks,
strip_encrypted_reasoning_blocks,
ANTHROPIC_BASE_URL_ENV, OauthHandling,
anthropic_request_errors_on_thinking_binding_mismatch, complete_anthropic_url,
get_auth_header, has_advisor_tool, has_anthropic_credential,
is_anthropic_invalid_thinking_block_error, is_anthropic_thinking_binding_mismatch_error,
is_tool_search_used, merge_beta_headers, optionally_handle_anthropic_oauth,
requires_native_compaction_beta, strip_advisor_blocks, strip_encrypted_reasoning_blocks,
strip_thinking_blocks_from_anthropic_messages_request,
},
base_llm::{
auth::AuthScheme,
@ -116,6 +119,34 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
provider_feature_betas(request, BetaProvider::Anthropic),
)
}
fn should_retry_anthropic_messages_on_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
_context: &MessagesTransformContext,
) -> bool {
if status != 400 {
return false;
}
let caller_asked_for_this_error =
anthropic_request_errors_on_thinking_binding_mismatch(request_body)
&& is_anthropic_thinking_binding_mismatch_error(error_text);
!caller_asked_for_this_error && is_anthropic_invalid_thinking_block_error(error_text)
}
fn transform_anthropic_messages_request_on_http_error(
&self,
status: u16,
error_text: &str,
request_body: Value,
) -> Value {
if status == 400 && is_anthropic_invalid_thinking_block_error(error_text) {
return strip_thinking_blocks_from_anthropic_messages_request(request_body);
}
request_body
}
}
pub(crate) fn transform_messages_request(
@ -1352,4 +1383,78 @@ mod tests {
.collect();
assert_eq!(undeclared, Vec::<&String>::new());
}
const INVALID_SIGNATURE: &str = "messages.1.content.0: Invalid `signature` in `thinking` block";
const BOUND_ELSEWHERE: &str = "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.";
fn replayed(thinking: Value) -> Value {
json!({
"model": "claude-sonnet-4-5",
"thinking": thinking,
"messages": [{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "s"},
{"type": "text", "text": "a"}
]}]
})
}
#[rstest]
#[case::invalid_signature(400, INVALID_SIGNATURE, json!({"type": "enabled"}), true)]
#[case::unpinned_binding(400, BOUND_ELSEWHERE, json!({"type": "adaptive"}), true)]
#[case::pinned_binding(
400,
BOUND_ELSEWHERE,
json!({"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}),
false
)]
#[case::pinned_binding_other_signature_error(
400,
INVALID_SIGNATURE,
json!({"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}),
true
)]
#[case::not_a_400(429, INVALID_SIGNATURE, json!({"type": "enabled"}), false)]
#[case::unrelated(400, "max_tokens: 4096 > 64000", json!({"type": "enabled"}), false)]
fn only_a_rejected_thinking_replay_is_retried(
#[case] status: u16,
#[case] error_text: &str,
#[case] thinking: Value,
#[case] expected: bool,
) {
assert_eq!(
ANTHROPIC_MESSAGES_CONFIG.should_retry_anthropic_messages_on_http_error(
status,
error_text,
&replayed(thinking),
&MessagesTransformContext::default(),
),
expected
);
}
#[rstest]
#[case::invalid_signature(400, INVALID_SIGNATURE, true)]
#[case::not_a_400(429, INVALID_SIGNATURE, false)]
#[case::unrelated(400, "max_tokens: 4096 > 64000", false)]
fn the_retry_drops_thinking_only_for_a_rejected_thinking_replay(
#[case] status: u16,
#[case] error_text: &str,
#[case] stripped: bool,
) {
let body = replayed(json!({"type": "enabled"}));
let retried = ANTHROPIC_MESSAGES_CONFIG.transform_anthropic_messages_request_on_http_error(
status,
error_text,
body.clone(),
);
let expected = if stripped {
json!({
"model": "claude-sonnet-4-5",
"messages": [{"role": "assistant", "content": [{"type": "text", "text": "a"}]}]
})
} else {
body
};
assert_eq!(retried, expected);
}
}

19
litellm-rust/crates/llms/src/base_llm/messages/transformation.rs
View file

@ -91,6 +91,25 @@ pub trait BaseMessagesConfig: Sync {
fn wire_body(&self, body: Value) -> Value {
body
}
fn should_retry_anthropic_messages_on_http_error(
&self,
_status: u16,
_error_text: &str,
_request_body: &Value,
_context: &MessagesTransformContext,
) -> bool {
false
}
fn transform_anthropic_messages_request_on_http_error(
&self,
_status: u16,
_error_text: &str,
request_body: Value,
) -> Value {
request_body
}
}
#[cfg(test)]

1
litellm/messages/dispatch.py
View file

@ -143,7 +143,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType(
Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity,
Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection,
Feature.STREAM_COST_INJECTION: route_detectors.stream_cost_injection,
Feature.THINKING_SIGNATURE_RETRY: route_detectors.thinking_signature_retry,
Feature.TOOL_CUSTOM_FIELD: route_detectors.tool_custom_field,
}
)

2
litellm/rust_bridge/catalog.py
View file

@ -96,7 +96,6 @@ class Feature(str, Enum):
"""Prompt caching injection that Python applies to messages and system."""
STREAM_COST_INJECTION = "stream_cost_injection"
"""litellm.include_cost_in_streaming_usage or the x-litellm-include-cost-in-streaming-usage header, where Python adds cost to streamed usage."""
THINKING_SIGNATURE_RETRY = "thinking_signature_retry"
"""Replayed thinking blocks, which Python strips and retries once if Anthropic rejects their signature."""
TOOL_CUSTOM_FIELD = "tool_custom_field"
"""Tools carrying a non-spec custom object, which Python drops or hoists for the provider."""
@ -184,7 +183,6 @@ MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VAR
Feature.WORKLOAD_IDENTITY,
Feature.CACHE_CONTROL_INJECTION,
Feature.STREAM_COST_INJECTION,
Feature.THINKING_SIGNATURE_RETRY,
Feature.TOOL_CUSTOM_FIELD,
Feature.STRICT_TOOLS_RETRY,
}

13
litellm/rust_bridge/host/callbacks_legacy_python.py
View file

@ -285,6 +285,19 @@ def record_response(logger: Logging, status: int, headers: Sequence[tuple[str, s
details["response_headers"] = dict(response.headers)
async def retried_request(
logger: Logging,
logged_body: dict[str, object], # mutable-ok: the retry rewrites the body pre_call logged, as Python's loop does
body: Mapping[str, object],
) -> None:
if logger.baseline_cache_context is not None:
await logger.invalidate_baseline_cache_estimate("retried_request")
for name in logged_body.keys() - body.keys():
del logged_body[name] # rebind-ok: the logged body is the one logging already holds
logged_body.update({name: value for name, value in body.items() if logged_body.get(name) != value})
logger.model_call_details.update(logged_body)
def record_api_window(logger: Logging, start: float, end: float) -> None:
"""Records the provider exchange's window the way Python's timed `AsyncHTTPHandler.post` does."""
from litellm.litellm_core_utils.logging_utils import (

13
litellm/rust_bridge/messages/detectors.py
View file

@ -42,19 +42,6 @@ def anthropic_workload_identity(request: NativeCall) -> bool:
)
_THINKING_BLOCKS: Final = frozenset({"thinking", "redacted_thinking"})
def _replays_thinking(message: object) -> bool:
content: Final = optional_sequence((optional_mapping(message) or {}).get("content")) or ()
return any((optional_mapping(block) or {}).get("type") in _THINKING_BLOCKS for block in content)
def thinking_signature_retry(request: NativeCall) -> bool:
"""Python may strip replayed thinking blocks and retry after a rejected signature."""
return any(_replays_thinking(message) for message in optional_sequence(request.resolved.get("messages")) or ())
def tool_custom_field(request: NativeCall) -> bool:
tools: Final = optional_sequence(request.resolved.get("tools")) or ()
return any("custom" in (optional_mapping(tool) or {}) for tool in tools)

124
tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py
View file

@ -305,6 +305,7 @@ CALLBACK_ASPECTS: Final = (
"logged_response_headers",
"logged_api_key",
"logged_call",
"logged_request",
)
ASPECTS: Final = ("dispatch", "request", "outcome", "hidden_params", "logged", *CALLBACK_ASPECTS)
@ -343,9 +344,6 @@ S15_COMPLETE_DEADLINE: Final = (
"so a body trickling in slower than the timeout succeeds in Python and times out natively"
)
F_MOCK: Final = "F: mock_response is Python-only (detector mock_response)"
F_THINKING_RETRY: Final = (
"F: retry without thinking after a rejected signature is Python-only (detector thinking_signature_retry)"
)
F_CACHE_INJECTION: Final = "F: prompt cache_control injection is Python-only (detector cache_control_injection)"
F_TOOL_CUSTOM: Final = "F: hoisting a tool's custom field is Python-only (detector tool_custom_field)"
F_STREAM_COST: Final = "F: cost injection into streamed usage is Python-only (detector stream_cost_injection)"
@ -363,6 +361,7 @@ KNOWN_GAPS: Final[Mapping[str, Mapping[str, str | None]]] = MappingProxyType(
logged_metadata=F_MOCK,
logged_api_key=F_MOCK,
logged_call=F_MOCK,
logged_request=F_MOCK,
),
"mock_response_rate_limit": _gaps(
logged_response_headers=F_MOCK,
@ -371,19 +370,12 @@ KNOWN_GAPS: Final[Mapping[str, Mapping[str, str | None]]] = MappingProxyType(
logged=F_MOCK,
logged_api_key=F_MOCK,
logged_call=F_MOCK,
logged_request=F_MOCK,
),
"thinking_signature_rejected_then_retried": _gaps(
request=F_THINKING_RETRY,
outcome=F_THINKING_RETRY,
hidden_params=F_THINKING_RETRY,
logged=F_THINKING_RETRY,
logged_metadata=F_THINKING_RETRY,
logged_response_headers=F_THINKING_RETRY,
),
"cache_control_injection_points": _gaps(request=F_CACHE_INJECTION),
"cache_control_injection_points": _gaps(request=F_CACHE_INJECTION, logged_request=F_CACHE_INJECTION),
"tool_custom_field": _gaps(request=F_TOOL_CUSTOM),
"stream_cost_injection": _gaps(outcome=F_STREAM_COST),
"enable_prompt_caching": _gaps(request=F_CACHE_INJECTION),
"enable_prompt_caching": _gaps(request=F_CACHE_INJECTION, logged_request=F_CACHE_INJECTION),
}
)
@ -685,6 +677,7 @@ FAILURE_CASES: Final = (
_kw(timeout=1),
upstream=(ResponseSpec(body=dict(TEXT_RESPONSE), chunk_delay=0.5, body_chunks=4),),
gaps=_gaps(
logged_request=S15_COMPLETE_DEADLINE,
outcome=S15_COMPLETE_DEADLINE,
hidden_params=S15_COMPLETE_DEADLINE,
logged=S15_COMPLETE_DEADLINE,
@ -732,18 +725,6 @@ SYNC_CASES: Final = (
FALLBACK_CASES: Final = (
Case("mock_response", _kw(mock_response="Hi from the mock"), upstream=()),
Case("mock_response_rate_limit", _kw(mock_response="litellm.RateLimitError"), upstream=()),
Case(
"thinking_replay",
_kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096),
),
Case(
"thinking_signature_rejected_then_retried",
_kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096),
upstream=(
_failure(400, "invalid_request_error", "messages.1.content.0: Invalid `signature` in `thinking` block"),
_ok(),
),
),
Case(
"cache_control_injection_points",
_kw(system="Answer in one word.", cache_control_injection_points=[{"location": "message", "role": "system"}]),
@ -769,7 +750,93 @@ 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)
INVALID_SIGNATURE: Final = "messages.1.content.0: Invalid `signature` in `thinking` block"
BOUND_ELSEWHERE: Final = f"{INVALID_SIGNATURE}. The block is bound to a different conversation."
THINKING: Final[Mapping[str, JsonValue]] = MappingProxyType({"type": "enabled", "budget_tokens": 2048})
BINDING_ERRORS: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}
)
THINKING_ONLY_REPLAY: Final[tuple[JsonValue, ...]] = (
{"role": "user", "content": "What is the capital of France?"},
{
"role": "assistant",
"content": [
{"type": "redacted_thinking", "data": "EmwKAhgBEgy"},
{"type": "thinking", "thinking": "Paris.", "signature": "EqQBCkYIARgCKkB"},
],
},
{"role": "assistant", "content": []},
{"role": "user", "content": "And Germany?"},
)
def _replay(**kwargs: object) -> Mapping[str, object]:
return _kw(messages=list(THINKING_REPLAY), thinking=dict(THINKING), max_tokens=4096, **kwargs)
def _rejected(message: str = INVALID_SIGNATURE, status: int = 400) -> ResponseSpec:
return _failure(status, "invalid_request_error", message)
THINKING_RETRY_CASES: Final = (
Case("thinking_replay", _replay()),
Case("thinking_signature_rejected_then_retried", _replay(), upstream=(_rejected(), _ok())),
Case(
"thinking_signature_rejected_then_retried_stream",
_replay(stream=True),
upstream=(_rejected(), _stream(TEXT_STREAM)),
),
Case(
"thinking_signature_missing_then_retried",
_replay(),
upstream=(_rejected("messages.1.content.0.thinking.signature.str: Input should be a valid string"), _ok()),
),
Case(
"empty_thinking_rejected_then_retried_without_thinking_only_turns",
_kw(messages=list(THINKING_ONLY_REPLAY), thinking=dict(THINKING), max_tokens=4096),
upstream=(_rejected("messages.1.content.1.thinking: each thinking block must contain thinking"), _ok()),
),
Case("thinking_signature_rejected_twice", _replay(), upstream=(_rejected(), _rejected())),
Case(
"thinking_signature_rejected_twice_stream",
_replay(stream=True),
upstream=(_rejected(), _rejected()),
),
Case("thinking_signature_text_on_a_429_is_not_retried", _replay(), upstream=(_rejected(status=429),)),
Case(
"unrelated_400_is_not_retried",
_replay(),
upstream=(_rejected("max_tokens: 4096 > 64000, which is the maximum allowed"),),
),
Case(
"binding_mismatch_retried_when_the_caller_did_not_pin_it",
_kw(messages=list(THINKING_REPLAY), thinking={"type": "adaptive"}, max_tokens=4096),
model=OPUS_ADAPTIVE,
upstream=(_rejected(BOUND_ELSEWHERE), _ok()),
),
Case(
"binding_mismatch_raised_when_the_caller_pinned_it",
_kw(messages=list(THINKING_REPLAY), thinking=dict(BINDING_ERRORS), max_tokens=4096),
model=OPUS_ADAPTIVE,
upstream=(_rejected(BOUND_ELSEWHERE),),
),
Case(
"invalid_signature_retried_even_when_the_caller_pinned_binding",
_kw(messages=list(THINKING_REPLAY), thinking=dict(BINDING_ERRORS), max_tokens=4096),
model=OPUS_ADAPTIVE,
upstream=(_rejected(), _ok()),
),
)
CASES: Final = (
*REQUEST_CASES,
*RESPONSE_CASES,
*FAILURE_CASES,
*STREAM_CASES,
*SYNC_CASES,
*THINKING_RETRY_CASES,
*FALLBACK_CASES,
)
@dataclass(frozen=True, slots=True)
@ -996,6 +1063,11 @@ def _logged(recorder: RecordingLogger, events: tuple[str, ...]) -> tuple[Mapping
"model": details.get("model"),
"stream": details.get("stream"),
},
"logged_request": {
"messages": details.get("messages"),
"thinking": details.get("thinking"),
"complete_input_dict": _mapping(details.get("additional_args")).get("complete_input_dict"),
},
}
return logged, callback_kwargs

48
tests/test_litellm_rust/messages/test_dispatch.py
View file

@ -18,7 +18,7 @@ from litellm.messages.dispatch import (
)
from litellm.proxy import proxy_server
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Feature, Route, RouteRule, Rules, Target
from litellm.rust_bridge.catalog import RULES, Feature, Route, RouteRule, Rules, Target
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES, NativeAmessages, NativeMessages
from litellm.rust_bridge.public_call import NativeCall
@ -486,12 +486,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict":
{},
),
(Feature.CACHE_CONTROL_INJECTION, "anthropic/claude-sonnet-4-5", {"enable_prompt_caching": True}, {}),
(
Feature.THINKING_SIGNATURE_RETRY,
"anthropic/claude-sonnet-4-5",
{"messages": [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "x"}]}]},
{},
),
(
Feature.TOOL_CUSTOM_FIELD,
"anthropic/claude-sonnet-4-5",
@ -578,6 +572,46 @@ def test_proxy_request_shape_stays_native(kwargs: Mapping[str, object]) -> None:
assert served is native_result
@pytest.mark.asyncio
@pytest.mark.parametrize(
"block",
(
pytest.param({"type": "thinking", "thinking": "x", "signature": "s"}, id="thinking"),
pytest.param({"type": "redacted_thinking", "data": "x"}, id="redacted-thinking"),
),
)
async def test_replayed_thinking_is_served_natively_under_the_shipped_rules(
monkeypatch: pytest.MonkeyPatch, block: Mapping[str, object]
) -> None:
monkeypatch.setenv("LITELLM_RUST", "1")
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "cache", None)
replay: Final = [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": [block, {"type": "text", "text": "hello"}]},
{"role": "user", "content": "again"},
]
native_result: Final = response()
python_result: Final = response()
async def python(*call_args: object, **python_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub
return python_result
async def native(request: NativeCall) -> AnthropicMessagesResponse:
return native_result
served: Final = await _ADISPATCH.arun(
(16, replay, "anthropic/claude-sonnet-4-5"),
{"api_key": "sk-ant-test"},
python=python,
binding=amessages_binding(native),
native=lambda hook, request, call_args, native_kwargs: hook(request),
rules=RULES,
)
assert served is native_result
@pytest.mark.asyncio
@pytest.mark.parametrize(
"kwargs",

50
tests/unit/rust_bridge/host/test_callbacks_legacy_python.py
View file

@ -4,7 +4,7 @@ import inspect
from collections.abc import Mapping
from pathlib import Path
from types import MappingProxyType
from typing import Final
from typing import Final, cast # noqa: TID251 # a duck-typed logger double stands in for Logging
import httpx
import pytest
@ -323,6 +323,54 @@ def test_native_provider_exchange_window_is_recorded_like_the_timed_post() -> No
assert metadata["llm_api_timing_windows"] == ((1_700_000_000.0, 1_700_000_000.25),)
@pytest.mark.asyncio
async def test_a_retried_request_rewrites_the_logged_body_in_place_like_the_python_retry_loop() -> None:
logger: Final = _messages_logger()
metadata: Final = {"user_id": "u-1"}
replayed: Final = [
{"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "s"}]},
{"role": "user", "content": "again"},
]
logged: Final[dict[str, object]] = {
"model": "claude-sonnet-4-5",
"messages": replayed,
"metadata": metadata,
"thinking": {"type": "enabled", "budget_tokens": 1024},
}
logger.model_call_details.update(logged)
retried: Final = {
"model": "claude-sonnet-4-5",
"messages": [{"role": "user", "content": "again"}],
"metadata": {"user_id": "u-1"},
}
await legacy.retried_request(logger, logged, retried)
assert logged == retried
assert logged["metadata"] is metadata
assert logger.model_call_details["messages"] == [{"role": "user", "content": "again"}]
assert logger.model_call_details["thinking"] == {"type": "enabled", "budget_tokens": 1024}
class _BaselineLogger:
def __init__(self) -> None:
self.baseline_cache_context: object = object()
self.model_call_details: dict[str, object] = {}
self.invalidated: list[str] = []
async def invalidate_baseline_cache_estimate(self, reason: str) -> None:
self.invalidated.append(reason)
@pytest.mark.asyncio
async def test_a_retried_request_invalidates_the_baseline_cache_estimate() -> None:
logger: Final = _BaselineLogger()
await legacy.retried_request(cast(Logging, logger), {"messages": []}, {"messages": []}) # cast-ok: duck-typed
assert logger.invalidated == ["retried_request"]
def test_stream_head_is_stamped_in_place_with_the_call_hidden_params() -> None:
logger: Final = _messages_logger()
head: Final[dict[str, object]] = {"additional_headers": {"request-id": "req_native"}}

14
tests/unit/rust_bridge/test_detectors.py
View file

@ -31,8 +31,6 @@ MAPPED_FILE: Final = {
}
MAPPED_FILE_INPUT: Final = [{"role": "user", "content": [{"type": "input_file", "file_id": "litellm-file"}]}]
CACHED_TEXT: Final = {"type": "input_text", "text": "hi", "cache_control": {"type": "ephemeral"}}
THINKING_TURN: Final = {"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "s"}]}
REDACTED: Final = {"type": "redacted_thinking", "data": "x"}
STRICT_TOOL: Final = {"name": "b", "input_schema": {"type": "object"}, "strict": True}
TOOL_SEARCH_TYPE: Final = sorted(ANTHROPIC_TOOL_SEARCH_TOOL_TYPES)[0]
@ -196,18 +194,6 @@ def _wrap_nested_level(value: object, depth: int) -> object:
(detectors.cache_control_injection, {"enable_prompt_caching": True}, True),
(detectors.cache_control_injection, {"cache_control_injection_points": [{"location": "message"}]}, True),
(detectors.cache_control_injection, {"cache_control_injection_points": []}, False),
(messages_detectors.thinking_signature_retry, {"messages": [THINKING_TURN]}, True),
(
messages_detectors.thinking_signature_retry,
{"messages": [{"role": "assistant", "content": [REDACTED]}]},
True,
),
(messages_detectors.thinking_signature_retry, {"messages": [{"role": "assistant", "content": "done"}]}, False),
(
messages_detectors.thinking_signature_retry,
{"messages": [{"role": "user", "content": [{"type": "text"}]}]},
False,
),
(messages_detectors.tool_custom_field, {"tools": [{"name": "a"}, {"name": "b", "custom": {}}]}, True),
(messages_detectors.tool_custom_field, {"tools": [{"name": "a", "defer_loading": True}]}, False),
(messages_detectors.tool_custom_field, {}, False),