From 45dcc9cc806eb9c9770224f2282bd7afbcecddc3 Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Tue, 10 Nov 2026 20:43:24 +0000 Subject: [PATCH] feat(messages): retry rejected thinking signatures on rust --- .../python_contract.json | 5 + .../callbacks-legacy-python/src/adapter.rs | 83 +++++++ .../callbacks-legacy-python/src/mapping.rs | 3 + .../callbacks-legacy-python/src/python.rs | 2 + .../src/test_support.rs | 4 + litellm-rust/crates/host-native/src/driver.rs | 7 + litellm-rust/crates/host-python/src/driver.rs | 8 + .../host-python/src/hooks/chain/dispatch.rs | 7 + .../crates/host-python/tests/support/mod.rs | 3 + litellm-rust/crates/host/src/interceptors.rs | 14 ++ litellm-rust/crates/host/src/lifecycle.rs | 13 +- .../crates/host/src/machine/context.rs | 11 + litellm-rust/crates/host/src/protocol.rs | 4 + .../inference-messages/src/constants.rs | 2 + .../crates/inference-messages/src/handler.rs | 114 +++++++--- .../inference-messages/tests/messages/main.rs | 1 + .../tests/messages/retry.rs | 209 ++++++++++++++++++ .../inference-messages/tests/support/mod.rs | 10 + litellm-rust/crates/inference/src/context.rs | 13 +- .../crates/llms/src/anthropic/common_utils.rs | 162 ++++++++++++++ .../src/anthropic/messages/transformation.rs | 113 +++++++++- .../src/base_llm/messages/transformation.rs | 19 ++ litellm/messages/dispatch.py | 1 - litellm/rust_bridge/catalog.py | 2 - .../host/callbacks_legacy_python.py | 13 ++ litellm/rust_bridge/messages/detectors.py | 13 -- .../sdk/messages/test_native_parity.py | 124 ++++++++--- .../messages/test_dispatch.py | 48 +++- .../host/test_callbacks_legacy_python.py | 50 ++++- tests/unit/rust_bridge/test_detectors.py | 14 -- 30 files changed, 976 insertions(+), 96 deletions(-) create mode 100644 litellm-rust/crates/inference-messages/tests/messages/retry.rs diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index 8a212223523..662a964be4f 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -43,6 +43,11 @@ "headers", "body" ], + "retried_request": [ + "logger", + "logged_body", + "body" + ], "record_api_window": [ "logger", "start", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index f68541c5c7b..b5b95fa0822 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -664,6 +664,38 @@ impl LegacyLogging { Ok(HookStep::Ready(())) } + pub(crate) fn request_retried( + &mut self, + py: Python<'_>, + wire: &WireRequest, + ) -> PyResult> { + 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>, + ) -> PyResult> { + 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, &[]) } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs index cfa9e53aa7b..7c6bc1a52df 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs @@ -179,6 +179,9 @@ impl CallHooks 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) } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 5b6faa3dbc2..7b96301832e 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -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")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index ee16e0c0059..2c088330426 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -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 ), diff --git a/litellm-rust/crates/host-native/src/driver.rs b/litellm-rust/crates/host-native/src/driver.rs index af3248848a2..80eaf91930d 100644 --- a/litellm-rust/crates/host-native/src/driver.rs +++ b/litellm-rust/crates/host-native/src/driver.rs @@ -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)); diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 052156cbc5d..3c7a2654d8f 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -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 = { diff --git a/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs b/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs index 4939995cefa..7f91a5c8f32 100644 --- a/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs +++ b/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs @@ -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 { diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index 15e864c0718..4adac754069 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -422,6 +422,9 @@ impl CallHooks 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)) } diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index 9398ea4bad6..e54180b8630 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -107,6 +107,13 @@ pub trait Interceptors: Send + Sync { &self, raw: RawResponse, ) -> impl Future> + Send; + + fn provider_request_retried( + &self, + _wire: WireRequest, + ) -> impl Future> + Send { + async { Ok(()) } + } } impl + ?Sized> Interceptors for &T { @@ -128,6 +135,13 @@ impl + ?Sized> Interceptors for &T { ) -> impl Future> + Send { (**self).after_provider_response(raw) } + + fn provider_request_retried( + &self, + wire: WireRequest, + ) -> impl Future> + Send { + (**self).provider_request_retried(wire) + } } impl Interceptors for () { diff --git a/litellm-rust/crates/host/src/lifecycle.rs b/litellm-rust/crates/host/src/lifecycle.rs index e6b31381494..fc4f716d63f 100644 --- a/litellm-rust/crates/host/src/lifecycle.rs +++ b/litellm-rust/crates/host/src/lifecycle.rs @@ -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 { }, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq)] pub enum ExecutionEvent { ResultReady { facts: crate::interceptors::ExecutionFacts, @@ -58,6 +61,9 @@ pub enum ExecutionEvent { ProviderResponseReceived { raw: Raw, }, + ProviderRequestRetried { + wire: WireRequest, + }, } impl> CallEvent { @@ -76,6 +82,9 @@ impl> CallEvent { + CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { wire: wire.clone() }) + } Self::Succeeded { timing, .. } => CallEvent::Succeeded { timing: *timing, response: (), diff --git a/litellm-rust/crates/host/src/machine/context.rs b/litellm-rust/crates/host/src/machine/context.rs index 1e53605ed7b..34c9505d8c9 100644 --- a/litellm-rust/crates/host/src/machine/context.rs +++ b/litellm-rust/crates/host/src/machine/context.rs @@ -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(Channel

); diff --git a/litellm-rust/crates/host/src/protocol.rs b/litellm-rust/crates/host/src/protocol.rs index c24557b0348..6c52b50508c 100644 --- a/litellm-rust/crates/host/src/protocol.rs +++ b/litellm-rust/crates/host/src/protocol.rs @@ -33,6 +33,10 @@ pub enum InterceptRequest { raw: RawResponse, reply: Reply>, }, + ProviderRequestRetried { + wire: Box, + reply: Reply>, + }, } pub enum StreamDelivery { diff --git a/litellm-rust/crates/inference-messages/src/constants.rs b/litellm-rust/crates/inference-messages/src/constants.rs index c0b023be7c7..867b2b50da4 100644 --- a/litellm-rust/crates/inference-messages/src/constants.rs +++ b/litellm-rust/crates/inference-messages/src/constants.rs @@ -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; diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index c7a7d93bf98..e34cbebaee8 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -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, timeout: Option, 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>, ) -> Result { + 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>, +) -> Result { + 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 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 { 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()), } } diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index fe203287766..e67556af5f3 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -22,6 +22,7 @@ use support::*; mod host; mod request; mod response; +mod retry; mod secrets; mod stream; mod tencent; diff --git a/litellm-rust/crates/inference-messages/tests/messages/retry.rs b/litellm-rust/crates/inference-messages/tests/messages/retry.rs new file mode 100644 index 00000000000..b1d072c8a3b --- /dev/null +++ b/litellm-rust/crates/inference-messages/tests/messages/retry.rs @@ -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, +) -> Result<(MessagesCallResponse, Vec), 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) -> Vec { + 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::::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 = 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::::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::::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()); +} diff --git a/litellm-rust/crates/inference-messages/tests/support/mod.rs b/litellm-rust/crates/inference-messages/tests/support/mod.rs index 3cecc865070..2bf1ea3f08f 100644 --- a/litellm-rust/crates/inference-messages/tests/support/mod.rs +++ b/litellm-rust/crates/inference-messages/tests/support/mod.rs @@ -117,6 +117,7 @@ pub struct RecordingCall { pub events: Arc, pub chunks: Mutex>, pub head: Mutex>, + pub retried: Mutex>, } #[derive(Default)] @@ -167,6 +168,7 @@ impl RecordingCall

{ events: Arc::new(CallEvents::default()), chunks: Mutex::new(Vec::new()), head: Mutex::new(None), + retried: Mutex::new(Vec::new()), } } } @@ -195,6 +197,14 @@ impl 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

RecordingCall

diff --git a/litellm-rust/crates/inference/src/context.rs b/litellm-rust/crates/inference/src/context.rs index ef2635d24a1..a1ef731f857 100644 --- a/litellm-rust/crates/inference/src/context.rs +++ b/litellm-rust/crates/inference/src/context.rs @@ -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> 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 + } } diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 0854820358b..783a463322d 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -650,6 +650,62 @@ pub fn flatten_unencrypted_web_search_results(messages: Vec) -> Vec 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 { + 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"]); + } } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 18789e5b2be..605fdcac193 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -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); + } } diff --git a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs index 744ce89afdb..b74ea8e4fa8 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs @@ -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)] diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index faa604a25ce..6ce20cfad98 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -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, } ) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index c493617003d..1e80c5526cb 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -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, } diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index ba557f769eb..ec73a587693 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -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 ( diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 99555bb4852..aeb7edbb29e 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -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) 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..fc689e2723c 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 @@ -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 diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index b580eca2e42..65c929ce6a9 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -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", diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index cb594767908..baf7a029357 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -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"}} diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index fccc01dd8e6..010427cc9b9 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -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),