Merge remote-tracking branch 'origin/main' into litellm_e2e_cost_calculation_scripted_provider

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	tests/e2e/conftest.py
#	tests/e2e/e2e_config.py

ARCHITECTURE.md

@@ -45,7 +45,7 @@ sequenceDiagram
     ProxyServer->>Auth: user_api_key_auth()
     Auth->>Redis: Check API key cache
     Redis-->>Auth: Key info + spend limits
-    ProxyServer->>Hooks: max_budget_limiter, parallel_request_limiter
+    ProxyServer->>Hooks: parallel_request_limiter, cache_control_check
     Hooks->>Redis: Check/increment rate limit counters
     ProxyServer->>Router: route_request()
     Router->>Main: litellm.acompletion()
@@ -145,7 +145,6 @@ graph TD
 
 | Hook | File | Purpose |
 |------|------|---------|
-| `max_budget_limiter` | `proxy/hooks/max_budget_limiter.py` | Enforce budget limits |
 | `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
 | `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
 | `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |

litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql

@@ -0,0 +1,5 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
+
+-- AlterTable
+ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -426,6 +426,7 @@ model LiteLLM_VerificationToken {
     key_alias   String?
     soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down
     spend      Float    @default(0.0)
+    total_spend Float   @default(0.0)
     expires    DateTime?
     models     String[]
     aliases    Json  @default("{}")
@@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken {
     key_alias               String?
     soft_budget_cooldown    Boolean  @default(false)
     spend                   Float    @default(0.0)
+    total_spend             Float    @default(0.0)
     expires                 DateTime?
     models                  String[]
     aliases                 Json     @default("{}")

litellm-rust/Cargo.lock

@@ -1951,6 +1951,7 @@ dependencies = [
 name = "litellm-python-bridge"
 version = "0.1.0"
 dependencies = [
+ "bytes",
  "criterion",
  "futures-util",
  "litellm-auth",

litellm-rust/crates/core/src/ocr/client.rs

@@ -54,9 +54,16 @@ impl OcrClient {
             match call.resume(result.take()).await? {
                 OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
                     result = Some(OcrHostResult::Request(Ok((
-                        Box::new(request.take().ok_or_else(|| {
-                            Error::InvalidRequest("OCR request was already projected".into())
-                        })?),
+                        Box::new(
+                            request
+                                .take()
+                                .ok_or_else(|| {
+                                    Error::InvalidRequest(
+                                        "OCR request was already projected".into(),
+                                    )
+                                })?
+                                .into(),
+                        ),
                         false,
                     ))))
                 }

litellm-rust/crates/core/src/ocr/document.rs

@@ -1,3 +1,6 @@
+use std::io::Read;
+use std::path::Path;
+
 use base64::{Engine, engine::general_purpose::STANDARD};
 use data_url::mime::Mime;
 use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError};
@@ -5,12 +8,52 @@ use reqwest::Url;
 use serde_json::Map;
 
 use super::error::{OcrError, OcrRequestError, OcrResponseError};
-use super::types::{OcrConnection, OcrDocument};
+use super::types::{OcrConnection, OcrDocument, OcrDocumentInput};
 use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS};
 use crate::media::Error as MediaError;
 use crate::media::{DownloadPolicy, MediaFetcher};
 use crate::transport::Error as TransportError;
 
+pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, super::Error> {
+    match input {
+        OcrDocumentInput::Document(document) => Ok(document),
+        OcrDocumentInput::Path { path, mime_type } => {
+            read_path_document(&path, mime_type.as_deref())
+        }
+        OcrDocumentInput::Bytes {
+            bytes,
+            file_name,
+            mime_type,
+        } => Ok(encode_file_document(
+            &bytes,
+            file_name.as_deref(),
+            mime_type.as_deref(),
+        )?),
+        OcrDocumentInput::HostReader { .. } => Err(super::Error::InvalidRequest(
+            "OCR file reader was not read by the host".into(),
+        )),
+    }
+}
+
+pub fn read_path_document(
+    path: &Path,
+    mime_type: Option<&str>,
+) -> Result<OcrDocument, super::Error> {
+    let mut bytes = Vec::new();
+    std::fs::File::open(path)
+        .and_then(|file| {
+            file.take(OCR_INLINE_MAX_BYTES as u64 + 1)
+                .read_to_end(&mut bytes)
+        })
+        .map_err(|source| super::Error::FileRead {
+            path: path.to_owned(),
+            kind: source.kind(),
+            message: source.to_string(),
+        })?;
+    let name = path.file_name().map(|name| name.to_string_lossy());
+    Ok(encode_file_document(&bytes, name.as_deref(), mime_type)?)
+}
+
 pub fn encode_file_document(
     bytes: &[u8],
     file_name: Option<&str>,
@@ -75,18 +118,6 @@ pub fn mime_type_for_name(name: &str) -> &'static str {
     }
 }
 
-pub fn upload_mime_type<'a>(file_name: Option<&str>, content_type: Option<&'a str>) -> &'a str {
-    match content_type
-        .and_then(|value| value.split(';').next())
-        .map(str::trim)
-    {
-        Some(value) if !value.is_empty() && value != "application/octet-stream" => value,
-        _ => file_name
-            .map(mime_type_for_name)
-            .unwrap_or("application/octet-stream"),
-    }
-}
-
 pub(crate) struct InlineDocument<'a>(DataUrl<'a>);
 
 impl<'a> InlineDocument<'a> {
@@ -230,24 +261,65 @@ mod tests {
     }
 
     #[test]
-    fn upload_mime_mapping_matches_python() {
+    fn path_documents_are_read_and_named_by_core() {
+        let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::<u64>()));
+        std::fs::create_dir_all(&dir).unwrap();
+        let path = dir.join("scan.png");
+        std::fs::write(&path, b"abc").unwrap();
         assert_eq!(
-            upload_mime_type(Some("report.pdf"), Some("application/octet-stream")),
-            "application/pdf"
-        );
-        assert_eq!(upload_mime_type(Some("image.png"), None), "image/png");
-        assert_eq!(upload_mime_type(None, None), "application/octet-stream");
-        assert_eq!(
-            upload_mime_type(Some("doc.pdf"), Some("application/pdf; charset=utf-8")),
-            "application/pdf"
+            prepare_document(OcrDocumentInput::Path {
+                path: path.clone(),
+                mime_type: None,
+            })
+            .unwrap(),
+            OcrDocument::ImageUrl {
+                image_url: "data:image/png;base64,YWJj".into(),
+                extra_fields: Map::new(),
+            }
         );
         assert_eq!(
-            upload_mime_type(
-                Some("img.png"),
-                Some("image/png; charset=utf-8; boundary=something")
-            ),
-            "image/png"
+            prepare_document(OcrDocumentInput::Path {
+                path: path.clone(),
+                mime_type: Some("application/pdf".into()),
+            })
+            .unwrap(),
+            document("data:application/pdf;base64,YWJj")
         );
+        std::fs::write(&path, vec![b'a'; OCR_INLINE_MAX_BYTES + 1]).unwrap();
+        assert_eq!(
+            prepare_document(OcrDocumentInput::Path {
+                path: path.clone(),
+                mime_type: None,
+            }),
+            Err(OcrRequestError::InlineDocumentTooLarge.into())
+        );
+        std::fs::remove_dir_all(&dir).unwrap();
+
+        let missing = dir.join("missing.pdf");
+        let Err(super::super::Error::FileRead { path, kind, .. }) =
+            prepare_document(OcrDocumentInput::Path {
+                path: missing.clone(),
+                mime_type: None,
+            })
+        else {
+            panic!("missing paths must surface a file read error");
+        };
+        assert_eq!(path, missing);
+        assert_eq!(kind, std::io::ErrorKind::NotFound);
+    }
+
+    #[test]
+    fn byte_documents_are_encoded_and_host_readers_must_be_read_first() {
+        assert_eq!(
+            prepare_document(OcrDocumentInput::Bytes {
+                bytes: b"abc".as_slice().into(),
+                file_name: Some("scan.pdf".into()),
+                mime_type: None,
+            })
+            .unwrap(),
+            document("data:application/pdf;base64,YWJj")
+        );
+        assert!(prepare_document(OcrDocumentInput::HostReader { mime_type: None }).is_err());
     }
 
     #[test]

litellm-rust/crates/core/src/ocr/error.rs

@@ -50,6 +50,12 @@ pub enum Error {
     Connect(String),
     #[error("routing error: {0}")]
     Routing(String),
+    #[error("Failed to read OCR file {}: {message}", path.display())]
+    FileRead {
+        path: std::path::PathBuf,
+        kind: std::io::ErrorKind,
+        message: String,
+    },
     /// The request is outside the surface this route covers in Rust. Hosts that
     /// keep a reference implementation treat this as "fall back", not "fail".
     #[error("unsupported by the rust path: {0}")]

litellm-rust/crates/core/src/ocr/lifecycle.rs

@@ -9,6 +9,7 @@ use super::hooks::{
     OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest,
     OcrPreCallRequest,
 };
+use super::types::{OcrDocumentInput, OcrFileContent};
 use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
 use crate::call_lifecycle::host::{
     HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase,
@@ -52,6 +53,7 @@ impl OcrAdmission {
 #[derive(Clone, Debug)]
 pub enum OcrHostOperation {
     ProjectRequest,
+    ReadDocument,
     Lifecycle(HostPhase),
     ConstructResponse(Arc<LiteLLMOcrResponse>),
     MapFailure(Error),
@@ -83,7 +85,8 @@ impl OcrHostOperation {
 }
 
 pub enum OcrHostResult {
-    Request(Result<(Box<LiteLLMOcrRequest>, bool), Error>),
+    Request(Result<(Box<LiteLLMOcrRequest<OcrDocumentInput>>, bool), Error>),
+    Document(Result<OcrFileContent, Error>),
     Lifecycle(Result<(), HostFailure<Error>>),
     AzureAdToken(Result<ResolvedCredential, AuthError>),
     PreCall(Result<OcrPreCallRequest, Error>),
@@ -313,7 +316,7 @@ struct PendingOperation {
 
 struct OcrExecution {
     client: Option<OcrClient>,
-    request: Option<LiteLLMOcrRequest>,
+    request: Option<LiteLLMOcrRequest<OcrDocumentInput>>,
     operations_tx: mpsc::UnboundedSender<PendingOperation>,
     operations_rx: mpsc::UnboundedReceiver<PendingOperation>,
     pending_result: Option<oneshot::Sender<OcrHostResult>>,
@@ -397,12 +400,14 @@ impl OcrExecution {
                 },
             )));
         }
-        request.hooks = Arc::new(ProtocolHooks {
+        let hooks = Arc::new(ProtocolHooks {
             operations: self.operations_tx.clone(),
             intercepts_requests,
             terminal: self.terminal.clone(),
         });
+        request.hooks = hooks.clone();
         self.execution = Some(tokio::spawn(async move {
+            let request = prepare_request_document(request, &hooks).await?;
             perform_ocr_request(&client, request).await
         }));
     }
@@ -423,6 +428,39 @@ impl OcrExecution {
     }
 }
 
+async fn prepare_request_document(
+    request: LiteLLMOcrRequest<OcrDocumentInput>,
+    hooks: &ProtocolHooks,
+) -> Result<LiteLLMOcrRequest, Error> {
+    let request = match &request.document {
+        OcrDocumentInput::HostReader { mime_type } => {
+            let mime_type = mime_type.clone();
+            let content = match hooks.invoke(OcrHostOperation::ReadDocument).await? {
+                OcrHostResult::Document(result) => result?,
+                _ => {
+                    return Err(Error::InvalidRequest(
+                        "invalid OCR document read host result".into(),
+                    ));
+                }
+            };
+            request.with_document(OcrDocumentInput::Bytes {
+                bytes: content.bytes,
+                file_name: content.file_name,
+                mime_type,
+            })
+        }
+        _ => request,
+    };
+    if let OcrDocumentInput::Document(_) = &request.document {
+        return request.map_document(super::document::prepare_document);
+    }
+    tokio::task::spawn_blocking(move || request.map_document(super::document::prepare_document))
+        .await
+        .map_err(|error| {
+            Error::InvalidRequest(format!("OCR document preparation task failed: {error}"))
+        })?
+}
+
 impl Drop for OcrExecution {
     fn drop(&mut self) {
         if let Some(execution) = &self.execution {
@@ -567,6 +605,9 @@ impl OcrHost for NoopOcrHost {
                 OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
                     Error::InvalidRequest("OCR host has no request projection".into()),
                 )),
+                OcrHostOperation::ReadDocument => OcrHostResult::Document(Err(
+                    Error::InvalidRequest("OCR host has no document reader".into()),
+                )),
                 OcrHostOperation::Lifecycle(_)
                 | OcrHostOperation::ConstructResponse(_)
                 | OcrHostOperation::MapFailure(_)
@@ -602,6 +643,9 @@ impl OcrHost for OcrHookHost {
                 OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
                     Error::InvalidRequest("OCR hook host has no request projection".into()),
                 )),
+                OcrHostOperation::ReadDocument => OcrHostResult::Document(Err(
+                    Error::InvalidRequest("OCR hook host has no document reader".into()),
+                )),
                 OcrHostOperation::Success {
                     context,
                     response,

litellm-rust/crates/core/src/ocr/mod.rs

@@ -13,12 +13,15 @@ pub mod types;
 pub mod wire;
 
 pub use client::{OcrClient, ocr};
-pub use document::{encode_file_document, mime_type_for_name, upload_mime_type};
+pub use document::{encode_file_document, mime_type_for_name, read_path_document};
 pub use lifecycle::{
     NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline,
     OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult,
 };
-pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument};
+pub use types::{
+    LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrDocumentInput,
+    OcrFileContent,
+};
 
 #[cfg(test)]
 #[path = "../../tests/azure_ai_ocr.rs"]

litellm-rust/crates/core/src/ocr/types.rs

@@ -1,7 +1,10 @@
 use std::collections::BTreeMap;
+use std::convert::Infallible;
+use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::Duration;
 
+use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 
@@ -50,6 +53,35 @@ impl OcrDocument {
     }
 }
 
+#[derive(Clone, Debug, PartialEq)]
+pub enum OcrDocumentInput {
+    Document(OcrDocument),
+    Path {
+        path: PathBuf,
+        mime_type: Option<String>,
+    },
+    Bytes {
+        bytes: Bytes,
+        file_name: Option<String>,
+        mime_type: Option<String>,
+    },
+    HostReader {
+        mime_type: Option<String>,
+    },
+}
+
+impl From<OcrDocument> for OcrDocumentInput {
+    fn from(document: OcrDocument) -> Self {
+        Self::Document(document)
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct OcrFileContent {
+    pub bytes: Bytes,
+    pub file_name: Option<String>,
+}
+
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "lowercase")]
 pub enum OcrResponseFormat {
@@ -89,9 +121,9 @@ impl Default for OcrConnection {
     }
 }
 
-pub struct LiteLLMOcrRequest {
+pub struct LiteLLMOcrRequest<D = OcrDocument> {
     pub model: String,
-    pub document: OcrDocument,
+    pub document: D,
     pub connection: OcrConnection,
     pub hooks: Arc<dyn OcrHooks>,
     pub litellm_call_id: Option<String>,
@@ -101,10 +133,10 @@ pub struct LiteLLMOcrRequest {
     pub(crate) adapter: OcrAdapterKind,
 }
 
-impl LiteLLMOcrRequest {
+impl<D> LiteLLMOcrRequest<D> {
     pub fn new(
         model: String,
-        document: OcrDocument,
+        document: D,
         custom_llm_provider: Option<&str>,
         optional_params: Map<String, Value>,
     ) -> Result<Self, Error> {
@@ -151,6 +183,36 @@ impl LiteLLMOcrRequest {
             ..self
         }
     }
+
+    pub fn map_document<T, E>(
+        self,
+        map: impl FnOnce(D) -> Result<T, E>,
+    ) -> Result<LiteLLMOcrRequest<T>, E> {
+        Ok(LiteLLMOcrRequest {
+            model: self.model,
+            document: map(self.document)?,
+            connection: self.connection,
+            hooks: self.hooks,
+            litellm_call_id: self.litellm_call_id,
+            optional_params: self.optional_params,
+            input_sources: self.input_sources,
+            azure_ad_token_provider: self.azure_ad_token_provider,
+            adapter: self.adapter,
+        })
+    }
+
+    pub fn with_document<T>(self, document: T) -> LiteLLMOcrRequest<T> {
+        let Ok(request) = self.map_document(|_| Ok::<T, Infallible>(document));
+        request
+    }
+}
+
+impl From<LiteLLMOcrRequest> for LiteLLMOcrRequest<OcrDocumentInput> {
+    fn from(request: LiteLLMOcrRequest) -> Self {
+        let Ok(request) = request
+            .map_document(|document| Ok::<_, Infallible>(OcrDocumentInput::Document(document)));
+        request
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]

litellm-rust/crates/core/src/ocr/wire.rs

@@ -68,9 +68,9 @@ pub struct DecodedOcrResponse<T> {
 
 #[derive(Deserialize)]
 #[serde(deny_unknown_fields)]
-pub struct OcrWireRequest {
+pub struct OcrWireRequest<D = Value> {
     pub model: String,
-    pub document: Value,
+    pub document: D,
     pub api_key: Option<String>,
     pub api_base: Option<String>,
     pub custom_llm_provider: Option<String>,
@@ -141,10 +141,34 @@ pub fn consumed_optional_params(
 }
 
 pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error> {
+    let OcrWireRequest {
+        model,
+        document,
+        api_key,
+        api_base,
+        custom_llm_provider,
+        extra_headers,
+        optional_params,
+        input_sources,
+        timeout_seconds,
+    } = wire;
+    decode_request_input(OcrWireRequest {
+        model,
+        document: decode_document(document)?,
+        api_key,
+        api_base,
+        custom_llm_provider,
+        extra_headers,
+        optional_params,
+        input_sources,
+        timeout_seconds,
+    })
+}
+
+pub fn decode_request_input<D>(wire: OcrWireRequest<D>) -> Result<LiteLLMOcrRequest<D>, Error> {
     let api_key_source = source_for(&wire.input_sources, "api_key");
     let api_base_source = source_for(&wire.input_sources, "api_base");
     let extra_headers_source = source_for(&wire.input_sources, "extra_headers");
-    let document = decode_document(wire.document)?;
     let headers = wire
         .extra_headers
         .unwrap_or_default()
@@ -183,7 +207,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
         .unwrap_or(defaults.max_response_bytes);
     let request = LiteLLMOcrRequest::new(
         wire.model,
-        document,
+        wire.document,
         wire.custom_llm_provider.as_deref(),
         wire.optional_params
             .into_iter()
@@ -209,14 +233,14 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
     })
 }
 
-fn decode_document(value: Value) -> Result<OcrDocument, OcrRequestError> {
+pub fn decode_document(value: Value) -> Result<OcrDocument, Error> {
     let kind = value.get("type").and_then(Value::as_str);
     let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none()
         || matches!(kind, Some("image_url")) && value.get("image_url").is_none();
     if missing_url {
-        return Err(OcrRequestError::MissingDocumentUrl);
+        return Err(OcrRequestError::MissingDocumentUrl.into());
     }
-    decode_request_value(value, "document")
+    Ok(decode_request_value(value, "document")?)
 }
 
 fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSource {
@@ -334,10 +358,7 @@ mod tests {
             serde_json::json!({"type": "document_url"}),
             serde_json::json!({"type": "image_url"}),
         ] {
-            assert_eq!(
-                decode_document(document),
-                Err(OcrRequestError::MissingDocumentUrl)
-            );
+            assert_eq!(decode_document(document), Err(Error::MissingDocumentUrl));
         }
     }
 }

litellm-rust/crates/core/tests/ocr.rs

@@ -348,13 +348,14 @@ async fn fallible_host_phases_do_not_replay_or_reach_transport() {
                     }
                     OcrHostOperation::ProjectRequest => {
                         result = Some(OcrHostResult::Request(Ok((
-                            Box::new(request.take().unwrap()),
+                            Box::new(request.take().unwrap().into()),
                             false,
                         ))))
                     }
                     OcrHostOperation::AcquireAzureAdToken => {
                         panic!("test request has no token provider")
                     }
+                    OcrHostOperation::ReadDocument => panic!("test request has no file reader"),
                     OcrHostOperation::PreCall(request) => {
                         phases.push("pre");
                         result = Some(OcrHostResult::PreCall(if failure_phase == "pre" {
@@ -405,7 +406,7 @@ async fn invalid_provider_response_runs_post_call_before_normalization_failure()
         match call.resume(result.take()).await {
             Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => {
                 result = Some(OcrHostResult::Request(Ok((
-                    Box::new(request.take().unwrap()),
+                    Box::new(request.take().unwrap().into()),
                     false,
                 ))));
             }
@@ -467,9 +468,10 @@ async fn direct_native_host_drives_the_same_state_machine() {
                     _ => panic!("unexpected OCR operation"),
                 });
                 result = Some(match operation {
-                    OcrHostOperation::ProjectRequest => {
-                        OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false)))
-                    }
+                    OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok((
+                        Box::new(request.take().unwrap().into()),
+                        false,
+                    ))),
                     operation => host.invoke(operation).await,
                 });
             }
@@ -501,6 +503,137 @@ async fn direct_native_host_drives_the_same_state_machine() {
     ));
 }
 
+async fn drive_native_file_call(
+    request: super::LiteLLMOcrRequest<super::OcrDocumentInput>,
+    content: Result<super::OcrFileContent, crate::ocr::Error>,
+) -> (Result<super::LiteLLMOcrResponse, crate::ocr::Error>, usize) {
+    let NativeOutcome::Completed(mut call) =
+        OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all())
+    else {
+        panic!("supported call declined")
+    };
+    let mut request = Some(request);
+    let mut content = Some(content);
+    let mut result = None;
+    let mut reads = 0;
+    let outcome = loop {
+        match call.resume(result.take()).await {
+            Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => {
+                result = Some(OcrHostResult::Request(Ok((
+                    Box::new(request.take().unwrap()),
+                    false,
+                ))));
+            }
+            Ok(OcrCallStep::Host(OcrHostOperation::ReadDocument)) => {
+                reads += 1;
+                result = Some(OcrHostResult::Document(content.take().unwrap()));
+            }
+            Ok(OcrCallStep::Host(operation)) => result = Some(NoopOcrHost.invoke(operation).await),
+            Ok(OcrCallStep::Complete(response)) => break Ok(response),
+            Err(error) => break Err(error),
+        }
+    };
+    (outcome, reads)
+}
+
+#[tokio::test]
+async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
+        "pages":[{"index":0,"markdown":"file"}]
+    }))])
+    .await;
+    let request = wire_request("mistral/model", &base, json!({})).with_document(
+        super::OcrDocumentInput::HostReader {
+            mime_type: Some("application/pdf".into()),
+        },
+    );
+    let (response, reads) = drive_native_file_call(
+        request,
+        Ok(super::OcrFileContent {
+            bytes: b"abc".as_slice().into(),
+            file_name: Some("scan.png".into()),
+        }),
+    )
+    .await;
+    server.await.unwrap();
+    assert_eq!(response.unwrap().pages[0]["markdown"], "file");
+    assert_eq!(reads, 1);
+    assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj"));
+}
+
+#[tokio::test]
+async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() {
+    let (base, seen, _server) = mock_server(vec![]).await;
+    let request = wire_request("mistral/model", &base, json!({}));
+    let failure = crate::ocr::Error::InvalidRequest("reader exploded".into());
+    let (response, reads) = drive_native_file_call(
+        request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }),
+        Err(failure.clone()),
+    )
+    .await;
+    assert_eq!(response.unwrap_err(), failure);
+    assert_eq!(reads, 1);
+
+    let request = wire_request("mistral/model", &base, json!({}));
+    let (response, _) = drive_native_file_call(
+        request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }),
+        Ok(super::OcrFileContent {
+            bytes: Default::default(),
+            file_name: None,
+        }),
+    )
+    .await;
+    assert!(matches!(
+        response.unwrap_err(),
+        crate::ocr::Error::InvalidRequest(_)
+    ));
+    assert!(seen.lock().unwrap().is_empty());
+}
+
+#[tokio::test]
+async fn path_documents_are_read_by_core_without_a_host_operation() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
+        "pages":[{"index":0,"markdown":"path"}]
+    }))])
+    .await;
+    let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::<u64>()));
+    std::fs::create_dir_all(&dir).unwrap();
+    let path = dir.join("scan.png");
+    std::fs::write(&path, b"abc").unwrap();
+    let request = wire_request("mistral/model", &base, json!({})).with_document(
+        super::OcrDocumentInput::Path {
+            path: path.clone(),
+            mime_type: None,
+        },
+    );
+    let (response, reads) = drive_native_file_call(
+        request,
+        Err(crate::ocr::Error::InvalidRequest("unused".into())),
+    )
+    .await;
+    server.await.unwrap();
+    std::fs::remove_dir_all(&dir).unwrap();
+    assert_eq!(response.unwrap().pages[0]["markdown"], "path");
+    assert_eq!(reads, 0);
+    assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj"));
+
+    let (base, seen, _server) = mock_server(vec![]).await;
+    let request = wire_request("mistral/model", &base, json!({}));
+    let (response, _) = drive_native_file_call(
+        request.with_document(super::OcrDocumentInput::Path {
+            path: path.clone(),
+            mime_type: None,
+        }),
+        Err(crate::ocr::Error::InvalidRequest("unused".into())),
+    )
+    .await;
+    assert!(matches!(
+        response.unwrap_err(),
+        crate::ocr::Error::FileRead { path: failed, kind: std::io::ErrorKind::NotFound, .. } if failed == path
+    ));
+    assert!(seen.lock().unwrap().is_empty());
+}
+
 #[tokio::test]
 async fn public_finalization_failure_never_dispatches_success_or_replays_provider() {
     use crate::call_lifecycle::host::{HostFailure, HostPhase};
@@ -543,9 +676,10 @@ async fn public_finalization_failure_never_dispatches_success_or_replays_provide
                     | OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => {
                         panic!("finalization failure used provider/success dispatch")
                     }
-                    OcrHostOperation::ProjectRequest => {
-                        OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false)))
-                    }
+                    OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok((
+                        Box::new(request.take().unwrap().into()),
+                        false,
+                    ))),
                     operation => host.invoke(operation).await,
                 });
             }
@@ -582,7 +716,7 @@ async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption
             OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break,
             OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
                 result = Some(OcrHostResult::Request(Ok((
-                    Box::new(request.take().unwrap()),
+                    Box::new(request.take().unwrap().into()),
                     false,
                 ))))
             }
@@ -799,7 +933,7 @@ async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_
                     _ = entered.notified() => break,
                     step = call.resume(result.take()) => {
                         result = Some(match step.unwrap() {
-                            OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))),
+                            OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap().into()), false))),
                             OcrCallStep::Host(operation) => NoopOcrHost.invoke(operation).await,
                             OcrCallStep::Complete(_) => panic!("pending provider completed"),
                         });

litellm-rust/crates/python-bridge/Cargo.toml

@@ -16,6 +16,7 @@ extension-module = ["pyo3/extension-module"]
 panic-test = []
 
 [dependencies]
+bytes.workspace = true
 futures-util.workspace = true
 litellm-core.workspace = true
 litellm-auth.workspace = true

litellm-rust/crates/python-bridge/src/routes/ocr/document.rs

@@ -1,97 +1,56 @@
-use std::io::Read;
 use std::path::PathBuf;
 
-use pyo3::exceptions::{PyFileNotFoundError, PyTypeError, PyValueError};
+use bytes::Bytes;
+use pyo3::exceptions::{PyTypeError, PyValueError};
+use pyo3::gc::{PyTraverseError, PyVisit};
 use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedBytes;
-#[cfg(test)]
-use pyo3::types::PyDict;
 use pyo3::types::{PyBytes, PyString};
 
-use litellm_core::constants::OCR_INLINE_MAX_BYTES;
-use litellm_core::ocr::{OcrDocument, encode_file_document, mime_type_for_name, upload_mime_type};
-use litellm_python_interop::to_py_preserving_errors;
+use litellm_core::ocr::{OcrDocumentInput, OcrFileContent};
 
-enum FileBytes {
-    Python(PyBackedBytes),
-    Native(Vec<u8>),
+#[derive(Debug)]
+pub(super) struct PythonFileReader {
+    reader: Py<PyAny>,
+    name: Option<String>,
 }
 
-impl AsRef<[u8]> for FileBytes {
-    fn as_ref(&self) -> &[u8] {
-        match self {
-            Self::Python(bytes) => bytes,
-            Self::Native(bytes) => bytes,
-        }
+impl PythonFileReader {
+    pub(super) fn read(&self, py: Python<'_>) -> PyResult<OcrFileContent> {
+        let value = self.reader.bind(py).call0()?;
+        let bytes = if value.is_instance_of::<PyString>() {
+            Bytes::from(value.extract::<String>()?)
+        } else if value.is_instance_of::<PyBytes>() {
+            extract_bytes(&value)?
+        } else {
+            return Err(PyTypeError::new_err(format!(
+                "OCR file read must return bytes or str, got {}",
+                value.get_type(),
+            )));
+        };
+        Ok(OcrFileContent {
+            bytes,
+            file_name: self.name.clone(),
+        })
+    }
+
+    pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.reader)
     }
 }
 
-fn read_file_input(
-    py: Python<'_>,
-    file: &Bound<'_, PyAny>,
-) -> PyResult<(FileBytes, Option<String>)> {
-    if file.is_instance_of::<PyString>() {
-        return Err(PyValueError::new_err(
-            "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.",
-        ));
+fn extract_bytes(value: &Bound<'_, PyAny>) -> PyResult<Bytes> {
+    if value.is_exact_instance_of::<PyBytes>() {
+        return Ok(Bytes::from_owner(value.extract::<PyBackedBytes>()?));
     }
-    if file.is_instance(&py.import("os")?.getattr("PathLike")?)? {
-        let path: PathBuf = file.extract()?;
-        let name = path
-            .file_name()
-            .map(|value| value.to_string_lossy().into_owned());
-        let bytes = py
-            .detach(|| {
-                let mut bytes = Vec::new();
-                std::fs::File::open(&path)?
-                    .take(OCR_INLINE_MAX_BYTES as u64 + 1)
-                    .read_to_end(&mut bytes)?;
-                Ok::<_, std::io::Error>(bytes)
-            })
-            .map_err(|error| {
-                if error.kind() == std::io::ErrorKind::NotFound {
-                    PyFileNotFoundError::new_err(format!("File not found: {}", path.display()))
-                } else {
-                    error.into()
-                }
-            })?;
-        return Ok((FileBytes::Native(bytes), name));
-    }
-    if file.is_instance_of::<PyBytes>() {
-        return Ok((FileBytes::Python(file.extract()?), None));
-    }
-    let reader = file
-        .getattr_opt("read")?
-        .filter(|value| value.is_callable());
-    let Some(reader) = reader else {
-        return Err(PyValueError::new_err(format!(
-            "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.",
-            file.get_type(),
-        )));
-    };
-    let name = file
-        .getattr_opt("name")?
-        .filter(|value| !value.is_none())
-        .map(|value| value.extract::<String>())
-        .transpose()?;
-    let value = reader.call0()?;
-    let bytes = if value.is_instance_of::<PyString>() {
-        FileBytes::Native(value.extract::<String>()?.into_bytes())
-    } else if value.is_instance_of::<PyBytes>() {
-        FileBytes::Python(value.extract()?)
-    } else {
-        return Err(PyTypeError::new_err(format!(
-            "OCR file read must return bytes or str, got {}",
-            value.get_type(),
-        )));
-    };
-    Ok((bytes, name))
+    Ok(Bytes::copy_from_slice(
+        value.extract::<PyBackedBytes>()?.as_ref(),
+    ))
 }
 
 pub(super) struct FileDocumentInput {
-    bytes: FileBytes,
-    name: Option<String>,
-    mime_type: Option<String>,
+    pub input: OcrDocumentInput,
+    pub reader: Option<PythonFileReader>,
 }
 
 impl FromPyObject<'_, '_> for FileDocumentInput {
@@ -104,79 +63,79 @@ impl FromPyObject<'_, '_> for FileDocumentInput {
             Err(error) if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) => None,
             Err(error) => return Err(error),
         };
+        let missing = || {
+            PyValueError::new_err(
+                "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes",
+            )
+        };
         let file = document.get_item("file").map_err(|error| {
             if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) {
-                PyValueError::new_err("document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes")
+                missing()
             } else {
                 error
             }
         })?;
         if file.is_none() {
+            return Err(missing());
+        }
+        if file.is_instance_of::<PyString>() {
             return Err(PyValueError::new_err(
-                "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes",
+                "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.",
             ));
         }
-        let (bytes, name) = read_file_input(py, &file)?;
+        if file.is_instance(&py.import("os")?.getattr("PathLike")?)? {
+            return Ok(Self {
+                input: OcrDocumentInput::Path {
+                    path: file.extract::<PathBuf>()?,
+                    mime_type,
+                },
+                reader: None,
+            });
+        }
+        if file.is_instance_of::<PyBytes>() {
+            return Ok(Self {
+                input: OcrDocumentInput::Bytes {
+                    bytes: extract_bytes(&file)?,
+                    file_name: None,
+                    mime_type,
+                },
+                reader: None,
+            });
+        }
+        let reader = file
+            .getattr_opt("read")?
+            .filter(|value| value.is_callable());
+        let Some(reader) = reader else {
+            return Err(PyValueError::new_err(format!(
+                "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.",
+                file.get_type(),
+            )));
+        };
+        let name = file
+            .getattr_opt("name")?
+            .filter(|value| !value.is_none())
+            .map(|value| value.extract::<String>())
+            .transpose()?;
         Ok(Self {
-            bytes,
-            name,
-            mime_type,
+            input: OcrDocumentInput::HostReader { mime_type },
+            reader: Some(PythonFileReader {
+                reader: reader.unbind(),
+                name,
+            }),
         })
     }
 }
 
-pub(super) fn file_document(py: Python<'_>, document: FileDocumentInput) -> PyResult<OcrDocument> {
-    py.detach(|| {
-        encode_file_document(
-            document.bytes.as_ref(),
-            document.name.as_deref(),
-            document.mime_type.as_deref(),
-        )
-    })
-    .map_err(|error| PyValueError::new_err(error.to_string()))
-}
-
-#[pyfunction]
-fn _ocr_file_document(py: Python<'_>, document: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
-    to_py_preserving_errors(py, &file_document(py, document.extract()?)?)
-}
-
-#[pyfunction]
-fn _ocr_mime_type(file_name: &str) -> String {
-    mime_type_for_name(file_name).into()
-}
-
-#[pyfunction]
-#[pyo3(signature = (file_content, file_name=None, content_type=None))]
-fn _ocr_upload_document(
-    py: Python<'_>,
-    file_content: &Bound<'_, PyBytes>,
-    file_name: Option<&str>,
-    content_type: Option<&str>,
-) -> PyResult<Py<PyAny>> {
-    let bytes: PyBackedBytes = file_content.extract()?;
-    let document = py
-        .detach(|| {
-            encode_file_document(
-                &bytes,
-                None,
-                Some(upload_mime_type(file_name, content_type)),
-            )
-        })
-        .map_err(|error| PyValueError::new_err(error.to_string()))?;
-    to_py_preserving_errors(py, &document)
-}
-
-pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    module.add("_OCR_MAX_FILE_BYTES", OCR_INLINE_MAX_BYTES)?;
-    module.add_function(wrap_pyfunction!(_ocr_upload_document, module)?)?;
-    module.add_function(wrap_pyfunction!(_ocr_file_document, module)?)?;
-    module.add_function(wrap_pyfunction!(_ocr_mime_type, module)?)
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
+    use pyo3::types::PyDict;
+
+    fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> {
+        let locals = PyDict::new(py);
+        py.run(source, Some(&locals), Some(&locals)).unwrap();
+        locals
+    }
 
     #[test]
     fn extraction_validates_required_file_and_optional_mime_type() {
@@ -196,69 +155,148 @@ mod tests {
                 let error = document.extract::<FileDocumentInput>().err().unwrap();
                 assert!(error.is_instance_of::<PyTypeError>(py));
             }
-            let document = py.eval(c"{'file': b'abc'}", None, None).unwrap();
+            let error = py
+                .eval(c"{'file': 'scan.pdf'}", None, None)
+                .unwrap()
+                .extract::<FileDocumentInput>()
+                .err()
+                .unwrap();
+            assert!(error.is_instance_of::<PyValueError>(py));
+            assert!(error.to_string().contains("bare str"));
+            let document = py
+                .eval(c"{'file': b'abc', 'mime_type': 'image/png'}", None, None)
+                .unwrap();
             let input: FileDocumentInput = document.extract().unwrap();
-            assert_eq!(input.bytes.as_ref(), b"abc");
-            assert_eq!(input.name, None);
-            assert_eq!(input.mime_type, None);
+            assert!(input.reader.is_none());
+            assert_eq!(
+                input.input,
+                OcrDocumentInput::Bytes {
+                    bytes: b"abc".as_slice().into(),
+                    file_name: None,
+                    mime_type: Some("image/png".into()),
+                }
+            );
         });
     }
 
     #[test]
-    fn extraction_validates_mime_type_before_consuming_file() {
+    fn paths_and_readers_are_projected_without_io() {
         Python::initialize();
         Python::attach(|py| {
-            let locals = PyDict::new(py);
-            py.run(
-                c"class Reader:
+            let locals = eval(
+                py,
+                c"from pathlib import Path
+class Reader:
+    name = 'scan.png'
     def __init__(self):
         self.reads = 0
     def read(self):
         self.reads += 1
         return b'abc'
 reader = Reader()
-document = {'file': reader, 'mime_type': 7}",
-                Some(&locals),
-                Some(&locals),
-            )
-            .unwrap();
+document = {'file': reader, 'mime_type': 7}
+reader_document = {'file': reader}
+path_document = {'file': Path('/nonexistent/ocr-projection-test.pdf'), 'mime_type': 'image/png'}",
+            );
             let document = locals.get_item("document").unwrap().unwrap();
             let error = document.extract::<FileDocumentInput>().err().unwrap();
             assert!(error.is_instance_of::<PyTypeError>(py));
-            let reads: usize = locals
-                .get_item("reader")
-                .unwrap()
-                .unwrap()
-                .getattr("reads")
-                .unwrap()
-                .extract()
-                .unwrap();
-            assert_eq!(reads, 0);
+
+            let document = locals.get_item("reader_document").unwrap().unwrap();
+            let input: FileDocumentInput = document.extract().unwrap();
+            assert_eq!(
+                input.input,
+                OcrDocumentInput::HostReader { mime_type: None }
+            );
+            let reads = || {
+                locals
+                    .get_item("reader")
+                    .unwrap()
+                    .unwrap()
+                    .getattr("reads")
+                    .unwrap()
+                    .extract::<usize>()
+                    .unwrap()
+            };
+            assert_eq!(reads(), 0);
+            let content = input.reader.unwrap().read(py).unwrap();
+            assert_eq!(reads(), 1);
+            assert_eq!(
+                content,
+                OcrFileContent {
+                    bytes: b"abc".as_slice().into(),
+                    file_name: Some("scan.png".into()),
+                }
+            );
+
+            let document = locals.get_item("path_document").unwrap().unwrap();
+            let input: FileDocumentInput = document.extract().unwrap();
+            assert!(input.reader.is_none());
+            assert_eq!(
+                input.input,
+                OcrDocumentInput::Path {
+                    path: PathBuf::from("/nonexistent/ocr-projection-test.pdf"),
+                    mime_type: Some("image/png".into()),
+                }
+            );
         });
     }
 
     #[test]
-    fn extraction_preserves_reader_key_error_identity() {
+    fn reader_results_are_normalized_and_exceptions_keep_their_identity() {
         Python::initialize();
         Python::attach(|py| {
-            let locals = PyDict::new(py);
-            py.run(
+            let locals = eval(
+                py,
                 c"failure = KeyError('reader failed')
-class Reader:
+class Raising:
     def read(self):
         raise failure
-document = {'file': Reader()}",
-                Some(&locals),
-                Some(&locals),
-            )
-            .unwrap();
-            let document = locals.get_item("document").unwrap().unwrap();
-            let error = document.extract::<FileDocumentInput>().err().unwrap();
+class Text:
+    def read(self):
+        return 'héllo'
+class Wrong:
+    def read(self):
+        return 7
+raising = {'file': Raising()}
+text = {'file': Text()}
+wrong = {'file': Wrong()}",
+            );
+            let reader = |name: &str| {
+                locals
+                    .get_item(name)
+                    .unwrap()
+                    .unwrap()
+                    .extract::<FileDocumentInput>()
+                    .unwrap()
+                    .reader
+                    .unwrap()
+            };
+            let error = reader("raising").read(py).unwrap_err();
             assert!(
                 error
                     .value(py)
                     .is(locals.get_item("failure").unwrap().unwrap())
             );
+            assert_eq!(
+                reader("text").read(py).unwrap().bytes.as_ref(),
+                "héllo".as_bytes()
+            );
+            let error = reader("wrong").read(py).unwrap_err();
+            assert!(error.is_instance_of::<PyTypeError>(py));
+            assert!(error.to_string().contains("bytes or str"));
         });
     }
+
+    #[test]
+    fn exact_python_bytes_transfer_without_copying_and_outlive_the_input() {
+        Python::initialize();
+        let (bytes, pointer) = Python::attach(|py| {
+            let value = PyBytes::new(py, b"document bytes");
+            let pointer = value.as_bytes().as_ptr() as usize;
+            (extract_bytes(value.as_any()).unwrap(), pointer)
+        });
+        assert_eq!(bytes.as_ptr() as usize, pointer);
+        assert_eq!(bytes.as_ref(), b"document bytes");
+    }
 }

litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs

@@ -1,4 +1,5 @@
 use litellm_core::ocr::Error;
+use pyo3::exceptions::{PyFileNotFoundError, PyOSError};
 use pyo3::prelude::*;
 
 use crate::errors::{RustUpstreamError, core_error_to_pyerr};
@@ -7,6 +8,12 @@ pub(super) fn to_pyerr(error: Error) -> PyErr {
     let status = error.http_status_code();
     let mapped = match error {
         Error::Http { status, body } => RustUpstreamError::new_err((status, body)),
+        Error::FileRead {
+            path,
+            kind: std::io::ErrorKind::NotFound,
+            ..
+        } => PyFileNotFoundError::new_err(format!("File not found: {}", path.display())),
+        Error::FileRead { message, .. } => PyOSError::new_err(message),
         other => core_error_to_pyerr(other.into()),
     };
     attach_status(mapped, status)

litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs

@@ -66,13 +66,27 @@ impl PythonOcrHost {
                 retained_fields.set_item(name, value)?;
             }
         }
-        retained_fields.set_item("document", &self.projected()?.fields.document)?;
         let projected = self.projected_mut()?;
+        let document = match &projected.fields.document {
+            Some(document) => document.clone_ref(py),
+            None => to_py(py, &request.document)?,
+        };
+        retained_fields.set_item("document", &document)?;
+        projected.fields.document = Some(document);
         projected.retained_fields = Some(retained_fields.unbind());
         projected.pre_call = Some((&request).into());
         Ok(request)
     }
 
+    fn read_document(&self, py: Python<'_>) -> PyResult<litellm_core::ocr::OcrFileContent> {
+        self.projected()?
+            .fields
+            .reader
+            .as_ref()
+            .ok_or_else(missing_state)?
+            .read(py)
+    }
+
     fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult<ResolvedCredential> {
         let provider = self
             .projected()?
@@ -193,7 +207,7 @@ impl PythonRoute for PythonOcrHost {
                 let OcrHostData::Unprojected { request } = &self.data else {
                     return Err(missing_state());
                 };
-                let projected = project_request(py, request.bind(py), self.state.kwargs.bind(py))?;
+                let projected = project_request(request.bind(py), self.state.kwargs.bind(py))?;
                 let has_token_provider = projected.fields.azure_ad_token_provider.is_some();
                 let request = projected.request;
                 self.data = OcrHostData::Projected(Box::new(ProjectedOcrHost {
@@ -205,6 +219,7 @@ impl PythonRoute for PythonOcrHost {
                 }));
                 OcrHostResult::Request(Ok((Box::new(request), has_token_provider)))
             }
+            OcrHostOperation::ReadDocument => OcrHostResult::Document(Ok(self.read_document(py)?)),
             OcrHostOperation::AcquireAzureAdToken => {
                 OcrHostResult::AzureAdToken(Ok(self.acquire_azure_ad_token(py)?))
             }
@@ -258,6 +273,9 @@ impl PythonRoute for PythonOcrHost {
             OcrHostData::Projected(projected) => {
                 visit.call(&projected.fields.boundary_request)?;
                 visit.call(&projected.fields.document)?;
+                if let Some(reader) = &projected.fields.reader {
+                    reader.traverse(visit)?;
+                }
                 visit.call(&projected.fields.api_key)?;
                 if let Some(provider) = &projected.fields.azure_ad_token_provider {
                     provider.traverse(visit)?;

litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs

@@ -9,6 +9,5 @@ use pyo3::prelude::*;
 
 pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
     value::register(module)?;
-    document::register(module)?;
     lifecycle::register(module)
 }

litellm-rust/crates/python-bridge/src/routes/ocr/project.rs

@@ -1,14 +1,15 @@
 use std::sync::Arc;
 
-use litellm_core::ocr::wire::{OcrWireRequest, consumed_optional_params, decode_request};
-use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall};
-use litellm_python_interop::{
-    from_py_preserving_errors as from_py, to_py_preserving_errors as to_py,
+use litellm_core::ocr::wire::{
+    OcrWireRequest, consumed_optional_params, decode_document, decode_request_input,
 };
+use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall, OcrDocumentInput};
+use litellm_python_interop::from_py_preserving_errors as from_py;
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
 use serde_json::{Map, Value};
 
+use super::document::{FileDocumentInput, PythonFileReader};
 use super::errors::to_pyerr as ocr_error_to_pyerr;
 use super::lifecycle::BridgeOcrHooks;
 use crate::auth::{AZURE_AD_TOKEN_PROVIDER, PythonTokenProvider};
@@ -17,7 +18,8 @@ use crate::marshal::{project_optional_fields, python_timeout_seconds, request_in
 
 pub(super) struct ProjectedOcrFields {
     pub boundary_request: Py<PyAny>,
-    pub document: Py<PyAny>,
+    pub document: Option<Py<PyAny>>,
+    pub reader: Option<PythonFileReader>,
     pub api_key: Py<PyAny>,
     pub azure_ad_token_provider: Option<PythonTokenProvider>,
     pub provider: &'static str,
@@ -25,7 +27,7 @@ pub(super) struct ProjectedOcrFields {
 }
 
 pub(super) struct ProjectedOcrCall {
-    pub request: LiteLLMOcrRequest,
+    pub request: LiteLLMOcrRequest<OcrDocumentInput>,
     pub fields: ProjectedOcrFields,
 }
 
@@ -80,12 +82,12 @@ impl<'py> OcrArguments<'_, 'py> {
 }
 
 enum ProjectedDocument {
-    File { wire: Value, retained: Py<PyAny> },
+    File(FileDocumentInput),
     Other { wire: Value, retained: Py<PyAny> },
 }
 
 impl ProjectedDocument {
-    fn project(py: Python<'_>, document: &Bound<'_, PyAny>) -> PyResult<Self> {
+    fn project(document: &Bound<'_, PyAny>) -> PyResult<Self> {
         let kind: String = document.get_item("type")?.extract()?;
         if kind != "file" {
             return Ok(Self::Other {
@@ -93,25 +95,28 @@ impl ProjectedDocument {
                 retained: document.clone().unbind(),
             });
         }
-        let input = document.extract()?;
-        let encoded = super::document::file_document(py, input)?;
-        let wire = serde_json::to_value(encoded)
-            .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?;
-        Ok(Self::File {
-            retained: to_py(py, &wire)?,
-            wire,
-        })
+        Ok(Self::File(document.extract()?))
     }
 
-    fn into_parts(self) -> (Value, Py<PyAny>) {
+    fn into_parts(
+        self,
+    ) -> PyResult<(
+        OcrDocumentInput,
+        Option<Py<PyAny>>,
+        Option<PythonFileReader>,
+    )> {
         match self {
-            Self::File { wire, retained } | Self::Other { wire, retained } => (wire, retained),
+            Self::File(FileDocumentInput { input, reader }) => Ok((input, None, reader)),
+            Self::Other { wire, retained } => Ok((
+                decode_document(wire).map_err(ocr_error_to_pyerr)?.into(),
+                Some(retained),
+                None,
+            )),
         }
     }
 }
 
 pub(super) fn project_request(
-    py: Python<'_>,
     request: &Bound<'_, PyAny>,
     kwargs: &Bound<'_, PyDict>,
 ) -> PyResult<ProjectedOcrCall> {
@@ -119,8 +124,7 @@ pub(super) fn project_request(
     let arguments = OcrArguments { request, kwargs };
     let model = arguments.model()?;
     let custom_llm_provider = arguments.custom_llm_provider()?;
-    let (wire_document, retained_document) =
-        ProjectedDocument::project(py, &arguments.document()?)?.into_parts();
+    let document = ProjectedDocument::project(&arguments.document()?)?;
     let api_key = arguments.api_key()?;
     let specs = consumed_optional_params(&model, custom_llm_provider.as_deref())
         .map_err(ocr_error_to_pyerr)?;
@@ -136,9 +140,10 @@ pub(super) fn project_request(
     let azure_ad_token_provider = kwargs
         .get_item("azure_ad_token_provider")?
         .and_then(|provider| PythonTokenProvider::select(provider, AZURE_AD_TOKEN_PROVIDER));
+    let (document, retained_document, reader) = document.into_parts()?;
     let wire = OcrWireRequest {
         model,
-        document: wire_document,
+        document,
         api_key: api_key.extract()?,
         api_base: arguments.api_base()?,
         custom_llm_provider,
@@ -147,13 +152,14 @@ pub(super) fn project_request(
         input_sources,
         timeout_seconds: arguments.timeout_seconds()?,
     };
-    let request = decode_request(wire).map_err(ocr_error_to_pyerr)?;
+    let request = decode_request_input(wire).map_err(ocr_error_to_pyerr)?;
     let provider = request.provider_name();
     Ok(ProjectedOcrCall {
         request: request.with_host_hooks(Arc::new(BridgeOcrHooks), None),
         fields: ProjectedOcrFields {
             boundary_request,
             document: retained_document,
+            reader,
             api_key: api_key.unbind(),
             azure_ad_token_provider,
             provider,
@@ -197,10 +203,21 @@ mod tests {
     }
 
     fn project_document(
-        py: Python<'_>,
         document: &Bound<'_, PyAny>,
-    ) -> PyResult<(Value, Py<PyAny>)> {
-        ProjectedDocument::project(py, document).map(ProjectedDocument::into_parts)
+    ) -> PyResult<(
+        OcrDocumentInput,
+        Option<Py<PyAny>>,
+        Option<PythonFileReader>,
+    )> {
+        ProjectedDocument::project(document)?.into_parts()
+    }
+
+    fn url_document(url: &str) -> OcrDocumentInput {
+        litellm_core::ocr::OcrDocument::DocumentUrl {
+            document_url: url.into(),
+            extra_fields: Map::new(),
+        }
+        .into()
     }
 
     fn stub_timeout_conversion(py: Python<'_>) {
@@ -374,7 +391,7 @@ kwargs = {}
     }
 
     #[test]
-    fn document_reader_mutations_are_visible_to_later_field_reads() {
+    fn document_readers_are_not_consumed_during_projection() {
         Python::initialize();
         Python::attach(|py| {
             stub_timeout_conversion(py);
@@ -406,7 +423,12 @@ kwargs = {}
                 .unwrap();
             let arguments = arguments(&request, &kwargs);
             let document = arguments.document().unwrap();
-            project_document(py, &document).unwrap();
+            let (input, retained, reader) = project_document(&document).unwrap();
+            assert_eq!(input, OcrDocumentInput::HostReader { mime_type: None });
+            assert!(retained.is_none());
+            assert_eq!(arguments.api_base().unwrap().as_deref(), Some("original"));
+            assert_eq!(arguments.timeout_seconds().unwrap(), Some(1.0));
+            reader.unwrap().read(py).unwrap();
             assert_eq!(arguments.api_base().unwrap().as_deref(), Some("mutated"));
             assert_eq!(arguments.timeout_seconds().unwrap(), Some(9.0));
         });
@@ -444,7 +466,7 @@ kwargs = {'api_key': key}
     }
 
     #[test]
-    fn file_documents_are_encoded_and_other_documents_keep_the_python_object() {
+    fn file_documents_become_typed_inputs_and_other_documents_keep_the_python_object() {
         Python::initialize();
         Python::attach(|py| {
             let file = py
@@ -454,13 +476,17 @@ kwargs = {'api_key': key}
                     None,
                 )
                 .unwrap();
+            let (input, retained, reader) = project_document(&file).unwrap();
             assert_eq!(
-                project_document(py, &file).unwrap().0,
-                serde_json::json!({
-                    "type": "document_url",
-                    "document_url": "data:application/pdf;base64,JVBERi0xLjQ=",
-                })
+                input,
+                OcrDocumentInput::Bytes {
+                    bytes: b"%PDF-1.4".as_slice().into(),
+                    file_name: None,
+                    mime_type: Some("application/pdf".into()),
+                }
             );
+            assert!(retained.is_none());
+            assert!(reader.is_none());
 
             let original = py
                 .eval(
@@ -469,44 +495,21 @@ kwargs = {'api_key': key}
                     None,
                 )
                 .unwrap();
-            let (wire, retained) = project_document(py, &original).unwrap();
-            assert_eq!(
-                wire,
-                serde_json::json!({
-                    "type": "document_url",
-                    "document_url": "https://example.com/a.pdf",
-                })
-            );
-            assert!(retained.bind(py).is(&original));
+            let (input, retained, _) = project_document(&original).unwrap();
+            assert_eq!(input, url_document("https://example.com/a.pdf"));
+            assert!(retained.unwrap().bind(py).is(&original));
         });
     }
 
     #[test]
-    fn unknown_document_types_reach_existing_downstream_validation() {
+    fn unknown_document_types_reach_existing_core_validation() {
         Python::initialize();
         Python::attach(|py| {
             let document = py
                 .eval(c"{'type': 'mystery', 'mystery': 'x'}", None, None)
                 .unwrap();
-            let wire_document = project_document(py, &document).unwrap().0;
-            assert_eq!(
-                wire_document,
-                serde_json::json!({"type": "mystery", "mystery": "x"})
-            );
-            let error = match decode_request(OcrWireRequest {
-                model: "mistral/mistral-ocr-latest".into(),
-                document: wire_document,
-                api_key: None,
-                api_base: None,
-                custom_llm_provider: None,
-                extra_headers: None,
-                optional_params: Map::new(),
-                input_sources: Default::default(),
-                timeout_seconds: None,
-            }) {
-                Ok(_) => panic!("unknown discriminators belong to core validation"),
-                Err(error) => error,
-            };
+            let error = project_document(&document).unwrap_err();
+            assert!(error.is_instance_of::<PyValueError>(py));
             assert!(error.to_string().contains("document"));
         });
     }
@@ -517,14 +520,14 @@ kwargs = {'api_key': key}
         Python::attach(|py| {
             let missing = py.eval(c"{}", None, None).unwrap();
             assert!(
-                project_document(py, &missing)
+                project_document(&missing)
                     .unwrap_err()
                     .is_instance_of::<PyKeyError>(py)
             );
 
             let non_string = py.eval(c"{'type': 1}", None, None).unwrap();
             assert!(
-                project_document(py, &non_string)
+                project_document(&non_string)
                     .unwrap_err()
                     .is_instance_of::<PyTypeError>(py)
             );
@@ -540,7 +543,7 @@ document = Document()
 ",
             );
             let error =
-                project_document(py, &locals.get_item("document").unwrap().unwrap()).unwrap_err();
+                project_document(&locals.get_item("document").unwrap().unwrap()).unwrap_err();
             assert!(
                 error
                     .value(py)
@@ -569,9 +572,9 @@ document = Document()
 ",
             );
             let document = locals.get_item("document").unwrap().unwrap();
-            let (wire, retained) = project_document(py, &document).unwrap();
-            assert_eq!(wire["type"], "document_url");
-            assert!(!retained.bind(py).is(&document));
+            let (input, retained, _) = project_document(&document).unwrap();
+            assert!(matches!(input, OcrDocumentInput::Bytes { .. }));
+            assert!(retained.is_none());
             let reads: Vec<String> = document.getattr("reads").unwrap().extract().unwrap();
             assert_eq!(reads, ["type", "mime_type", "file"]);
         });

litellm/constants.py

@@ -40,6 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
         "router_general_settings",
         "ignore_invalid_deployments",
         "fallback_access_check",
+        "fallback_budget_check",
         "auto_router_capability_limit",
     }
 )
@@ -53,6 +54,7 @@ S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64
 S3_PREFIX_DIGEST_CHARS: Final = 16
 # s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against
 MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
+S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY"
 MAX_FILE_LIST_LIMIT: Final = 10000
 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
 DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
@@ -1499,6 +1501,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to
 CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
 CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
 ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
+ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
 INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
 SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
 SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"

litellm/integrations/callback_configs.json

@@ -446,6 +446,12 @@
         "ui_name": "S3 Path Prefix",
         "description": "Path prefix within the bucket for organizing logs",
         "required": false
+      },
+      "s3_log_prompts_only": {
+        "type": "boolean",
+        "ui_name": "Log Prompts Only",
+        "description": "Log request messages to S3 but drop the model response from each logged object",
+        "required": false
       }
     },
     "description": "S3 Bucket (AWS) Logging Integration"

litellm/integrations/custom_logger.py

@@ -118,6 +118,7 @@ class CustomLogger:  # https://docs.litellm.ai/docs/observability/custom_callbac
 
         alias_map: Final = {
             "langfuse_otel": "langfuse",
+            "s3_v2": "s3",
         }
         lookup_name: Final = alias_map.get(normalized_name, normalized_name)
 

litellm/integrations/opentelemetry.py

@@ -2965,16 +2965,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
         )
 
         propagator: Final = TraceContextTextMapPropagator()
-        carrier: Final = {"traceparent": _traceparent}
+        carrier: Final = {key: headers[key] for key in ("traceparent", "tracestate") if headers.get(key) is not None}
         _parent_context: Final = propagator.extract(carrier=carrier)
 
         return _parent_context
 
     def _get_span_context(self, kwargs, default_span: Span | None = None):
         from opentelemetry import context, trace
-        from opentelemetry.trace.propagation.tracecontext import (
-            TraceContextTextMapPropagator,
-        )
 
         litellm_params: Final = kwargs.get("litellm_params", {}) or {}
         proxy_server_request: Final = litellm_params.get("proxy_server_request", {}) or {}
@@ -2998,11 +2995,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
         # Priority 2: HTTP traceparent header
         if traceparent is not None:
             verbose_logger.debug("OpenTelemetry: Using traceparent header for context propagation")
-            carrier: Final = {"traceparent": traceparent}
-            return (
-                TraceContextTextMapPropagator().extract(carrier=carrier),
-                None,
-            )
+            return self.get_traceparent_from_header(headers=headers), None
 
         # Priority 3: Active span from global context (auto-detection)
         try:

litellm/integrations/otel/plumbing/context.py

@@ -26,6 +26,7 @@ if TYPE_CHECKING:
     from litellm.integrations.otel.model.destination import OtelDestination
 
 _PROPAGATOR: Final = TraceContextTextMapPropagator()
+_W3C_TRACE_HEADERS: Final = frozenset(("traceparent", "tracestate"))
 
 # The request's root span — the FastAPI-owned SERVER span — captured ONCE when the
 # proxy first resolves it, so request-level spans (the LLM call, guardrails) can
@@ -310,6 +311,37 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None:
     return _PROPAGATOR.extract(carrier)
 
 
+def _outgoing_trace_context(parent_span: object) -> Context | None:
+    if isinstance(parent_span, Span) and is_recordable_span(parent_span):
+        return context_from_span(parent_span)
+
+    root: Final = request_root_span()
+    if root is not None:
+        return context_from_span(root)
+
+    current: Final = get_current()
+    if is_recordable_span(get_current_span(current)):
+        return current
+    return None
+
+
+def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]:
+    """``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span.
+
+    Parent preference: ``parent_span`` (the request span auth stashed on the key), then
+    the anchored request root span, then the ambient active span. Only trace context is
+    injected, never Baggage. Unchanged when no valid span exists anywhere.
+    """
+    context: Final = _outgoing_trace_context(parent_span)
+    if context is None:
+        return dict(headers)  # mutable-ok: OpenTelemetry propagator requires a mutable carrier
+    carrier: Final = {  # mutable-ok: OpenTelemetry propagator requires a mutable carrier
+        key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS
+    }
+    _PROPAGATOR.inject(carrier, context=context)
+    return carrier
+
+
 # The OTLP destinations this request's key or team pointed its traces at, resolved
 # once during auth. A ``ContextVar`` for the same reason the root span above is one:
 # it rides the request task's context into the ``asyncio.create_task`` children that

litellm/integrations/s3.py

@@ -2,19 +2,42 @@
 #    On success + failure, log events to Supabase
 
 import hashlib
+import os
+from collections.abc import Mapping
 from datetime import datetime
 from typing import Final, cast
 
+from pydantic import TypeAdapter, ValidationError
+
 import litellm
 from litellm._logging import print_verbose, verbose_logger
 from litellm.constants import (
     MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES,
     MAX_S3_OBJECT_KEY_BYTES,
     S3_BOUNDED_OBJECT_KEY_HEAD_BYTES,
+    S3_LOG_PROMPTS_ONLY_ENV_VAR,
     S3_PREFIX_DIGEST_CHARS,
 )
 from litellm.types.utils import StandardLoggingPayload
 
+_S3_LOG_PROMPTS_ONLY: Final = TypeAdapter(bool)
+
+
+def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] | None = None) -> bool:
+    env: Final = os.environ if environ is None else environ
+    raw: Final = env.get(S3_LOG_PROMPTS_ONLY_ENV_VAR) if configured is None else configured
+    if raw is None or raw == "":
+        return False
+    try:
+        return _S3_LOG_PROMPTS_ONLY.validate_python(raw.strip() if isinstance(raw, str) else raw)
+    except ValidationError:
+        verbose_logger.warning("s3 logging: s3_log_prompts_only=%r is not a boolean, logging prompts only", raw)
+        return True
+
+
+def prompts_only_payload(payload: StandardLoggingPayload) -> StandardLoggingPayload:
+    return {**payload, "response": None}
+
 
 class S3Logger:
     # Class variables or attributes
@@ -33,6 +56,7 @@ class S3Logger:
         s3_config=None,
         s3_server_side_encryption: str | None = None,
         s3_sse_kms_key_id: str | None = None,
+        s3_log_prompts_only: bool | None = None,
         **kwargs,
     ):
         import boto3
@@ -41,29 +65,30 @@ class S3Logger:
             verbose_logger.debug("in init s3 logger - s3_callback_params %s", litellm.s3_callback_params)
 
             s3_use_team_prefix = False
+            params: Final = {
+                key: litellm.get_secret(value) if isinstance(value, str) and value.startswith("os.environ/") else value
+                for key, value in (litellm.s3_callback_params or {}).items()
+            }
 
             if litellm.s3_callback_params is not None:
-                # read in .env variables - example os.environ/AWS_BUCKET_NAME
-                for key, value in litellm.s3_callback_params.items():
-                    if isinstance(value, str) and value.startswith("os.environ/"):
-                        litellm.s3_callback_params[key] = litellm.get_secret(value)
-                # now set s3 params from litellm.s3_logger_params
-                s3_bucket_name = litellm.s3_callback_params.get("s3_bucket_name")
-                s3_region_name = litellm.s3_callback_params.get("s3_region_name")
-                s3_api_version = litellm.s3_callback_params.get("s3_api_version")
-                s3_use_ssl = litellm.s3_callback_params.get("s3_use_ssl", True)
-                s3_verify = litellm.s3_callback_params.get("s3_verify")
-                s3_endpoint_url = litellm.s3_callback_params.get("s3_endpoint_url")
-                s3_aws_access_key_id = litellm.s3_callback_params.get("s3_aws_access_key_id")
-                s3_aws_secret_access_key = litellm.s3_callback_params.get("s3_aws_secret_access_key")
-                s3_aws_session_token = litellm.s3_callback_params.get("s3_aws_session_token")
-                s3_config = litellm.s3_callback_params.get("s3_config")
-                s3_path = litellm.s3_callback_params.get("s3_path")
-                s3_server_side_encryption = litellm.s3_callback_params.get("s3_server_side_encryption")
-                s3_sse_kms_key_id = litellm.s3_callback_params.get("s3_sse_kms_key_id")
-                # done reading litellm.s3_callback_params
-                s3_use_team_prefix = bool(litellm.s3_callback_params.get("s3_use_team_prefix", False))
+                s3_bucket_name = params.get("s3_bucket_name")
+                s3_region_name = params.get("s3_region_name")
+                s3_api_version = params.get("s3_api_version")
+                s3_use_ssl = params.get("s3_use_ssl", True)
+                s3_verify = params.get("s3_verify")
+                s3_endpoint_url = params.get("s3_endpoint_url")
+                s3_aws_access_key_id = params.get("s3_aws_access_key_id")
+                s3_aws_secret_access_key = params.get("s3_aws_secret_access_key")
+                s3_aws_session_token = params.get("s3_aws_session_token")
+                s3_config = params.get("s3_config")
+                s3_path = params.get("s3_path")
+                s3_server_side_encryption = params.get("s3_server_side_encryption")
+                s3_sse_kms_key_id = params.get("s3_sse_kms_key_id")
+                s3_use_team_prefix = bool(params.get("s3_use_team_prefix", False))
             self.s3_use_team_prefix = s3_use_team_prefix
+            self.s3_log_prompts_only: object = (
+                params.get("s3_log_prompts_only") if s3_log_prompts_only is None else s3_log_prompts_only
+            )
             self.bucket_name = s3_bucket_name
             self.s3_path = s3_path
             self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params(
@@ -144,7 +169,9 @@ class S3Logger:
 
             from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
-            payload_str: Final = safe_dumps(payload)
+            payload_str: Final = safe_dumps(
+                prompts_only_payload(payload) if resolve_s3_log_prompts_only(self.s3_log_prompts_only) else payload
+            )
 
             print_verbose(f"\ns3 Logger - Logging payload = {payload_str}")
 

litellm/integrations/s3_v2.py

@@ -21,6 +21,8 @@ from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_S
 from litellm.integrations.s3 import (
     get_s3_object_download_filename,
     get_s3_object_key,
+    prompts_only_payload,
+    resolve_s3_log_prompts_only,
     resolve_sse_params,
 )
 from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
@@ -68,6 +70,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
         s3_use_virtual_hosted_style: bool = False,
         s3_server_side_encryption: str | None = None,
         s3_sse_kms_key_id: str | None = None,
+        s3_log_prompts_only: bool | None = None,
         s3_callback_params_override: dict | None = None,
         **kwargs,
     ):
@@ -108,6 +111,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
                 s3_use_virtual_hosted_style=s3_use_virtual_hosted_style,
                 s3_server_side_encryption=s3_server_side_encryption,
                 s3_sse_kms_key_id=s3_sse_kms_key_id,
+                s3_log_prompts_only=s3_log_prompts_only,
             )
             verbose_logger.debug("s3 logger using endpoint url %s", s3_endpoint_url)
 
@@ -163,6 +167,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
         s3_use_virtual_hosted_style: bool = False,
         s3_server_side_encryption: str | None = None,
         s3_sse_kms_key_id: str | None = None,
+        s3_log_prompts_only: bool | None = None,
         params_source: dict | None = None,
     ):
         """
@@ -212,6 +217,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
             bool(params.get("s3_use_virtual_hosted_style", False)) or s3_use_virtual_hosted_style
         )
 
+        self.s3_log_prompts_only: object = (
+            params.get("s3_log_prompts_only") if s3_log_prompts_only is None else s3_log_prompts_only
+        )
+
         self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params(
             params.get("s3_server_side_encryption") or s3_server_side_encryption,
             params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id,
@@ -489,8 +498,13 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
 
         s3_object_download_filename: Final = get_s3_object_download_filename(start_time, standard_logging_payload["id"])
 
+        payload: Final = (
+            prompts_only_payload(standard_logging_payload)
+            if resolve_s3_log_prompts_only(self.s3_log_prompts_only)
+            else standard_logging_payload
+        )
         return s3BatchLoggingElement(
-            payload=dict(standard_logging_payload),
+            payload=dict(payload),
             s3_object_key=s3_object_key,
             s3_object_download_filename=s3_object_download_filename,
         )

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -1561,10 +1561,12 @@ class AnthropicMessagesHandler(BaseTranslation):
 
     def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
         stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
+        tool_use_fingerprints: Final = self._streamed_tool_use_fingerprints(responses_so_far)
         return StreamingScanKey(
             texts=(self.get_streaming_string_so_far(responses_so_far),),
-            tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (),
+            tool_calls=tool_use_fingerprints if stream_ended else (),
             stream_ended=stream_ended,
+            tool_calls_in_flight=bool(tool_use_fingerprints) and not stream_ended,
         )
 
     @classmethod

litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py

@@ -78,6 +78,7 @@ _PROPAGATED_METADATA_KEYS: Final = (
     "user_api_key_end_user_id",
     "user_api_end_user_max_budget",
     "user_api_key_model_max_budget",
+    "user_api_key_team_model_max_budget",
     "user_api_key_user_model_max_budget",
     "user_api_key_end_user_model_max_budget",
     "litellm_call_id",
@@ -395,9 +396,9 @@ async def _check_summary_model_budget(
     ``user_api_key_auth`` runs for the client-requested model. Returns True outside the proxy or when no
     per-model budget is configured.
 
-    All three scopes are checked because the summary's spend is charged to all
-    three: this file propagates the key, user and end-user budgets into the
-    subrequest's metadata, so enforcing only two of them would let compaction
+    Every scope is checked because the summary's spend is charged to every
+    scope: this file propagates the key, team, user and end-user budgets into the
+    subrequest's metadata, so skipping one of them would let compaction
     increment a counter it can never be refused by.
     """
     if user_api_key_auth is None:
@@ -444,6 +445,26 @@ async def _check_summary_model_budget(
             )
             return False
 
+    team_model_max_budget: Final = user_api_key_auth.team_model_max_budget
+    team_id: Final = user_api_key_auth.team_id
+    if isinstance(team_model_max_budget, dict) and team_model_max_budget and team_id is not None:
+        try:
+            await model_max_budget_limiter.is_team_within_model_budget(
+                team_id=team_id,
+                team_model_max_budget=team_model_max_budget,
+                key_model_max_budget=model_max_budget if isinstance(model_max_budget, dict) else None,
+                model=summary_model,
+            )
+        except litellm.BudgetExceededError:
+            return False
+        except Exception as e:  # noqa: BLE001  # a budget gate denies on any failure, as the other scopes do
+            verbose_logger.warning(
+                "compact_20260112: unexpected error during team model-budget check for summary_model=%s; denying: %s",
+                summary_model,
+                e,
+            )
+            return False
+
     end_user_model_max_budget: Final[dict[str, object] | None] = getattr(
         user_api_key_auth, "end_user_model_max_budget", None
     )

litellm/llms/base_llm/guardrail_translation/base_translation.py

@@ -40,11 +40,15 @@ class StreamingScanKey:
     """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys
     compare equal when the round would scan the same content again; ``stream_ended``
     stays out of the comparison and only says whether the handler is on its
-    end-of-stream path, where an empty payload is still scanned today."""
+    end-of-stream path, where an empty payload is still scanned today.
+    ``tool_calls_in_flight`` also stays out of the comparison: it flags that tool
+    calls have streamed which this round cannot scan yet, so a buffered window
+    holding them must stay withheld until the end-of-stream scan covers them."""
 
     texts: tuple[str, ...]
     tool_calls: tuple[str, ...] = ()
     stream_ended: bool = field(default=False, compare=False)
+    tool_calls_in_flight: bool = field(default=False, compare=False)
 
     @property
     def has_nothing_to_scan(self) -> bool:

litellm/llms/openai/chat/guardrail_translation/handler.py

@@ -792,10 +792,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
         chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream))
         stream_ended: Final = self._first_choice_has_finished(responses_so_far)
+        tool_call_fingerprints: Final = self._streamed_tool_call_fingerprints(responses_so_far)
         return StreamingScanKey(
             texts=tuple(self._combine_streaming_texts(chunks).values()),
-            tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (),
+            tool_calls=tool_call_fingerprints if stream_ended else (),
             stream_ended=stream_ended,
+            tool_calls_in_flight=bool(tool_call_fingerprints) and not stream_ended,
         )
 
     @staticmethod
@@ -804,7 +806,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             stream_item_fingerprint(tool_call)
             for chunk in responses_so_far
             for choice in _stream_chunk_choices(chunk)
-            for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls")
+            for tool_call in _streamed_delta_tool_calls(stream_item_field(choice, "delta"))
         )
 
     @staticmethod
@@ -1342,6 +1344,12 @@ def _stream_chunk_choices(item: object) -> Sequence[object]:
     return ()
 
 
+def _streamed_delta_tool_calls(delta: object) -> tuple[object, ...]:
+    function_call: Final = stream_item_field(delta, "function_call")
+    legacy: Final = () if function_call is None else (function_call,)
+    return stream_item_items(delta, "tool_calls") + legacy
+
+
 def _blocked_stream_identity(
     exc: "ModifyResponseException", responses_so_far: Sequence[object]
 ) -> tuple[str, int, str]:

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -1175,11 +1175,22 @@ class OpenAIResponsesHandler(BaseTranslation):
         last_event_type: Final = stream_item_field(last_event, "type")
         if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value:
             return None
-        if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value:
+        if last_event_type in _TERMINAL_ENVELOPE_EVENT_TYPES:
             return self._completed_response_scan_key(stream_item_field(last_event, "response"))
         return StreamingScanKey(
             texts=(self.get_streaming_string_so_far(responses_so_far),),
-            stream_ended=self._check_streaming_has_ended(responses_so_far),
+            tool_calls_in_flight=self._has_streamed_tool_call_events(responses_so_far),
+        )
+
+    @staticmethod
+    def _has_streamed_tool_call_events(responses_so_far: Sequence[object]) -> bool:
+        return any(
+            stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
+            or (
+                stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES
+                and stream_item_field(stream_item_field(event, "item"), "type") in _TOOL_CALL_ITEM_TYPES
+            )
+            for event in responses_so_far
         )
 
     @staticmethod

litellm/models/verification_token.py

@@ -18,6 +18,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
     key_name: str | None = None
     key_alias: str | None = None
     spend: float = 0.0
+    total_spend: float = 0.0
     max_budget: float | None = None
     expires: str | datetime | None = None
     models: list = []
@@ -69,6 +70,7 @@ class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken):
     """Audit record for deleted keys; mirrors the token plus deletion metadata."""
 
     id: str | None = None
+    organization_id: str | None = None
     deleted_at: datetime | None = None
     deleted_by: str | None = None
     deleted_by_api_key: str | None = None

litellm/ocr/input.py

@@ -1,112 +0,0 @@
-from collections.abc import Mapping
-from os import PathLike
-from typing import Final, Literal, Protocol, cast  # noqa: TID251  # native callables are validated when loaded
-
-from typing_extensions import NotRequired, ReadOnly, TypedDict
-
-from litellm.rust_bridge.bindings import NativeBinding
-from litellm.rust_bridge.configuration import rust_ocr_enabled
-
-
-class FileReader(Protocol):
-    def read(self) -> bytes | str: ...
-
-
-class FileDocument(TypedDict):
-    type: ReadOnly[Literal["file"]]
-    file: ReadOnly[bytes | PathLike[str] | FileReader]
-    mime_type: ReadOnly[NotRequired[str]]
-
-
-class NativeFileDocument(Protocol):
-    def __call__(self, document: Mapping[str, object]) -> dict[str, str]: ...
-
-
-class NativeUploadDocument(Protocol):
-    def __call__(self, file_content: bytes, file_name: str | None, content_type: str | None) -> dict[str, str]: ...
-
-
-class NativeMimeType(Protocol):
-    def __call__(self, file_name: str) -> str: ...
-
-
-_FILE_DOCUMENT: Final = NativeBinding(
-    "_ocr_file_document",
-    validate=lambda value: (
-        cast(  # cast-ok: native export owns the callable signature
-            NativeFileDocument, value
-        )
-        if callable(value)
-        else None
-    ),
-)
-_UPLOAD_DOCUMENT: Final = NativeBinding(
-    "_ocr_upload_document",
-    validate=lambda value: (
-        cast(  # cast-ok: native export owns the callable signature
-            NativeUploadDocument, value
-        )
-        if callable(value)
-        else None
-    ),
-)
-_MAX_FILE_BYTES: Final = NativeBinding(
-    "_OCR_MAX_FILE_BYTES", validate=lambda value: value if isinstance(value, int) and value > 0 else None
-)
-_MIME_TYPE: Final = NativeBinding(
-    "_ocr_mime_type",
-    validate=lambda value: (
-        cast(  # cast-ok: native export owns the callable signature
-            NativeMimeType, value
-        )
-        if callable(value)
-        else None
-    ),
-)
-_PYTHON_MAX_FILE_BYTES: Final = 50 * 1024 * 1024
-
-
-def get_mime_type(file_path: str) -> str:
-    native: Final = _MIME_TYPE.load() if rust_ocr_enabled() else None
-    if native is None:
-        from litellm.ocr import legacy
-
-        return legacy.get_mime_type(file_path)
-    return native(file_path)
-
-
-def get_max_file_bytes() -> int:
-    limit: Final = _MAX_FILE_BYTES.load() if rust_ocr_enabled() else None
-    if limit is None:
-        return _PYTHON_MAX_FILE_BYTES
-    return limit
-
-
-def convert_file_document_to_url_document(document: FileDocument) -> dict[str, str]:
-    native: Final = _FILE_DOCUMENT.load() if rust_ocr_enabled() else None
-    if native is None:
-        from litellm.ocr import legacy
-
-        return legacy.convert_file_document_to_url_document(document)
-    return native(document)
-
-
-def convert_upload_to_url_document(
-    file_content: bytes, filename: str | None, content_type: str | None
-) -> dict[str, str]:
-    native: Final = _UPLOAD_DOCUMENT.load() if rust_ocr_enabled() else None
-    if native is None:
-        from litellm.ocr import legacy
-
-        if len(file_content) > _PYTHON_MAX_FILE_BYTES:
-            raise ValueError("OCR file exceeds the size limit")
-        content_mime: Final = content_type.split(";")[0].strip() if content_type else None
-        mime_type: Final = (
-            legacy.get_mime_type(filename)
-            if filename and (not content_mime or content_mime == "application/octet-stream")
-            else content_mime or "application/octet-stream"
-        )
-        return legacy.convert_file_document_to_url_document(
-            {"type": "file", "file": file_content, "mime_type": mime_type}
-        )
-    return native(file_content, filename, content_type)

litellm/ocr/legacy.py

@@ -11,7 +11,7 @@ from collections.abc import Coroutine, Mapping
 from dataclasses import dataclass
 from io import IOBase
 from types import MappingProxyType
-from typing import Final, cast  # noqa: TID251  # adapters preserve the legacy untyped contracts
+from typing import Final, Protocol, cast  # noqa: TID251  # adapters preserve the legacy untyped contracts
 
 import httpx
 
@@ -26,7 +26,6 @@ from litellm.llms.base_llm.ocr.transformation import (
     parse_ocr_request_format,
 )
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
-from litellm.ocr.input import FileReader
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import CustomPricingLiteLLMParams
 from litellm.utils import ProviderConfigManager, client
@@ -34,6 +33,10 @@ from litellm.utils import ProviderConfigManager, client
 base_llm_http_handler: Final = BaseLLMHTTPHandler()
 
 
+class FileReader(Protocol):
+    def read(self) -> bytes | str: ...
+
+
 @dataclass(frozen=True, slots=True)
 class _PreparedOCRRequest:
     model: str

litellm/ocr/main.py

@@ -5,7 +5,7 @@ import httpx
 
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
 from litellm.ocr import legacy
-from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type
+from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type
 from litellm.rust_bridge.bindings import native_exception_types
 from litellm.rust_bridge.configuration import rust_ocr_enabled
 from litellm.rust_bridge.ocr import LiteLLMOcrRequest

litellm/proxy/_types.py

@@ -2004,6 +2004,13 @@ RouterSettingsDict = Annotated[
 class NewTeamRequest(TeamBase):
     router_settings: RouterSettingsDict | None = None
     model_aliases: dict | None = None
+    model_max_budget: GenericBudgetConfigType | None = Field(
+        default=None,
+        description=(
+            "Max budget per model for every key on the team, overridable per key "
+            "(e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}})"
+        ),
+    )
     tags: list | None = None
     guardrails: list[str] | None = None
     policies: list[str] | None = None
@@ -2105,6 +2112,13 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
     access_group_ids: list[str] | None = None
     budget_limits: list[BudgetLimitEntry] | None = None  # multiple concurrent budget windows
     default_team_member_models: list[str] | None = None  # default allowed_models seeded onto new team members
+    model_max_budget: GenericBudgetConfigType | None = Field(
+        default=None,
+        description=(
+            "Max budget per model for every key on the team, overridable per key "
+            "(e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}})"
+        ),
+    )
 
 
 class PatchTeamRequest(UpdateTeamRequest):
@@ -3032,6 +3046,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
     team_tpd_limit: int | None = None
     team_max_budget: float | None = None
     team_soft_budget: float | None = None
+    team_model_max_budget: dict[str, object] | None = None
     team_models: list = []
     team_blocked: bool = False
     soft_budget: float | None = None
@@ -3710,6 +3725,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
             "AWS_ACCESS_KEY_ID",
             "AWS_SECRET_ACCESS_KEY",
             "AWS_REGION_NAME",
+            "S3_LOG_PROMPTS_ONLY",
         ],
     )
 
@@ -4462,6 +4478,7 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
     # Parent org's model ceiling, reported only to callers who can manage the team.
     # None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.
     organization_models: list[str] | None = None
+    model_max_budget_usage: Mapping[str, Mapping[str, object]] | None = None
 
 
 class TeamInfoResponseObject(TypedDict):

litellm/proxy/auth/fallback_budget.py

@@ -0,0 +1,166 @@
+"""
+Enforce the caller's budget against router fallback targets.
+
+Budget is checked once, during auth, against the *requested* model group. A zero-cost group takes
+`_is_model_cost_zero`'s bypass and waives every budget check; the router then picks a fallback
+target after auth, inside `run_async_fallback`, and nothing re-checks budget on the group that
+actually bills. So a free model with a paid fallback spends without a gate.
+
+This predicate is injected into the router to re-check budget for each fallback target before it is
+attempted, mirroring `fallback_model_access.py`. It deliberately leaves the primary attempt alone:
+a zero-cost model is never blocked by budget, and only the paid fallback is refused. On by default;
+set `general_settings.enforce_fallback_budget: false` to restore the unguarded behaviour.
+
+Scope: the key's and the user's `max_budget`. Not covered yet, and each needs a read-only evaluation
+path before it can be: team, team-member, end-user, org, global and per-model budgets, whose
+auth-path functions enforce rather than report (they raise), so reusing them would fire threshold
+alerts and take spend reservations for a target that is then skipped; and the key's rolling
+`budget_limits` windows, whose accumulated spend lives only in per-window counters
+(`spend:key:{token}:window:{budget_duration}`), so enforcing them means more counter reads on the
+fallback path rather than reusing state auth already loaded.
+
+Two known limitations of that narrow scope, both shared with `fallback_model_access.py`:
+
+* This reads the spend counter, it does not reserve against it. Requests already in flight all
+  observe the same pre-billing figure, so a cap can be crossed by roughly the number of concurrent
+  fallbacks times their cost. Auth-time enforcement avoids this by pre-filling the counter through
+  `reserve_budget_for_request`, which the zero-cost bypass skips. Turning the soft cap into a hard
+  one means reserving per fallback attempt and reconciling on completion.
+* A request that reaches the router without `metadata["user_api_key_auth"]` is not restricted.
+  Only `add_litellm_data_to_request` populates that key, so endpoints that assemble metadata by
+  hand (for example `/queue/chat/completions`) fall through as unauthenticated.
+"""
+
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Final
+
+from pydantic import BaseModel, ValidationError
+
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.auth.auth_checks import (
+    _is_model_cost_zero,  # pyright: ignore[reportPrivateUsage]  # the zero-cost predicate the auth-time budget checks use; no public equivalent
+)
+from litellm.router import Router
+
+
+class _RequestMetadata(BaseModel):
+    user_api_key_auth: UserAPIKeyAuth | None = None
+
+
+class _FallbackBudgetSettings(BaseModel):
+    enforce_fallback_budget: bool = True
+
+
+def _token_in_metadata(metadata: object) -> UserAPIKeyAuth | None:
+    try:
+        return _RequestMetadata.model_validate(metadata).user_api_key_auth
+    except ValidationError:
+        return None
+
+
+def _user_api_key_auth_from_request(request_kwargs: Mapping[str, object]) -> UserAPIKeyAuth | None:
+    return next(
+        (
+            token
+            for field in ("metadata", "litellm_metadata")
+            if (token := _token_in_metadata(request_kwargs.get(field))) is not None
+        ),
+        None,
+    )
+
+
+def _enforced_by_general_settings() -> bool:
+    from litellm.proxy.proxy_server import general_settings
+
+    return _FallbackBudgetSettings.model_validate(general_settings).enforce_fallback_budget
+
+
+def _applies_user_budget_to_team_keys() -> bool:
+    from litellm.proxy.proxy_server import general_settings
+
+    return general_settings.get("apply_user_budget_to_team_keys") is True
+
+
+async def _counter_spend(counter_key: str, fallback_spend: float, max_budget: float) -> float:
+    """
+    Read a spend counter the same way the auth-time budget checks do.
+
+    `max_budget` is not advisory: it makes `get_current_spend` re-check the counter against the
+    authoritative recorded spend before admitting. A counter restored from an older Redis snapshot
+    reads as a hit rather than a clean miss, so without this the reseed path never runs and a
+    stale-low counter would keep admitting paid fallbacks past the cap.
+    """
+    from litellm.proxy.proxy_server import get_current_spend
+
+    return await get_current_spend(
+        counter_key=counter_key,
+        fallback_spend=fallback_spend,
+        max_budget=max_budget,
+    )
+
+
+async def is_token_within_budget_for_model(*, model: str, valid_token: UserAPIKeyAuth, llm_router: Router) -> bool:
+    """
+    True when the key and the user behind it can still pay for `model`.
+
+    A zero-cost fallback target is always allowed: refusing it would deny a request on spend some
+    other model accrued, which is the same reasoning behind the auth-time bypass.
+    """
+    if _is_model_cost_zero(model=model, llm_router=llm_router):
+        return True
+
+    key_budget: Final = valid_token.max_budget
+    if key_budget is not None and valid_token.token is not None:
+        key_spend: Final = await _counter_spend(
+            counter_key=f"spend:key:{valid_token.token}",
+            fallback_spend=valid_token.spend or 0.0,
+            max_budget=key_budget,
+        )
+        if key_spend >= key_budget:
+            return False
+
+    # Mirrors `_PROXY_MaxBudgetLimiter`: a team key does not carry the key owner's personal budget
+    # unless the proxy opts in, so the personal cap must not gate the fallback either.
+    user_budget: Final = valid_token.user_max_budget
+    if (
+        user_budget is not None
+        and valid_token.user_id is not None
+        and (valid_token.team_id is None or _applies_user_budget_to_team_keys())
+    ):
+        user_spend: Final = await _counter_spend(
+            counter_key=f"spend:user:{valid_token.user_id}",
+            fallback_spend=valid_token.user_spend or 0.0,
+            max_budget=user_budget,
+        )
+        if user_spend >= user_budget:
+            return False
+
+    return True
+
+
+@dataclass(frozen=True, slots=True)
+class RouterFallbackBudgetCheck:
+    """
+    `FallbackBudgetCheck` for the proxy's router: while `is_enforced()` is true, a paid fallback
+    target is attempted only when the caller is still within budget. Requests that carry no key
+    (for example internal health checks) are not restricted.
+    """
+
+    is_enforced: Callable[[], bool]
+
+    async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool:
+        if not self.is_enforced():
+            return True
+        valid_token: Final = _user_api_key_auth_from_request(request_kwargs)
+        if valid_token is None:
+            return True
+        try:
+            return await is_token_within_budget_for_model(model=model, valid_token=valid_token, llm_router=llm_router)
+        except Exception as e:  # noqa: BLE001  # fail closed: a spend lookup failure must not bill the caller
+            verbose_proxy_logger.warning("Skipping fallback to model=%s: budget lookup failed: %s", model, e)
+            return False
+
+
+router_fallback_budget_check: Final = RouterFallbackBudgetCheck(is_enforced=_enforced_by_general_settings)

litellm/proxy/auth/team_grants.py

@@ -59,6 +59,7 @@ class TeamGrants(TypedDict, total=False):
     team_tpd_limit: ReadOnly[int | None]
     team_max_budget: ReadOnly[float | None]
     team_soft_budget: ReadOnly[float | None]
+    team_model_max_budget: ReadOnly[dict[str, object] | None]
     team_spend: ReadOnly[float | None]
     team_models: ReadOnly[Sequence[str]]
     team_blocked: ReadOnly[bool]
@@ -101,6 +102,7 @@ def team_grants(
         team_tpd_limit=team_object.tpd_limit,
         team_max_budget=team_object.max_budget,
         team_soft_budget=team_object.soft_budget,
+        team_model_max_budget=team_object.model_max_budget,
         team_spend=team_object.spend,
         team_models=tuple(team_object.models),
         team_blocked=team_object.blocked,

litellm/proxy/auth/user_api_key_auth.py

@@ -304,6 +304,16 @@ class _UserModelBudgetLimiter(Protocol):
     ) -> bool: ...
 
 
+class _TeamModelBudgetLimiter(Protocol):
+    async def is_team_within_model_budget(
+        self,
+        team_id: str,
+        team_model_max_budget: Mapping[str, object],
+        key_model_max_budget: Mapping[str, object] | None,
+        model: str,
+    ) -> bool: ...
+
+
 class _TokenTeamModels(Protocol):
     @property
     def team_models(self) -> list[str]: ...
@@ -374,6 +384,25 @@ async def _check_user_model_budget(
         )
 
 
+async def _check_team_model_budget(
+    valid_token: UserAPIKeyAuth,
+    model_max_budget_limiter: _TeamModelBudgetLimiter,
+    models: list[str],
+) -> None:
+    """Enforce the team's `model_max_budget` for every requested model the key does not override."""
+    team_model_max_budget: Final = valid_token.team_model_max_budget
+    if valid_token.team_id is None or not team_model_max_budget:
+        return
+    key_model_max_budget: Final[Mapping[str, object] | None] = valid_token.model_max_budget
+    for model_name in models:
+        await model_max_budget_limiter.is_team_within_model_budget(
+            team_id=valid_token.team_id,
+            team_model_max_budget=team_model_max_budget,
+            key_model_max_budget=key_model_max_budget,
+            model=model_name,
+        )
+
+
 async def _check_key_model_budget_with_fallback(
     valid_token: UserAPIKeyAuth,
     model_max_budget_limiter: _KeyModelBudgetLimiter,
@@ -2376,6 +2405,7 @@ async def _user_api_key_auth_builder(
                         team_id=valid_token.team_id,
                         max_budget=valid_token.team_max_budget,
                         soft_budget=valid_token.team_soft_budget,
+                        model_max_budget=valid_token.team_model_max_budget,
                         spend=valid_token.team_spend,
                         tpm_limit=valid_token.team_tpm_limit,
                         rpm_limit=valid_token.team_rpm_limit,
@@ -2530,6 +2560,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached
         team_id=valid_token.team_id,
         max_budget=valid_token.team_max_budget,
         soft_budget=valid_token.team_soft_budget,
+        model_max_budget=valid_token.team_model_max_budget,
         spend=valid_token.team_spend,
         tpm_limit=valid_token.team_tpm_limit,
         rpm_limit=valid_token.team_rpm_limit,
@@ -2606,6 +2637,7 @@ async def _run_centralized_common_checks(
         litellm_proxy_admin_name,
         llm_router,
         master_key,
+        model_max_budget_limiter,
         prisma_client,
         proxy_logging_obj,
         user_api_key_cache,
@@ -2874,6 +2906,21 @@ async def _run_centralized_common_checks(
     finally:
         release_spend_counter_batch()
 
+    if not skip_budget_checks:
+        await _check_team_model_budget(
+            valid_token=user_api_key_auth_obj,
+            model_max_budget_limiter=model_max_budget_limiter,
+            models=_get_model_names_for_budget_checks(
+                model=_get_model_from_request_context(
+                    request_data=request_data,
+                    route=route,
+                    request=request,
+                    llm_router=llm_router,
+                    team_id=user_api_key_auth_obj.team_id,
+                )
+            ),
+        )
+
     await _reserve_budget_after_common_checks(
         user_api_key_auth_obj=user_api_key_auth_obj,
         request=request,

litellm/proxy/common_utils/proxy_rate_limit_error.py

@@ -11,7 +11,7 @@ exception types:
   an upstream LLM provider returns 429.
 * :class:`fastapi.HTTPException` (status 429) — raised directly by proxy hooks
   such as ``parallel_request_limiter``, ``dynamic_rate_limiter``,
-  ``batch_rate_limiter``, ``max_budget_limiter``, ``max_iterations_limiter``,
+  ``batch_rate_limiter``, ``max_iterations_limiter``,
   etc.
 * :class:`litellm.llms.base_llm.chat.transformation.BaseLLMException` (status
   429) — raised by some provider transports.

litellm/proxy/db/create_views.py

@@ -78,6 +78,7 @@ async def create_missing_views(db: SupportsRawQueries) -> None:
                 v.*,
                 t.spend AS team_spend,
                 t.max_budget AS team_max_budget,
+                t.model_max_budget AS team_model_max_budget,
                 t.tpm_limit AS team_tpm_limit,
                 t.rpm_limit AS team_rpm_limit,
                 t.tpd_limit AS team_tpd_limit,

litellm/proxy/db/db_spend_update_writer.py

@@ -18,6 +18,8 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
 from urllib.parse import quote, unquote
 
+from typing_extensions import ReadOnly, TypedDict
+
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching import RedisCache
@@ -109,6 +111,10 @@ def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool
     return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS})
 
 
+class _SpendIncrement(TypedDict):
+    increment: ReadOnly[float]
+
+
 class _SpendBatch(Protocol):
     litellm_usertable: BatchTable
     litellm_verificationtoken: BatchTable
@@ -1615,10 +1621,12 @@ class DBSpendUpdateWriter:
                         async with transaction.batch_() as batcher:
                             # Sort by token for consistent lock ordering across pods to prevent deadlocks.
                             for token, response_cost in sorted(key_list_transactions.items()):
+                                spend_increment: _SpendIncrement = {"increment": response_cost}
                                 batcher.litellm_verificationtoken.update_many(  # 'update_many' prevents error from being raised if no row exists
                                     where={"token": token},
                                     data={
-                                        "spend": {"increment": response_cost},
+                                        "spend": spend_increment,
+                                        "total_spend": spend_increment,
                                         "last_active": datetime.now(timezone.utc),
                                     },
                                 )

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -37,6 +37,7 @@ from litellm.types.guardrails import (
     ApplyGuardrailResponse,
     BaseLitellmParams,
     BedrockGuardrailConfigModel,
+    BedrockGuardrailStreamingParams,
     Guardrail,
     GuardrailEventHooks,
     GuardrailInfoResponse,
@@ -1959,7 +1960,10 @@ async def get_provider_specific_params():
     ```
     """
     # Get fields from the models
-    bedrock_fields: Final = _get_fields_from_model(BedrockGuardrailConfigModel)
+    bedrock_fields: Final = {
+        **_get_fields_from_model(BedrockGuardrailConfigModel),
+        **_get_fields_from_model(BedrockGuardrailStreamingParams),
+    }
     presidio_fields: Final = _get_fields_from_model(PresidioPresidioConfigModelUserInterface)
     lakera_v2_fields: Final = _get_fields_from_model(LakeraV2GuardrailConfigModel)
     tool_permission_fields: Final = _get_fields_from_model(ToolPermissionGuardrailConfigModel)

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -248,6 +248,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         streaming_buffer_until_moderated: bool | None = None,
         streaming_sampling_rate: int | None = None,
         streaming_end_of_stream_only: bool | None = None,
+        streaming_buffer_release_on_scan: bool | None = None,
         **kwargs,
     ):
         self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
@@ -258,6 +259,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
                         "streaming_buffer_until_moderated": streaming_buffer_until_moderated,
                         "streaming_sampling_rate": streaming_sampling_rate,
                         "streaming_end_of_stream_only": streaming_end_of_stream_only,
+                        "streaming_buffer_release_on_scan": streaming_buffer_release_on_scan,
                     }
                 )
             )
@@ -321,13 +323,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         self.streaming_buffer_until_moderated = streaming_params.streaming_buffer_until_moderated
         self.streaming_sampling_rate = streaming_params.streaming_sampling_rate
         self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only
+        self.streaming_buffer_release_on_scan = streaming_params.streaming_buffer_release_on_scan
 
     def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
         super().update_in_memory_litellm_params(litellm_params)
         self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra))
 
     def _streams_incrementally(self) -> bool:
-        return not self.streaming_buffer_until_moderated and not self.mask_response_content
+        if self.mask_response_content:
+            return False
+        if not self.streaming_buffer_until_moderated:
+            return True
+        return self.streaming_buffer_release_on_scan and not self.streaming_end_of_stream_only
 
     @classmethod
     def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:

litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py

@@ -23,6 +23,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
         event_hook=litellm_params.mode,
         default_on=litellm_params.default_on,
         fail_on_error=litellm_params.fail_on_error,
+        streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated,
+        streaming_buffer_release_on_scan=streaming_params.streaming_buffer_release_on_scan,
         streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
         streaming_sampling_rate=streaming_params.streaming_sampling_rate,
     )

litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py

@@ -260,6 +260,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
         api_key: str | None = None,
         api_base: str | None = None,
         fail_on_error: bool | None = True,
+        streaming_buffer_until_moderated: bool | None = None,
+        streaming_buffer_release_on_scan: bool | None = None,
         streaming_end_of_stream_only: bool | None = None,
         streaming_sampling_rate: int | None = None,
         async_handler: AsyncHTTPHandler | None = None,
@@ -287,6 +289,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
             CrowdStrikeAIDRGuardrailConfigModelOptionalParams(
                 streaming_end_of_stream_only=streaming_end_of_stream_only,
                 streaming_sampling_rate=streaming_sampling_rate,
+                streaming_buffer_until_moderated=streaming_buffer_until_moderated,
+                streaming_buffer_release_on_scan=streaming_buffer_release_on_scan,
             )
         )
 
@@ -310,6 +314,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
         )
 
     def _set_streaming_params(self, streaming_params: CrowdStrikeAIDRGuardrailConfigModelOptionalParams) -> None:
+        self.streaming_buffer_until_moderated: bool = streaming_params.streaming_buffer_until_moderated or False
+        self.streaming_buffer_release_on_scan: bool = streaming_params.streaming_buffer_release_on_scan or False
         self.streaming_end_of_stream_only: bool = streaming_params.streaming_end_of_stream_only or False
         self.streaming_sampling_rate: int = streaming_params.streaming_sampling_rate or 5
 

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -956,6 +956,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         buffer_until_moderated: bool = _streaming_flag(
             "streaming_buffer_until_moderated", buffer_until_moderated_default
         )
+        release_on_scan: Final[bool] = _streaming_flag("streaming_buffer_release_on_scan", False)
 
         if (
             buffer_until_moderated
@@ -970,9 +971,7 @@ class UnifiedLLMGuardrails(CustomLogger):
             )
             buffer_until_moderated = False
 
-        # Buffering can only moderate the assembled response, so it always
-        # defers to end-of-stream.
-        if buffer_until_moderated:
+        if buffer_until_moderated and not release_on_scan:
             end_of_stream_only = True
 
         if guardrail_to_apply is None:
@@ -1026,12 +1025,14 @@ class UnifiedLLMGuardrails(CustomLogger):
         chunk_counter = 0
         responses_so_far: Final[list[object]] = []
         responses_yielded: Final[list[object]] = []
+        withheld_items: Final[list[object]] = []  # mutable-ok: streaming window must be released incrementally
         pending_end_of_stream_items: Final[list[object]] = []
         # Whether any real response chunk has been forwarded to the client.
         # Drives how a block terminates the stream: continue the in-progress
         # message (True) vs emit a standalone block message (False, buffered).
         chunks_yielded = False
         last_scan_key: StreamingScanKey | None = None  # rebind-ok: replaced after every scan round
+        tool_calls_in_flight = False  # rebind-ok: tracks the latest scan key's unscanned tool calls
 
         async for item in response:
             chunk_counter += 1
@@ -1069,21 +1070,37 @@ class UnifiedLLMGuardrails(CustomLogger):
                         chunks_yielded = True
                         responses_yielded.append(item)
                         yield item
+                else:
+                    withheld_items.append(item)
                 continue
 
             # Process chunk based on sampling rate
+            if buffer_until_moderated:
+                withheld_items.append(item)
             if chunk_counter % sampling_rate == 0:
                 endpoint_translation = mappings[CallTypes(call_type)]()
                 scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
+                if scan_key is not None:
+                    tool_calls_in_flight = scan_key.tool_calls_in_flight
+                hold_window = buffer_until_moderated and (scan_key is None or tool_calls_in_flight)
                 if _is_redundant_scan(scan_key, last_scan_key):
                     verbose_proxy_logger.debug(
                         "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round",
                         chunk_counter,
                         guardrail_to_apply.guardrail_name,
                     )
-                    chunks_yielded = True
-                    responses_yielded.append(item)
-                    yield item
+                    if buffer_until_moderated:
+                        if hold_window:
+                            continue
+                        for withheld_item in withheld_items:
+                            chunks_yielded = True
+                            responses_yielded.append(withheld_item)
+                            yield withheld_item
+                        withheld_items.clear()
+                    else:
+                        chunks_yielded = True
+                        responses_yielded.append(item)
+                        yield item
                     continue
 
                 verbose_proxy_logger.debug(
@@ -1093,13 +1110,9 @@ class UnifiedLLMGuardrails(CustomLogger):
                     guardrail_to_apply.guardrail_name,
                 )
 
-                # Deep-copy the current chunk before guardrail processing.
-                # process_output_streaming_response modifies responses_so_far
-                # in-place: it puts the combined guardrailed text in the first
-                # chunk and clears all subsequent chunks to "". Without this
-                # copy, yielding processed_items[-1] would yield an empty
-                # string, permanently losing this chunk's content.
-                original_item = copy.deepcopy(item)
+                original_items = (
+                    tuple(copy.deepcopy(withheld_items)) if buffer_until_moderated else (copy.deepcopy(item),)
+                )
 
                 try:
                     await endpoint_translation.process_output_streaming_response(
@@ -1144,13 +1157,24 @@ class UnifiedLLMGuardrails(CustomLogger):
                     return
                 if scan_key is not None:
                     last_scan_key = scan_key
-                chunks_yielded = True
-                responses_yielded.append(original_item)
-                yield original_item
+                if hold_window:
+                    verbose_proxy_logger.debug(
+                        "Holding %s buffered chunks for guardrail %s: this round could not scan the whole window",
+                        len(withheld_items),
+                        guardrail_to_apply.guardrail_name,
+                    )
+                    withheld_items[:] = original_items
+                    continue
+                for original_item in original_items:
+                    chunks_yielded = True
+                    responses_yielded.append(original_item)
+                    yield original_item
+                withheld_items.clear()
             else:
-                chunks_yielded = True
-                responses_yielded.append(item)
-                yield item
+                if not buffer_until_moderated:
+                    chunks_yielded = True
+                    responses_yielded.append(item)
+                    yield item
 
         # Stream has ended - do final processing with all collected chunks
         if call_type is not None and CallTypes(call_type) in mappings:
@@ -1162,14 +1186,13 @@ class UnifiedLLMGuardrails(CustomLogger):
 
             endpoint_translation = mappings[CallTypes(call_type)]()
 
-            # When buffering, snapshot the original chunks before moderation.
-            # A shallow copy suffices: end-of-stream
-            # process_output_streaming_response builds a separate assembled
-            # response (it does not mutate the individual chunks in place), and
-            # the chunks themselves are replayed verbatim -- so we only need to
-            # preserve the list, not clone every chunk (deepcopy would double
-            # peak memory for large responses).
-            buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None
+            buffered_items: Final = (
+                tuple(copy.deepcopy(withheld_items))
+                if buffer_until_moderated and release_on_scan and not end_of_stream_only
+                else tuple(withheld_items)
+                if buffer_until_moderated
+                else None
+            )
             end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far)
             if _is_redundant_scan(end_scan_key, last_scan_key):
                 verbose_proxy_logger.debug(

litellm/proxy/guardrails/guardrail_initializers.py

@@ -44,6 +44,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
         streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated,
         streaming_sampling_rate=streaming_params.streaming_sampling_rate,
         streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
+        streaming_buffer_release_on_scan=streaming_params.streaming_buffer_release_on_scan,
     )
     litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback)
     return _bedrock_callback

litellm/proxy/hooks/__init__.py

@@ -4,7 +4,6 @@ from typing import Final, Literal
 from . import *
 from .cache_control_check import _PROXY_CacheControlCheck
 from .litellm_skills import SkillsInjectionHook
-from .max_budget_limiter import _PROXY_MaxBudgetLimiter
 from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler
 from .max_iterations_limiter import _PROXY_MaxIterationsHandler
 from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler
@@ -18,7 +17,6 @@ from .sensitive_data_routing import _PROXY_SensitiveDataRoutingHandler
 # transitively through `enterprise.enterprise_hooks` can resolve `PROXY_HOOKS`
 # and `get_proxy_hook` from this partially-initialized module without circling.
 PROXY_HOOKS: Final = {
-    "max_budget_limiter": _PROXY_MaxBudgetLimiter,
     "parallel_request_limiter": _PROXY_MaxParallelRequestsHandler_v3,
     "cache_control_check": _PROXY_CacheControlCheck,
     "responses_id_security": ResponsesIDSecurity,
@@ -35,7 +33,7 @@ if os.getenv("LEGACY_MULTI_INSTANCE_RATE_LIMITING", "false").lower() == "true":
 
 
 def get_proxy_hook(
-    hook_name: Literal["max_budget_limiter", "managed_files", "parallel_request_limiter", "cache_control_check"] | str,
+    hook_name: Literal["managed_files", "parallel_request_limiter", "cache_control_check"] | str,
 ):
     """
     Factory method to get a proxy hook instance by name

litellm/proxy/hooks/max_budget_limiter.py

@@ -1,84 +0,0 @@
-from typing import Final
-
-from fastapi import HTTPException
-
-from litellm import verbose_logger
-from litellm._logging import verbose_proxy_logger
-from litellm.caching.caching import DualCache
-from litellm.exceptions import RateLimitType
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.proxy._types import UserAPIKeyAuth
-from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
-from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
-
-
-class _PROXY_MaxBudgetLimiter(CustomLogger):
-    # Class variables or attributes
-    def __init__(self):
-        pass
-
-    async def async_pre_call_hook(
-        self,
-        user_api_key_dict: UserAPIKeyAuth,
-        cache: DualCache,
-        data: dict,
-        call_type: str,
-    ):
-        try:
-            verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook")
-            max_budget: Final = user_api_key_dict.user_max_budget
-            user_id: Final = user_api_key_dict.user_id
-
-            if max_budget is None or user_id is None:
-                return
-
-            from litellm.proxy.proxy_server import general_settings
-
-            if (
-                user_api_key_dict.team_id is not None
-                and general_settings.get("apply_user_budget_to_team_keys") is not True
-            ):
-                return
-
-            # The reservation path admits at the strict-`<` boundary and
-            # atomically pre-fills the same counter we'd read here. Re-checking
-            # with `>=` would reject a request the reservation already admitted
-            # when the reservation fills the counter to exactly max_budget.
-            # Imported lazily to avoid a circular import via proxy.utils.
-            from litellm.proxy.spend_tracking.budget_reservation import (
-                get_reserved_counter_keys,
-            )
-
-            user_counter_key: Final = f"spend:user:{user_id}"
-            if user_counter_key in get_reserved_counter_keys(user_api_key_dict.budget_reservation):
-                return
-
-            from litellm.proxy.proxy_server import get_current_spend
-
-            curr_spend: Final = await get_current_spend(
-                counter_key=user_counter_key,
-                fallback_spend=user_api_key_dict.user_spend or 0.0,
-            )
-
-            verbose_proxy_logger.debug(
-                "MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f",
-                user_id,
-                curr_spend,
-                max_budget,
-            )
-
-            # CHECK IF REQUEST ALLOWED
-            if curr_spend >= max_budget:
-                resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None)
-                raise ProxyRateLimitError(
-                    detail="Max budget limit reached.",
-                    rate_limit_type=RateLimitType.BUDGET,
-                    model=resolved_model,
-                    llm_provider=llm_provider,
-                )
-        except HTTPException as e:
-            raise e
-        except Exception as e:
-            verbose_logger.exception(
-                "litellm.proxy.hooks.max_budget_limiter.py::async_pre_call_hook(): Exception occured - %s", e
-            )

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -19,12 +19,14 @@ from litellm.types.utils import BudgetConfig, StandardLoggingPayload
 VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend"
 END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend"
 USER_SPEND_CACHE_KEY_PREFIX: Final = "user_model_spend"
+TEAM_SPEND_CACHE_KEY_PREFIX: Final = "team_model_spend"
 
 _SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType(
     {
         Litellm_EntityType.KEY: VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
         Litellm_EntityType.USER: USER_SPEND_CACHE_KEY_PREFIX,
         Litellm_EntityType.END_USER: END_USER_SPEND_CACHE_KEY_PREFIX,
+        Litellm_EntityType.TEAM: TEAM_SPEND_CACHE_KEY_PREFIX,
     }
 )
 
@@ -37,6 +39,7 @@ _BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType(
         Litellm_EntityType.KEY: "virtual_key_budget_start_time",
         Litellm_EntityType.USER: "user_model_budget_start_time",
         Litellm_EntityType.END_USER: "end_user_budget_start_time",
+        Litellm_EntityType.TEAM: "team_model_budget_start_time",
     }
 )
 
@@ -139,6 +142,18 @@ def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) ->
     return None
 
 
+def team_model_budget_applies(model: str, key_model_max_budget: Mapping[str, object] | None) -> bool:
+    """A key entry that spend-gates `model` overrides the team cap: it is then gated on and billed to the key alone."""
+    if not key_model_max_budget:
+        return True
+    resolved: Final = resolve_model_budget(model=model, model_max_budget=key_model_max_budget)
+    return resolved is None or not _spend_gated(resolved.budget_config)
+
+
+def _spend_gated(budget_config: BudgetConfig) -> bool:
+    return budget_config.max_budget is not None and budget_config.max_budget >= 0
+
+
 def _budget_model_candidates(model: str) -> tuple[str, ...]:
     """Names a budget may be configured under for a request on `model`, most specific first.
 
@@ -346,6 +361,30 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
             exceeded_message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}",
         )
 
+    async def is_team_within_model_budget(
+        self,
+        team_id: str,
+        team_model_max_budget: Mapping[str, object],
+        key_model_max_budget: Mapping[str, object] | None,
+        model: str,
+    ) -> bool:
+        """
+        Check if the team is within the model budget, unless the key's own
+        `model_max_budget` overrides it for `model`
+
+        Raises:
+            BudgetExceededError: If the team has exceeded the model budget
+        """
+        if not team_model_budget_applies(model=model, key_model_max_budget=key_model_max_budget):
+            return True
+        return await self._is_entity_within_model_budget(
+            entity_type=Litellm_EntityType.TEAM,
+            entity_id=team_id,
+            model_max_budget=team_model_max_budget,
+            model=model,
+            exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}",
+        )
+
     async def _is_entity_within_model_budget(
         self,
         entity_type: Litellm_EntityType,
@@ -456,11 +495,26 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
             return
 
         response_cost: Final[float] = standard_logging_payload.get("response_cost", 0)
+        key_model_max_budget: Final = _metadata.get("user_api_key_model_max_budget")
         entity_budgets: Final = (
             (
                 Litellm_EntityType.KEY,
                 payload_metadata.get("user_api_key_hash"),
-                _metadata.get("user_api_key_model_max_budget"),
+                key_model_max_budget,
+            ),
+            (
+                Litellm_EntityType.TEAM,
+                payload_metadata.get("user_api_key_team_id"),
+                (
+                    _metadata.get("user_api_key_team_model_max_budget")
+                    if team_model_budget_applies(
+                        model=model,
+                        key_model_max_budget=(
+                            key_model_max_budget if isinstance(key_model_max_budget, Mapping) else None
+                        ),
+                    )
+                    else None
+                ),
             ),
             (
                 Litellm_EntityType.USER,
@@ -478,7 +532,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
         if not resolved_budgets:
             verbose_proxy_logger.debug(
                 "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: "
-                "no key, user or end-user model_max_budget covers model=%s",
+                "no key, team, user or end-user model_max_budget covers model=%s",
                 model,
             )
             return

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -531,6 +531,7 @@ class RequestRateLimiterStash:
     owner_litellm_call_id: str | None = None
     rate_limit_response: RateLimitResponse | None = None
     parallel_slot: ParallelSlotAcquisition | None = None
+    parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False)
     reserved_tokens: int = 0
     reserved_model: str | None = None
     reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
@@ -1620,6 +1621,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             statuses.append(self._gauge_status(gauge, in_flight + 1, "OK"))
         return RateLimitResponse(overall_code="OK", statuses=statuses)
 
+    async def _release_stashed_parallel_slot(
+        self,
+        stash: RequestRateLimiterStash | None,
+        parent_otel_span: Span | None,
+    ) -> None:
+        if stash is None:
+            return
+        async with stash.parallel_slot_release_lock:
+            acquisition: Final = stash.parallel_slot
+            if acquisition is None:
+                return
+            await self._release_parallel_request_slots(acquisition, parent_otel_span)
+            stash.parallel_slot = None  # rebind-ok: marks this request's slot as released
+
     async def _release_parallel_request_slots(
         self,
         acquisition: ParallelSlotAcquisition,
@@ -3379,13 +3394,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                     parent_otel_span=user_api_key_dict.parent_otel_span,
                 )
                 stash.reservation_released = True
-            acquisition: Final = stash.parallel_slot
-            if acquisition is not None:
-                await self._release_parallel_request_slots(
-                    acquisition=acquisition,
-                    parent_otel_span=user_api_key_dict.parent_otel_span,
-                )
-                stash.parallel_slot = None
+            await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
             self._handle_rate_limit_error(
                 response=io_response,
                 descriptors=descriptors,
@@ -3700,13 +3709,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                 )
 
                 if tpm_response["overall_code"] == "OVER_LIMIT":
-                    acquisition: Final = stash.parallel_slot
-                    if acquisition is not None:
-                        await self._release_parallel_request_slots(
-                            acquisition=acquisition,
-                            parent_otel_span=user_api_key_dict.parent_otel_span,
-                        )
-                        stash.parallel_slot = None
+                    await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
                     self._handle_rate_limit_error(
                         response=tpm_response,
                         descriptors=descriptors,
@@ -4524,13 +4527,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING")
 
             stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs))
-            acquisition: Final = stash.parallel_slot if stash is not None else None
-            if stash is not None and acquisition is not None:
-                await self._release_parallel_request_slots(
-                    acquisition=acquisition,
-                    parent_otel_span=litellm_parent_otel_span,
-                )
-                stash.parallel_slot = None
+            await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span)
 
             pipeline_operations: Final = self._build_success_event_pipeline_operations(
                 kwargs=kwargs,
@@ -4650,13 +4647,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = []
 
             stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs))
-            acquisition: Final = stash.parallel_slot if stash is not None else None
-            if stash is not None and acquisition is not None:
-                await self._release_parallel_request_slots(
-                    acquisition=acquisition,
-                    parent_otel_span=litellm_parent_otel_span,
-                )
-                stash.parallel_slot = None
+            await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span)
 
             # Skip the reservation refund if async_post_call_failure_hook
             # already released it (proxy-level rejection that also bubbles up
@@ -4764,23 +4755,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         object's current max_parallel_requests configuration, which can
         change mid-request) decides whether there is anything to release.
         """
-        stash: Final = get_request_stash()
-        if stash is None or stash.parallel_slot is None:
-            return
-
-        await self._release_parallel_request_slots(
-            acquisition=stash.parallel_slot,
-            parent_otel_span=None,
-        )
-        stash.parallel_slot = None
+        await self._release_stashed_parallel_slot(get_request_stash(), None)
 
     async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response):
         """
-        Post-call hook to update rate limit headers in the response.
+        Release completed-request slots and update rate limit headers in the response.
         """
         try:
-            stash: Final = get_request_stash()
-            litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None
+            slot_stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(data))
+            await self._release_stashed_parallel_slot(slot_stash, user_api_key_dict.parent_otel_span)
+        except Exception as e:
+            verbose_proxy_logger.exception("Error releasing parallel request slot in post-call hook: %s", e)
+
+        try:
+            header_stash: Final = get_request_stash()
+            litellm_proxy_rate_limit_response: Final = (
+                header_stash.rate_limit_response if header_stash is not None else None
+            )
 
             if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response):
                 additional_headers: Final = ensure_response_additional_headers(response)
@@ -4848,12 +4839,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             stash: Final = get_request_stash()
             if stash is None:
                 return
-            if stash.parallel_slot is not None:
-                await self._release_parallel_request_slots(
-                    acquisition=stash.parallel_slot,
-                    parent_otel_span=user_api_key_dict.parent_otel_span,
-                )
-                stash.parallel_slot = None
+            await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
 
             if stash.batch_enqueued_reservation is not None:
                 await self.batch_enqueued_token_store.refund(

litellm/proxy/litellm_pre_call_utils.py

@@ -25,6 +25,7 @@ from litellm.constants import (
     LITELLM_PROXY_MASTER_KEY_ALIAS,
     OTEL_SERVICE_NAME_METADATA_KEYS,
     PRE_CALL_EXECUTED_GUARDRAILS_KEY,
+    ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
     ROUTING_REQUEST_TAGS_METADATA_KEY,
     SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
     SESSION_ID_GENERATED_METADATA_KEY,
@@ -369,7 +370,13 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg
 # and read by spend logs as fact; a client value has no legitimate meaning and no
 # key or team setting keeps it, so the strip is never gated.
 _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
-    {"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY}
+    {
+        "attempted_fallbacks",
+        "original_model_group",
+        "request_retry_count",
+        CLIENT_OUTPUT_CEILING_METADATA_KEY,
+        ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
+    }
 )
 _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override"
 
@@ -2327,6 +2334,7 @@ async def add_litellm_data_to_request(
     # Team spend, budget - used by prometheus.py
     data[_metadata_variable_name]["user_api_key_team_max_budget"] = user_api_key_dict.team_max_budget
     data[_metadata_variable_name]["user_api_key_team_spend"] = user_api_key_dict.team_spend
+    data[_metadata_variable_name]["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget
     data[_metadata_variable_name]["user_api_key_request_route"] = user_api_key_dict.request_route
 
     # API Key spend, budget - used by prometheus.py

litellm/proxy/management_endpoints/common_utils.py

@@ -55,6 +55,7 @@ def validate_budget_duration(budget_duration: str | None, status_code: int = 400
 from litellm._logging import verbose_proxy_logger
 from litellm.caching import DualCache
 from litellm.proxy._types import (
+    CommonProxyErrors,
     KeyRequestBase,
     LiteLLM_ManagementEndpoint_MetadataFields,
     LiteLLM_ManagementEndpoint_MetadataFields_Premium,
@@ -73,12 +74,62 @@ from litellm.proxy._types import (  # noqa: F401  re-exported
 from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
 from litellm.proxy.utils import _premium_user_check
 from litellm.repositories.team_repository import TeamRepository
+from litellm.types.utils import BudgetConfig
 
 if TYPE_CHECKING:
     from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest
     from litellm.proxy.utils import PrismaClient, ProxyLogging
 
 
+def validate_team_model_max_budget(
+    model_max_budget: Mapping[str, BudgetConfig] | None,
+    premium_user: bool,
+) -> None:
+    """Reject a team `model_max_budget` the limiter could not enforce (no duration, bad cap, tpm/rpm limits)."""
+    if not model_max_budget:
+        return
+    if premium_user is not True:
+        raise HTTPException(
+            status_code=403,
+            detail={
+                "error": f"Setting model_max_budget on a team is an enterprise feature. {CommonProxyErrors.not_premium_user.value}"
+            },
+        )
+    for model_name, budget_config in model_max_budget.items():
+        if not model_name.strip():
+            raise HTTPException(
+                status_code=400,
+                detail={"error": "model_max_budget keys must be non-empty model names"},
+            )
+        max_budget = budget_config.max_budget
+        if max_budget is None or not math.isfinite(max_budget) or max_budget < 0:
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "error": (
+                        f"model_max_budget[{model_name!r}].max_budget must be a non-negative finite number. "
+                        f"Received: {max_budget}"
+                    )
+                },
+            )
+        if budget_config.budget_duration is None:
+            raise HTTPException(
+                status_code=400,
+                detail={"error": f"model_max_budget[{model_name!r}] requires a budget_duration, e.g. '1d' or '30d'"},
+            )
+        validate_budget_duration(budget_config.budget_duration)
+        if budget_config.tpm_limit is not None or budget_config.rpm_limit is not None:
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "error": (
+                        f"model_max_budget[{model_name!r}] tpm_limit/rpm_limit are not enforced on a team; "
+                        "set per-model rate limits on the key instead"
+                    )
+                },
+            )
+
+
 def require_caller_user_id_for_non_admin(
     user_api_key_dict: UserAPIKeyAuth,
 ) -> str:

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -4166,7 +4166,10 @@ async def info_key_fn(
 
     Returns:
     - key: str - The key that was looked up, echoed back as it was passed in
-    - info: dict - The key's row, minus the hashed token
+    - info: dict - The key's row, minus the hashed token. Deleted keys are served from the
+      LiteLLM_DeletedVerificationToken archive and carry deleted_at / deleted_by
+        - status: "active" | "expired" | "revoked" | "deleted" - Derived from blocked, expires and
+          whether the row came from the archive
         - key_alias: str | None - User-friendly key alias
         - spend: float - Amount spent by the key. When budget_duration is set this covers only the
           current budget window, not the key's lifetime
@@ -4220,10 +4223,15 @@ async def info_key_fn(
         hashed_key: str | None = key
         if key is not None:
             hashed_key = _hash_token_if_needed(token=key)
-        key_info = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
+        live_key_info: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
             where={"token": hashed_key},
             include={"litellm_budget_table": True},
         )
+        key_info: Final = (
+            live_key_info
+            if live_key_info is not None
+            else await _find_deleted_key_info(prisma_client=prisma_client, hashed_key=hashed_key)
+        )
         if key_info is None:
             raise ProxyException(
                 message="Key not found in database",
@@ -4231,7 +4239,6 @@ async def info_key_fn(
                 param="key",
                 code=status.HTTP_404_NOT_FOUND,
             )
-
         if (
             await _can_user_query_key_info(
                 user_api_key_dict=user_api_key_dict,
@@ -4245,38 +4252,46 @@ async def info_key_fn(
                 detail=f"You are not allowed to access this key's info. Your role={user_api_key_dict.user_role}",
             )
         ## REMOVE HASHED TOKEN INFO BEFORE RETURNING ##
-        try:
-            key_info = key_info.model_dump()
-        except Exception:
-            # if using pydantic v1
-            key_info = key_info.dict()  # pyright: ignore[reportDeprecated]  # deliberate pydantic v1 fallback
-        key_token_hash: Final[str | None] = key_info.pop("token")
+        key_info_dict: Final = key_info.model_dump()
+        key_token_hash: Final[str | None] = key_info_dict.pop("token")
+        key_info_dict["status"] = (
+            "deleted" if live_key_info is None else _derive_key_status(key_info_dict, now=datetime.now(timezone.utc))
+        )
 
-        model_max_budget = key_info.get("model_max_budget") or {}
-        budget_table: Final = key_info.get("litellm_budget_table") or {}
+        model_max_budget = key_info_dict.get("model_max_budget") or {}
+        budget_table: Final = key_info_dict.get("litellm_budget_table") or {}
         if not model_max_budget and isinstance(budget_table, dict):
             model_max_budget = budget_table.get("model_max_budget") or {}
         if model_max_budget and key_token_hash:
-            key_info["model_max_budget_usage"] = await _build_model_max_budget_usage(
+            key_info_dict["model_max_budget_usage"] = await _build_model_max_budget_usage(
                 api_key_hash=key_token_hash,
                 model_max_budget=model_max_budget,
                 user_api_key_cache=model_max_budget_limiter.dual_cache,
             )
         budget_limits_usage: Final = await _build_budget_limits_usage(
-            budget_limits=key_info.get("budget_limits"),
+            budget_limits=key_info_dict.get("budget_limits"),
             api_key_hash=key_token_hash,
         )
         if budget_limits_usage is not None:
-            key_info["budget_limits_usage"] = budget_limits_usage
+            key_info_dict["budget_limits_usage"] = budget_limits_usage
 
-        # Attach object_permission if object_permission_id is set
-        key_info = await attach_object_permission_to_dict(key_info, prisma_client)
-
-        return {"key": key, "info": key_info}
+        return {"key": key, "info": await attach_object_permission_to_dict(key_info_dict, prisma_client)}
     except Exception as e:
         raise handle_exception_on_proxy(e)
 
 
+async def _find_deleted_key_info(
+    prisma_client: PrismaClient, hashed_key: str | None
+) -> LiteLLM_DeletedVerificationToken | None:
+    archived_row: Final = await _deleted_verification_token_table(prisma_client).find_first(
+        where={"token": hashed_key},
+        order={"deleted_at": "desc"},
+    )
+    if archived_row is None:
+        return None
+    return LiteLLM_DeletedVerificationToken.model_validate(archived_row.model_dump())
+
+
 def _check_model_access_group(models: list[str] | None, llm_router: Router | None, premium_user: bool) -> Literal[True]:
     """
     if is_model_access_group is True + is_wildcard_route is True, check if user is a premium user
@@ -6216,6 +6231,24 @@ async def get_member_team_ids(
 
 VALID_EXPIRES_FILTER_VALUES: Final = frozenset({"active", "expired"})
 
+KeyStatus = Literal["active", "expired", "revoked", "deleted"]
+VALID_STATUS_FILTER_VALUES: Final[frozenset[KeyStatus]] = frozenset({"active", "expired", "revoked", "deleted"})
+
+
+class _KeyStatusSource(BaseModel):
+    blocked: bool | None = None
+    expires: datetime | None = None
+
+
+def _derive_key_status(row: Mapping[str, object], now: datetime) -> KeyStatus:
+    source: Final = _KeyStatusSource.model_validate(row)
+    if source.blocked is True:
+        return "revoked"
+    if source.expires is None:
+        return "active"
+    expires_utc: Final = source.expires if source.expires.tzinfo else source.expires.replace(tzinfo=timezone.utc)
+    return "expired" if expires_utc < now else "active"
+
 
 @router.get(
     "/key/list",
@@ -6252,7 +6285,10 @@ async def list_keys(
     ),
     sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"),
     expand: list[str] | None = Query(None, description="Expand related objects (e.g. 'user')"),
-    status: str | None = Query(None, description="Filter by status (e.g. 'deleted')"),
+    status: str | None = Query(
+        None,
+        description="Filter by status: 'active' (not blocked, not expired), 'expired' (not blocked, past expiry), 'revoked' (blocked) or 'deleted' (archived keys). Omit to return live keys regardless of status.",
+    ),
     project_id: str | None = Query(None, description="Filter keys by project ID"),
     access_group_id: str | None = Query(None, description="Filter keys by access group ID"),
     agent_id: str | None = Query(None, description="Filter keys by agent ID"),
@@ -6270,7 +6306,9 @@ async def list_keys(
 
     Parameters:
         expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information)
-        status: Optional[str] - Filter by status. Currently supports "deleted" to query deleted keys.
+        status: Optional[str] - Filter by status: "active", "expired", "revoked" (blocked) or "deleted".
+        "deleted" reads the LiteLLM_DeletedVerificationToken archive; the other values partition the
+        live key table, so every live key matches exactly one of them.
 
     Returns:
         {
@@ -6292,11 +6330,10 @@ async def list_keys(
             verbose_proxy_logger.error("Database not connected")
             raise Exception("Database not connected")
 
-        # Validate status parameter
-        if status is not None and status != "deleted":
+        if status is not None and status not in VALID_STATUS_FILTER_VALUES:
             raise HTTPException(
                 status_code=400,
-                detail={"error": "Invalid status value. Currently only 'deleted' is supported."},
+                detail={"error": "Invalid status value. Supported: 'active', 'expired', 'revoked', 'deleted'."},
             )
 
         if isinstance(expires, str) and expires not in VALID_EXPIRES_FILTER_VALUES:
@@ -6608,6 +6645,18 @@ def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str,
     return {"OR": [{"expires": None}, {"expires": {"gte": now}}]}
 
 
+def _not_blocked_where_clause() -> dict[str, object]:
+    return {"OR": [{"blocked": None}, {"blocked": False}]}
+
+
+def _build_status_where_clause(status_filter: str | None, now: datetime) -> dict[str, object] | None:
+    if status_filter == "revoked":
+        return {"blocked": True}
+    if status_filter in ("expired", "active"):
+        return {"AND": [_not_blocked_where_clause(), _build_expires_where_clause(status_filter, now)]}
+    return None
+
+
 def _build_key_search_where(search: str) -> KeySearchWhere:
     search_where: Final[KeySearchWhere] = {
         "OR": (
@@ -6635,6 +6684,7 @@ def _build_key_filter_conditions(
     use_key_alias_substring_matching: bool = False,
     expires_filter: str | None = None,
     search: str | None = None,
+    status_filter: str | None = None,
 ) -> Mapping[str, object]:
     """Build filter conditions for key listing.
 
@@ -6724,6 +6774,8 @@ def _build_key_filter_conditions(
 
     # Apply team_id, project_id and access_group_id as global AND filters so they
     # narrow results across all visibility conditions (own keys, team keys, etc.)
+    now: Final = datetime.now(timezone.utc)
+    status_where: Final = _build_status_where_clause(status_filter, now)
     global_filters: Final[tuple[Mapping[str, object], ...]] = (
         *(
             (
@@ -6741,10 +6793,11 @@ def _build_key_filter_conditions(
         *(({"access_group_ids": {"hasSome": [access_group_id]}},) if access_group_id else ()),
         *(({"agent_id": agent_id},) if agent_id and isinstance(agent_id, str) else ()),
         *(
-            (_build_expires_where_clause(expires_filter, datetime.now(timezone.utc)),)
+            (_build_expires_where_clause(expires_filter, now),)
             if expires_filter is not None and expires_filter in VALID_EXPIRES_FILTER_VALUES
             else ()
         ),
+        *((status_where,) if status_where is not None else ()),
     )
     combined_where: Final[Mapping[str, object]] = {"AND": [where, *global_filters]} if global_filters else where
     verbose_proxy_logger.debug("Filter conditions: %s", combined_where)
@@ -6817,6 +6870,7 @@ async def _list_key_helper(
         use_key_alias_substring_matching=use_key_alias_substring_matching,
         expires_filter=expires_filter,
         search=search,
+        status_filter=status,
     )
 
     # Calculate skip for pagination

litellm/proxy/management_endpoints/team_endpoints.py

@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protoc
 
 import fastapi
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
-from pydantic import BaseModel, JsonValue
+from pydantic import BaseModel, JsonValue, ValidationError
 from typing_extensions import ReadOnly, TypedDict
 
 import litellm
@@ -38,6 +38,7 @@ from litellm.proxy._types import (
     DeleteTeamRequest,
     LiteLLM_AuditLogs,
     LiteLLM_DeletedTeamTable,
+    Litellm_EntityType,
     LiteLLM_ManagementEndpoint_MetadataFields,
     LiteLLM_ManagementEndpoint_MetadataFields_Premium,
     LiteLLM_ModelTable,
@@ -95,6 +96,10 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
 from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch
 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+from litellm.proxy.hooks.model_max_budget_limiter import (
+    build_model_max_budget_usage,
+    resolve_model_budget,
+)
 from litellm.proxy.management_endpoints.common_daily_activity import (
     get_daily_activity_aggregated,
 )
@@ -108,6 +113,7 @@ from litellm.proxy.management_endpoints.common_utils import (
     _upsert_budget_and_membership,
     _user_has_admin_view,
     validate_budget_duration,
+    validate_team_model_max_budget,
 )
 from litellm.proxy.management_endpoints.organization_endpoints import (
     add_member_to_organization,
@@ -177,6 +183,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
     TeamUserSpendRow,
     UpdateTeamMemberPermissionsRequest,
 )
+from litellm.types.utils import BudgetConfig
 
 if TYPE_CHECKING:
     from prisma import Prisma
@@ -1170,6 +1177,62 @@ def _check_team_budget_update_authority(
         )
 
 
+def _existing_model_cap(raw_budget_config: object) -> BudgetConfig | None:
+    try:
+        return BudgetConfig.model_validate(raw_budget_config)
+    except ValidationError:
+        return None
+
+
+def _check_team_model_budget_update_authority(
+    data: UpdateTeamRequest,
+    user_api_key_dict: UserAPIKeyAuth,
+    existing_model_max_budget: Mapping[str, object] | None,
+) -> None:
+    """Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap."""
+    if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
+        return
+    if "model_max_budget" not in data.model_fields_set or not existing_model_max_budget:
+        return
+    requested: Final[Mapping[str, BudgetConfig]] = data.model_max_budget or {}
+    for model_name, raw_existing in existing_model_max_budget.items():
+        existing = _existing_model_cap(raw_existing)
+        if existing is None or existing.max_budget is None or model_name in requested:
+            continue
+        raise HTTPException(
+            status_code=403,
+            detail={
+                "error": (
+                    f"Only a proxy admin can remove a team's model_max_budget for {model_name!r}. "
+                    f"Current max_budget={existing.max_budget}."
+                )
+            },
+        )
+    for model_name, proposed in requested.items():
+        governing = resolve_model_budget(model=model_name, model_max_budget=existing_model_max_budget)
+        if governing is None:
+            continue
+        cap = governing.budget_config
+        if cap.max_budget is None:
+            continue
+        if (
+            proposed.max_budget is None
+            or proposed.max_budget > cap.max_budget
+            or proposed.budget_duration != cap.budget_duration
+        ):
+            raise HTTPException(
+                status_code=403,
+                detail={
+                    "error": (
+                        f"Only a proxy admin can raise a team's model_max_budget for {model_name!r} or change its "
+                        f"budget_duration. Current max_budget={cap.max_budget} per {cap.budget_duration} "
+                        f"(entry {governing.budget_model!r}), requested={proposed.max_budget} per "
+                        f"{proposed.budget_duration}."
+                    )
+                },
+            )
+
+
 def _should_auto_add_team_creator(
     user_api_key_dict: UserAPIKeyAuth,
     general_settings: Mapping[str, object],
@@ -1230,6 +1293,7 @@ async def new_team(
     - prompts: Optional[List[str]] - List of prompts that the team is allowed to use.
     - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
     - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
+    - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
     - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
     - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
     - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
@@ -1291,6 +1355,7 @@ async def new_team(
             general_settings,
             litellm_proxy_admin_name,
             llm_router,
+            premium_user,
             prisma_client,
             user_api_key_cache,
         )
@@ -1321,6 +1386,7 @@ async def new_team(
 
         validate_budget_duration(data.budget_duration)
         validate_budget_duration(data.team_member_budget_duration)
+        validate_team_model_max_budget(model_max_budget=data.model_max_budget, premium_user=premium_user)
 
         if data.soft_budget is not None:
             if data.max_budget is not None:
@@ -1980,6 +2046,7 @@ async def update_team(
     - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
     - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
     - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
+    - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}}
     - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
     - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
     - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
@@ -2031,6 +2098,7 @@ async def update_team(
         from litellm.proxy.proxy_server import (
             litellm_proxy_admin_name,
             llm_router,
+            premium_user,
             prisma_client,
             proxy_logging_obj,
             user_api_key_cache,
@@ -2069,6 +2137,7 @@ async def update_team(
 
         validate_budget_duration(data.budget_duration)
         validate_budget_duration(data.team_member_budget_duration)
+        validate_team_model_max_budget(model_max_budget=data.model_max_budget, premium_user=premium_user)
 
         existing_team_row = await _raw_team_db(TeamRepository(prisma_client)).find_unique(
             where={"team_id": data.team_id}
@@ -2204,8 +2273,15 @@ async def update_team(
                 user_api_key_dict=user_api_key_dict,
                 existing_team_max_budget=existing_team_row.max_budget,
             )
+        _check_team_model_budget_update_authority(
+            data=data,
+            user_api_key_dict=user_api_key_dict,
+            existing_model_max_budget=existing_team_row.model_max_budget,
+        )
 
         updated_kv = data.json(exclude_unset=True)
+        if "model_max_budget" in updated_kv and updated_kv["model_max_budget"] is None:
+            updated_kv["model_max_budget"] = {}
 
         # Drop server-owned metadata keys from caller input so they can only
         # be written by the same code path that creates the underlying rows.
@@ -4473,7 +4549,7 @@ async def team_info(
     ```
     """
     from litellm.proxy._types import TeamInfoResponseObjectTeamTable
-    from litellm.proxy.proxy_server import prisma_client
+    from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client
 
     try:
         if prisma_client is None:
@@ -4573,6 +4649,12 @@ async def team_info(
             update={  # mutable-ok: pydantic update payload
                 "members_with_roles": hydrated_members,
                 "organization_models": organization_models,
+                "model_max_budget_usage": await build_model_max_budget_usage(
+                    entity_type=Litellm_EntityType.TEAM,
+                    entity_id=team_id,
+                    model_max_budget=resolved_team_info.model_max_budget,
+                    cache=model_max_budget_limiter.dual_cache,
+                ),
             }
         )
 

litellm/proxy/ocr_endpoints/endpoints.py

@@ -15,12 +15,13 @@ from litellm.llms.base_llm.ocr.transformation import (
     OCRResponse,
     parse_ocr_request_format,
 )
-from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes
+from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
 from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
 
 router: Final = APIRouter()
+_MAX_FILE_BYTES: Final = 50 * 1024 * 1024
 
 
 def _build_document_from_upload(
@@ -28,7 +29,15 @@ def _build_document_from_upload(
     filename: str | None,
     content_type: str | None,
 ) -> dict[str, str]:
-    return convert_upload_to_url_document(file_content, filename, content_type)
+    supplied_mime: Final = content_type.split(";")[0].strip() if content_type else None
+    mime_type: Final = (
+        get_mime_type(filename)
+        if filename and (not supplied_mime or supplied_mime == "application/octet-stream")
+        else supplied_mime
+    )
+    return convert_file_document_to_url_document(
+        {"type": "file", "file": file_content, "mime_type": mime_type or "application/octet-stream"}
+    )
 
 
 def _with_request_format(data: Mapping[str, Any], request: Request) -> Mapping[str, Any]:
@@ -103,9 +112,11 @@ async def _parse_multipart_form(request: Request) -> dict[str, Any]:
 
     # Seek to start in case the file was already partially read by middleware
     await uploaded_file.seek(0)
-    file_content: Final = await uploaded_file.read(get_max_file_bytes() + 1)
+    file_content: Final = await uploaded_file.read(_MAX_FILE_BYTES + 1)
     if not file_content:
         raise ValueError("Uploaded file is empty")
+    if len(file_content) > _MAX_FILE_BYTES:
+        raise ValueError("OCR file exceeds the size limit")
 
     document: Final = _build_document_from_upload(
         file_content=file_content,

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -609,6 +609,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
         # merely shares the name.
         if not request_dispatched_to_pass_through_endpoint(request):
             _metadata["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget
+            _metadata["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget
             _metadata["user_api_key_user_model_max_budget"] = user_api_key_dict.user_model_max_budget
             _metadata["user_api_key_end_user_model_max_budget"] = user_api_key_dict.end_user_model_max_budget
         _metadata.update(
@@ -985,6 +986,7 @@ async def pass_through_request(
             headers=headers,
             forward_headers=forward_headers,
         )
+        upstream_headers: Final = _with_trace_context(headers, parent_span=user_api_key_dict.parent_otel_span)
 
         requested_query_params: dict | None = query_params or dict(request.query_params)
 
@@ -1018,7 +1020,7 @@ async def pass_through_request(
         verbose_proxy_logger.debug(
             "Pass through endpoint sending request to \nURL %s\nheaders: %s\nbody: %s\n",
             url,
-            headers,
+            upstream_headers,
             _parsed_body,
         )
 
@@ -1256,7 +1258,7 @@ async def pass_through_request(
             additional_args={
                 "complete_input_dict": _parsed_body,
                 "api_base": str(logging_url),
-                "headers": headers,
+                "headers": upstream_headers,
             },
         )
         stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
@@ -1273,7 +1275,7 @@ async def pass_through_request(
                     request=request,
                     async_client=async_client,
                     url=url,
-                    headers=headers,
+                    headers=upstream_headers,
                     requested_query_params=requested_query_params,
                     stream=True,
                 )
@@ -1285,7 +1287,7 @@ async def pass_through_request(
                         request.method,
                         url,
                         params=requested_query_params,
-                        headers=headers,
+                        headers=upstream_headers,
                         content=state_raw_body,
                     )
                     if state_raw_body is not None
@@ -1293,7 +1295,7 @@ async def pass_through_request(
                         request.method,
                         url,
                         params=requested_query_params,
-                        headers=headers,
+                        headers=upstream_headers,
                         json=_parsed_body,
                     )
                 )
@@ -1370,7 +1372,7 @@ async def pass_through_request(
             raw_body_request: Final = async_client.build_request(
                 request.method,
                 url,
-                headers=headers,
+                headers=upstream_headers,
                 params=requested_query_params,
                 content=state_raw_body,
             )
@@ -1380,7 +1382,7 @@ async def pass_through_request(
                 request=request,
                 async_client=async_client,
                 url=url,
-                headers=headers,
+                headers=upstream_headers,
                 requested_query_params=requested_query_params,
                 _parsed_body=_parsed_body,
                 forward_multipart=is_multipart,
@@ -2157,6 +2159,17 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None:
     return upstream_close
 
 
+_WEBSOCKET_FORWARDED_HEADERS: Final = frozenset(("authorization", "x-api-key", "x-goog-user-project"))
+
+
+def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict[str, str]:
+    try:
+        from litellm.integrations.otel.plumbing.context import inject_trace_context
+    except ImportError:
+        return dict(headers)  # mutable-ok: matches inject_trace_context's carrier return type
+    return inject_trace_context(headers, parent_span=parent_span)
+
+
 async def websocket_passthrough_request(
     websocket: WebSocket,
     target: str,
@@ -2199,20 +2212,15 @@ async def websocket_passthrough_request(
         await websocket.accept()
         verbose_proxy_logger.debug("WebSocket passthrough (%s): WebSocket connection accepted", endpoint)
 
-    # Prepare headers for the upstream connection
-    upstream_headers: Final = custom_headers.copy()
-
-    if forward_headers:
-        # Forward relevant headers from the incoming request
-        incoming_headers: Final = dict(websocket.headers)
-        for header_name, header_value in incoming_headers.items():
-            # Only forward certain headers to avoid conflicts
-            if header_name.lower() in [
-                "authorization",
-                "x-api-key",
-                "x-goog-user-project",
-            ]:
-                upstream_headers[header_name] = header_value
+    forwarded_headers: Final = {  # mutable-ok: one-shot upstream header dict, read as a Mapping
+        **custom_headers,
+        **{
+            header_name: header_value
+            for header_name, header_value in websocket.headers.items()
+            if forward_headers and header_name.lower() in _WEBSOCKET_FORWARDED_HEADERS
+        },
+    }
+    upstream_headers: Final = _with_trace_context(forwarded_headers, parent_span=user_api_key_dict.parent_otel_span)
 
     # Initialize logging object similar to HTTP passthrough
     team_callbacks: Final = _resolve_team_callback_wiring(

litellm/proxy/proxy_server.py

@@ -323,6 +323,7 @@ from litellm.proxy.auth.auth_utils import (
     log_once_if_budget_reservation_disabled,
     warn_once_if_custom_auth_skips_common_checks,
 )
+from litellm.proxy.auth.fallback_budget import router_fallback_budget_check
 from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
 from litellm.proxy.auth.handle_jwt import JWTHandler
 from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck
@@ -6161,6 +6162,7 @@ class ProxyConfig:
             ),
             ignore_invalid_deployments=True,  # don't raise an error if a deployment is invalid
             fallback_access_check=router_fallback_access_check,
+            fallback_budget_check=router_fallback_budget_check,
             auto_router_capability_limit=_license_check.auto_router_capability_limit,
         )
 
@@ -6622,6 +6624,7 @@ class ProxyConfig:
                         search_tools=search_tools,
                         ignore_invalid_deployments=True,
                         fallback_access_check=router_fallback_access_check,
+                        fallback_budget_check=router_fallback_budget_check,
                         auto_router_capability_limit=_license_check.auto_router_capability_limit,
                     )
                     verbose_proxy_logger.debug("updated llm_router: %s", llm_router)

litellm/proxy/schema.prisma

@@ -426,6 +426,7 @@ model LiteLLM_VerificationToken {
     key_alias   String?
     soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down
     spend      Float    @default(0.0)
+    total_spend Float   @default(0.0)
     expires    DateTime?
     models     String[]
     aliases    Json  @default("{}")
@@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken {
     key_alias               String?
     soft_budget_cooldown    Boolean  @default(false)
     spend                   Float    @default(0.0)
+    total_spend             Float    @default(0.0)
     expires                 DateTime?
     models                  String[]
     aliases                 Json     @default("{}")

litellm/proxy/spend_tracking/carried_budget_state.py

@@ -25,6 +25,7 @@ def carry_team_and_user_budget_state(
             budget_reset_at=team_object.budget_reset_at,
             max_budget=team_object.max_budget,
         )
+        valid_token.team_model_max_budget = team_object.model_max_budget  # rebind-ok: caller keeps this object
     if user_object is not None:
         valid_token.user_budget_snapshot = UserBudgetSnapshot(  # rebind-ok: same object the caller keeps using
             budget_reset_at=user_object.budget_reset_at,

litellm/proxy/utils.py

@@ -164,7 +164,6 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai
 )
 from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook
 from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
-from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
 from litellm.proxy.hooks.parallel_request_limiter import (
     _PROXY_MaxParallelRequestsHandler,
 )
@@ -982,7 +981,6 @@ class ProxyLogging:
             dual_cache=DualCache(default_in_memory_ttl=1)  # ping redis cache every 1s
         )
         self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache)
-        self.max_budget_limiter = _PROXY_MaxBudgetLimiter()
         self.cache_control_check = _PROXY_CacheControlCheck()
         self.alerting: list[str] | None = None
         self.alerting_threshold: float = 300  # default to 5 min. threshold
@@ -3580,7 +3578,7 @@ class ProxyLogging:
         caps: Final = ProxyLogging._callback_capabilities()
         post_call_pipelines: Final = _streamable_post_call_pipelines(request_data, user_api_key_dict)
         # Fast path: no real overrides. Internal proxy CustomLogger callbacks
-        # (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default
+        # (e.g. _PROXY_CacheControlCheck, ManagedFiles) inherit the default
         # ``async for chunk: yield chunk`` body, so wrapping the iterator
         # through each of them adds N pass-through trampolines per chunk for
         # zero behavior change. Skip the chain entirely and stream through.
@@ -4340,6 +4338,7 @@ class PrismaClient:
                             v.*,
                             t.spend AS team_spend,
                             t.max_budget AS team_max_budget,
+                            t.model_max_budget AS team_model_max_budget,
                             t.tpm_limit AS team_tpm_limit,
                             t.rpm_limit AS team_rpm_limit,
                             t.tpd_limit AS team_tpd_limit
@@ -4779,6 +4778,7 @@ class PrismaClient:
                             t.spend AS team_spend, 
                             t.max_budget AS team_max_budget,
                             t.soft_budget AS team_soft_budget,
+                            t.model_max_budget AS team_model_max_budget,
                             t.tpm_limit AS team_tpm_limit,
                             t.rpm_limit AS team_rpm_limit,
                             t.tpd_limit AS team_tpd_limit,

litellm/router.py

@@ -63,6 +63,7 @@ from litellm.constants import (
     DEFAULT_MAX_LRU_CACHE_SIZE,
     INTERNAL_CALL_ORIGIN_METADATA_KEY,
     OUTPUT_TOKEN_CEILING_PARAMS,
+    ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
     ROUTING_REQUEST_TAGS_METADATA_KEY,
     RUNTIME_UPDATABLE_ROUTER_SETTINGS,
     SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
@@ -132,7 +133,7 @@ from litellm.router_utils.add_retry_fallback_headers import (
     get_hidden_params_dict,
     prepare_response_for_header_attachment,
     replace_complexity_router_headers,
-    response_in_flight_token_count,
+    response_total_token_count,
 )
 from litellm.router_utils.auto_router_model_naming import (
     AUTO_ROUTER_MODEL_PREFIX,
@@ -215,6 +216,8 @@ from litellm.router_utils.reasoning_effort_capability import (
     resolve_supported_reasoning_efforts,
 )
 from litellm.router_utils.router_callbacks.track_deployment_metrics import (
+    find_deployment_metadata,
+    get_counted_usage_tokens,
     increment_deployment_failures_for_current_minute,
     increment_deployment_successes_for_current_minute,
 )
@@ -240,6 +243,7 @@ from litellm.types.router import (
     DeploymentModelListingInfo,
     DeploymentTypedDict,
     FallbackAccessCheck,
+    FallbackBudgetCheck,
     GuardrailTypedDict,
     LiteLLM_Params,
     MockRouterTestingParams,
@@ -777,6 +781,7 @@ class Router:
         background_health_check_model_groups: Sequence[str] | None = None,
         enable_weighted_failover: bool = False,
         fallback_access_check: FallbackAccessCheck | None = None,
+        fallback_budget_check: FallbackBudgetCheck | None = None,
         auto_router_capability_limit: AutoRouterCapabilityLimit | None = None,
     ) -> None:
         """
@@ -815,6 +820,7 @@ class Router:
             ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error.
             enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False.
             fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted).
+            fallback_budget_check (Optional[FallbackBudgetCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over budget is skipped. Defaults to None (budget is not re-checked on fallback).
         Returns:
             Router: An instance of the litellm.Router class.
 
@@ -856,6 +862,7 @@ class Router:
         self.ignore_invalid_deployments = ignore_invalid_deployments
         self.auto_router_capability_limit = auto_router_capability_limit
         self.fallback_access_check: Final = fallback_access_check
+        self.fallback_budget_check: Final = fallback_budget_check
         self.debug_level = debug_level
         self.enable_pre_call_checks = enable_pre_call_checks
         self.enable_tag_filtering = enable_tag_filtering
@@ -7937,6 +7944,7 @@ class Router:
         response = original_function(*args, **kwargs)
         if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
             response = await response
+        await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs)
         ## PROCESS RESPONSE HEADERS
         response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs)
 
@@ -8153,8 +8161,6 @@ class Router:
         """
         Track remaining tpm/rpm quota for model in model_list
         """
-        from litellm.types.caching import RedisPipelineIncrementOperation
-
         try:
             # WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
             if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
@@ -8162,114 +8168,135 @@ class Router:
             standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
             if standard_logging_object is None:
                 raise ValueError("standard_logging_object is None")
-            if kwargs["litellm_params"].get("metadata") is None:
-                pass
-            else:
-                deployment_name: Final = kwargs["litellm_params"]["metadata"].get(
-                    "deployment", None
-                )  # stable name - works for wildcard routes as well
-                # Get model_group and id from kwargs like the sync version does
-                model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
-                model_info: Final = kwargs["litellm_params"].get("model_info", {}) or {}
-                id = model_info.get("id", None)
-                if model_group is None or id is None:
-                    return
-                elif isinstance(id, int):
-                    id = str(id)
+            litellm_params: Final = kwargs["litellm_params"]
+            metadata: Final = litellm_params.get("metadata")
+            if metadata is None:
+                return
+            model_group: Final = metadata.get("model_group", None)
+            model_info: Final = litellm_params.get("model_info", {}) or {}
+            deployment_id: Final = model_info.get("id", None)
+            if model_group is None or deployment_id is None or self.get_deployment(model_id=str(deployment_id)) is None:
+                return
 
-                ## get deployment info
-                deployment_info: Final = self.get_deployment(model_id=id)
+            # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits
+            increment_deployment_successes_for_current_minute(
+                litellm_router_instance=self,
+                deployment_id=str(deployment_id),
+            )
 
-                if deployment_info is None:
-                    return
-                else:
-                    deployment_model_info: Final = self.get_router_model_info(
-                        deployment=deployment_info,
-                        received_model_name=model_group,
-                    )
-                    # get tpm/rpm from deployment info
-                    tpm: Final = deployment_info.get("tpm", None)
-                    rpm: Final = deployment_info.get("rpm", None)
-
-                    ## check tpm/rpm in litellm_params
-                    tpm_litellm_params: Final = deployment_info.litellm_params.tpm
-                    rpm_litellm_params: Final = deployment_info.litellm_params.rpm
-
-                    ## check tpm/rpm in model_info
-                    tpm_model_info: Final = deployment_model_info.get("tpm", None)
-                    rpm_model_info: Final = deployment_model_info.get("rpm", None)
-
-                # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits
-                increment_deployment_successes_for_current_minute(
-                    litellm_router_instance=self,
-                    deployment_id=id,
-                )
-
-                deployment_dict = deployment_info if isinstance(deployment_info, dict) else deployment_info.model_dump()
-                has_io_token_limits: Final = deployment_has_io_token_limits(deployment_dict)
-
-                ## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are
-                ## set. IO deployments still record TPM/RPM usage here so TPM-aware
-                ## routing strategies see their real load in mixed model groups; their
-                ## itpm/otpm enforcement runs separately in ModelRateLimitingCheck.
-                if (
-                    tpm is None
-                    and rpm is None
-                    and tpm_litellm_params is None
-                    and rpm_litellm_params is None
-                    and tpm_model_info is None
-                    and rpm_model_info is None
-                    and not has_io_token_limits
-                ):
-                    return
-
-                parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs)
-                total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0)
-
-                # ------------
-                # Setup values
-                # ------------
-                dt: Final = get_utc_datetime()
-                current_minute: Final = dt.strftime("%H-%M")  # use the same timezone regardless of system clock
-
-                tpm_key = RouterCacheEnum.TPM.value.format(id=id, current_minute=current_minute, model=deployment_name)
-                # ------------
-                # Update usage
-                # ------------
-                # update cache
-                pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = []
-
-                ## TPM
-                pipeline_operations.append(
-                    RedisPipelineIncrementOperation(
-                        key=tpm_key,
-                        increment_value=total_tokens,
-                        ttl=RoutingArgs.ttl.value,
-                    )
-                )
-
-                ## RPM
-                rpm_key = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name)
-                pipeline_operations.append(
-                    RedisPipelineIncrementOperation(
-                        key=rpm_key,
-                        increment_value=1,
-                        ttl=RoutingArgs.ttl.value,
-                    )
-                )
-
-                await self.cache.async_increment_cache_pipeline(
-                    increment_list=pipeline_operations,
-                    parent_otel_span=parent_otel_span,
-                )
-
-                return tpm_key
+            total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0)
+            counted_tokens: Final = get_counted_usage_tokens(litellm_params)
+            deployment_name: Final = metadata.get("deployment", None)
+            return await self._increment_deployment_usage(
+                deployment_id=str(deployment_id),
+                deployment_name=deployment_name if isinstance(deployment_name, str) else None,
+                model_group=model_group,
+                total_tokens=total_tokens if counted_tokens is None else max(0, total_tokens - counted_tokens),
+                rpm_increment=1 if counted_tokens is None else 0,
+                parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs),
+            )
 
         except Exception as e:
             verbose_router_logger.debug(
                 "litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e
             )
 
+    async def increment_deployment_usage_for_response(
+        self,
+        response: object,
+        request_kwargs: dict[str, object],
+    ) -> None:
+        if response is None:
+            return
+        try:
+            deployment_metadata: Final = find_deployment_metadata(request_kwargs)
+            model_group: Final = request_kwargs.get("model")
+            if deployment_metadata is None or not isinstance(model_group, str):
+                return
+            model_info: Final = deployment_metadata["model_info"]
+            deployment_id: Final = model_info.get("id") if isinstance(model_info, dict) else None
+            if deployment_id is None:
+                return
+            total_tokens: Final = response_total_token_count(response)
+            deployment_name: Final = deployment_metadata.get("deployment")
+            deployment_metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] = total_tokens
+            try:
+                await self._increment_deployment_usage(
+                    deployment_id=str(deployment_id),
+                    deployment_name=deployment_name if isinstance(deployment_name, str) else None,
+                    model_group=model_group,
+                    total_tokens=total_tokens,
+                    rpm_increment=1,
+                    parent_otel_span=_get_parent_otel_span_from_kwargs(request_kwargs),
+                )
+            except Exception:
+                deployment_metadata.pop(ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, None)
+                raise
+        except Exception as e:
+            verbose_router_logger.debug(
+                "litellm.router.Router::increment_deployment_usage_for_response(): Exception occured - %s", e
+            )
+
+    async def _increment_deployment_usage(
+        self,
+        *,
+        deployment_id: str,
+        deployment_name: str | None,
+        model_group: str,
+        total_tokens: float,
+        rpm_increment: int,
+        parent_otel_span: Span | None,
+    ) -> str | None:
+        from litellm.types.caching import RedisPipelineIncrementOperation
+
+        deployment_info: Final = self.get_deployment(model_id=deployment_id)
+        if deployment_info is None:
+            return None
+        deployment_model_info: Final = self.get_router_model_info(
+            deployment=deployment_info,
+            received_model_name=model_group,
+        )
+        configured_limits: Final = (
+            deployment_info.get("tpm", None),
+            deployment_info.get("rpm", None),
+            deployment_info.litellm_params.tpm,
+            deployment_info.litellm_params.rpm,
+            deployment_model_info.get("tpm", None),
+            deployment_model_info.get("rpm", None),
+        )
+        ## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are
+        ## set. IO deployments still record TPM/RPM usage here so TPM-aware
+        ## routing strategies see their real load in mixed model groups; their
+        ## itpm/otpm enforcement runs separately in ModelRateLimitingCheck.
+        if all(limit is None for limit in configured_limits) and not deployment_has_io_token_limits(
+            deployment_info.model_dump()
+        ):
+            return None
+        if total_tokens <= 0 and rpm_increment <= 0:
+            return None
+
+        current_minute: Final = get_utc_datetime().strftime("%H-%M")  # use the same timezone regardless of system clock
+        tpm_key: Final = RouterCacheEnum.TPM.value.format(
+            id=deployment_id, current_minute=current_minute, model=deployment_name
+        )
+        rpm_key: Final = RouterCacheEnum.RPM.value.format(
+            id=deployment_id, current_minute=current_minute, model=deployment_name
+        )
+        pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [
+            RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value)
+            for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment))
+        ]
+        post_increment_values: Final = await self.cache.async_increment_cache_pipeline(
+            increment_list=pipeline_operations,
+            parent_otel_span=parent_otel_span,
+        )
+        if post_increment_values is not None and self.cache.redis_cache is not None:
+            for operation, value in zip(pipeline_operations, post_increment_values):
+                await self.cache.async_set_cache(
+                    operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value
+                )
+        return tpm_key
+
     def sync_deployment_callback_on_success(
         self,
         kwargs,  # kwargs to completion
@@ -11205,15 +11232,7 @@ class Router:
 
         if model_group is not None:
             remaining_usage: Final = await self.get_remaining_model_group_usage(model_group)
-            # get_remaining_model_group_usage reads the router's TPM/RPM counter,
-            # which is incremented post-response by deployment_callback_on_success.
-            # Replay the in-flight increment for TPM/RPM only (LIT-2719); ITPM/OTPM
-            # counters are incremented at reservation time and must not be adjusted.
-            apply_remaining_usage_headers(
-                additional_headers,
-                remaining_usage,
-                response_in_flight_token_count(response),
-            )
+            apply_remaining_usage_headers(additional_headers, remaining_usage)
         return response
 
     def _build_model_name_index(self, model_list: list) -> None:

litellm/router_utils/add_retry_fallback_headers.py

@@ -151,7 +151,7 @@ def apply_quality_router_decision_headers(
             additional_headers[header] = str(decision[field])
 
 
-def response_in_flight_token_count(response: object) -> int:
+def response_total_token_count(response: object) -> int:
     usage: Final = response.get("usage") if isinstance(response, dict) else getattr(response, "usage", None)
     if usage is None:
         return 0
@@ -166,15 +166,10 @@ def response_in_flight_token_count(response: object) -> int:
 def apply_remaining_usage_headers(
     additional_headers: dict[str, object],
     remaining_usage: dict[str, int],
-    in_flight_tokens: int,
 ) -> None:
-    in_flight_delta: Final = {
-        "x-ratelimit-remaining-tokens": in_flight_tokens,
-        "x-ratelimit-remaining-requests": 1,
-    }
     for header, value in remaining_usage.items():
         if value is not None and header not in additional_headers:
-            additional_headers[header] = value - in_flight_delta.get(header, 0)
+            additional_headers[header] = value
 
 
 def _normalize_hidden_params(hidden_params: object) -> dict[str, object]:

litellm/router_utils/fallback_event_handlers.py

@@ -421,6 +421,25 @@ async def _is_fallback_target_authorized(
     return False
 
 
+async def _is_fallback_target_within_budget(
+    litellm_router: LitellmRouter,
+    fallback_entry: str | Mapping[str, object],
+    original_model_group: str,
+    kwargs: Mapping[str, object],
+) -> bool:
+    budget_check: Final = litellm_router.fallback_budget_check
+    target: Final = _get_fallback_target_model_group(fallback_entry)
+    if budget_check is None or target is None or target == original_model_group:
+        return True
+    if await budget_check(model=target, request_kwargs=kwargs, llm_router=litellm_router):
+        return True
+    verbose_router_logger.info(
+        "Skipping fallback to model_group = %s: caller is over budget",
+        mask_sensitive_structure(fallback_entry),
+    )
+    return False
+
+
 def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool:
     """
     True when a file, batch, or fine-tuning job operation names an id that only exists
@@ -528,6 +547,8 @@ async def run_async_fallback(
             continue
         if not await _is_fallback_target_authorized(litellm_router, mg, original_model_group, kwargs):
             continue
+        if not await _is_fallback_target_within_budget(litellm_router, mg, original_model_group, kwargs):
+            continue
         attempt_key = fallback_attempt_key(mg)
         if attempt_key is not None:
             if attempt_key in attempted:

litellm/router_utils/router_callbacks/track_deployment_metrics.py

@@ -9,8 +9,11 @@ get_deployment_failures_for_current_minute
 get_deployment_successes_for_current_minute
 """
 
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Final
 
+from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
+
 if TYPE_CHECKING:
     from litellm.router import Router as _Router
 
@@ -18,6 +21,26 @@ if TYPE_CHECKING:
 else:
     LitellmRouter = Any
 
+_METADATA_CHANNELS: Final = ("litellm_metadata", "metadata")
+
+
+def find_deployment_metadata(kwargs: Mapping[str, object]) -> dict[str, object] | None:
+    buckets: Final = (kwargs.get(channel) for channel in _METADATA_CHANNELS)
+    return next((bucket for bucket in buckets if isinstance(bucket, dict) and "model_info" in bucket), None)
+
+
+def get_counted_usage_tokens(litellm_params: Mapping[str, object]) -> int | None:
+    buckets: Final = (litellm_params.get(channel) for channel in _METADATA_CHANNELS)
+    counted: Final = next(
+        (
+            bucket[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY]
+            for bucket in buckets
+            if isinstance(bucket, dict) and ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY in bucket
+        ),
+        None,
+    )
+    return counted if isinstance(counted, int) and not isinstance(counted, bool) else None
+
 
 def increment_deployment_successes_for_current_minute(
     litellm_router_instance: LitellmRouter,

litellm/rust_bridge/_native.pyi

@@ -33,15 +33,6 @@ def aocr(
     timeout_seconds: float | None = None,
 ) -> Future[dict[str, object]]: ...
 
-_OCR_MAX_FILE_BYTES: int
-
-def _ocr_upload_document(
-    file_content: bytes,
-    file_name: str | None = None,
-    content_type: str | None = None,
-) -> dict[str, str]: ...
-def _ocr_file_document(document: Mapping[str, object]) -> dict[str, str]: ...
-def _ocr_mime_type(file_name: str) -> str: ...
 def _ocr_lifecycle(
     request: LiteLLMOcrRequest,
     args: tuple[object, ...],
@@ -139,15 +130,11 @@ class TokenCounter:
 def gil_stats() -> dict[str, int]: ...
 
 __all__ = [
-    "_OCR_MAX_FILE_BYTES",
     "ResponsesWebSocketConnection",
     "RustBridgeDeclined",
     "RustUpstreamError",
     "TokenCounter",
-    "_ocr_file_document",
     "_ocr_lifecycle",
-    "_ocr_mime_type",
-    "_ocr_upload_document",
     "achat_completions",
     "amessages",
     "aocr",

litellm/types/guardrails.py

@@ -682,6 +682,12 @@ class BedrockGuardrailStreamingParams(BaseModel):
         "and the scan result lands in guardrail_information; a flagged response still ends the "
         "stream with a block message (disable_exception_on_block=true) or an error frame.",
     )
+    streaming_buffer_release_on_scan: bool = Field(
+        default=False,
+        description="When buffering, scan the accumulated response every streaming_sampling_rate chunks "
+        "and release the withheld chunks once the scan passes, instead of holding everything to end of stream. "
+        "Flagged content is never released. Ignored when streaming_end_of_stream_only is true.",
+    )
 
     @classmethod
     def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams":

litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py

@@ -4,6 +4,14 @@ from .base import GuardrailConfigModel
 
 
 class CrowdStrikeAIDRGuardrailConfigModelOptionalParams(BaseModel):
+    streaming_buffer_until_moderated: bool | None = Field(
+        default=None,
+        description="When True, withhold streamed chunks until moderation passes. Defaults to False when unset.",
+    )
+    streaming_buffer_release_on_scan: bool | None = Field(
+        default=None,
+        description="When buffering, release withheld chunks after each passing scan. Defaults to False when unset.",
+    )
     streaming_end_of_stream_only: bool | None = Field(
         default=None,
         description="If False (default when unset), post_call scans the accumulated streamed response every "

litellm/types/router.py

@@ -963,6 +963,19 @@ class FallbackAccessCheck(Protocol):
     async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ...
 
 
+class FallbackBudgetCheck(Protocol):
+    """
+    Decides whether the caller behind `request_kwargs` is still within budget for fallback `model`.
+
+    Budget is enforced once during auth, against the *requested* model group. A fallback target is
+    chosen later, inside the router, so a zero-cost group that falls back to a priced one bills
+    without any budget gate. The router runs this before every cross-model-group fallback attempt
+    and skips targets it rejects, leaving the free attempt itself untouched.
+    """
+
+    async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ...
+
+
 class AutoRouterCapabilityLimit(Protocol):
     """
     Resolves how many complexity routers may claim each licensed capability right now; None means unlimited.

schema.prisma

@@ -426,6 +426,7 @@ model LiteLLM_VerificationToken {
     key_alias   String?
     soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down
     spend      Float    @default(0.0)
+    total_spend Float   @default(0.0)
     expires    DateTime?
     models     String[]
     aliases    Json  @default("{}")
@@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken {
     key_alias               String?
     soft_budget_cooldown    Boolean  @default(false)
     spend                   Float    @default(0.0)
+    total_spend             Float    @default(0.0)
     expires                 DateTime?
     models                  String[]
     aliases                 Json     @default("{}")

tests/_live_test_helpers.py

@@ -1,6 +1,8 @@
 import os
+from datetime import date
 
 import pytest
+from pydantic import BaseModel, ConfigDict
 
 
 def _skip_live_prompt_caching_test():
@@ -8,3 +10,55 @@ def _skip_live_prompt_caching_test():
         pytest.skip("Live prompt-caching E2E tests are opt-in")
     if os.environ.get("CASSETTE_REDIS_URL"):
         pytest.skip("Live prompt-caching E2E tests cannot run under VCR replay")
+
+
+
+class TogetherCostEntry(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+
+    litellm_provider: str | None = None
+    mode: str | None = None
+    deprecation_date: str | None = None
+    input_cost_per_token: float | None = None
+    output_cost_per_token: float | None = None
+    supports_function_calling: bool | None = None
+    supports_response_schema: bool | None = None
+
+
+def cheapest_together_chat_model(
+    *, function_calling: bool = False, response_schema: bool = False
+) -> str:
+    import litellm
+
+    today = date.today().isoformat()
+
+    def qualifies(name: str, entry: TogetherCostEntry) -> bool:
+        return (
+            name.startswith("together_ai/")
+            and entry.litellm_provider == "together_ai"
+            and entry.mode == "chat"
+            and (entry.deprecation_date is None or entry.deprecation_date > today)
+            and (entry.input_cost_per_token or 0.0) > 0
+            and (entry.output_cost_per_token or 0.0) > 0
+            and (not function_calling or bool(entry.supports_function_calling))
+            and (not response_schema or bool(entry.supports_response_schema))
+        )
+
+    registry: dict[str, TogetherCostEntry] = {
+        name: TogetherCostEntry.model_validate(raw)
+        for name, raw in litellm.model_cost.items()
+        if isinstance(raw, dict) and name.startswith("together_ai/")
+    }
+    candidates = sorted(
+        (name for name, entry in registry.items() if qualifies(name, entry)),
+        key=lambda name: (
+            registry[name].input_cost_per_token or 0.0,
+            registry[name].output_cost_per_token or 0.0,
+            name,
+        ),
+    )
+    assert candidates, (
+        "no live together_ai chat model in the cost map satisfies "
+        f"function_calling={function_calling} response_schema={response_schema}"
+    )
+    return candidates[0]

tests/e2e/PROVIDER_CACHE.md

@@ -1,11 +1,41 @@
 # Shared provider-response cache
 
-`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live
+`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live
 
-The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away. Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies
+The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored
+
+Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic
+
+Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way
+
+## Request identity
+
+A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is
+
+Requests that differ only by their markers therefore share a canonical identity, which is what makes the cache reusable across builds: every e2e test salts its prompt afresh, so an exact-byte key would miss on every call. Within one test, calls that share a canonical identity are still recorded and replayed separately, by a FIFO slot index appended to the key. That matters because a replayed response carries the recorded provider response id, `LiteLLM_SpendLogs.request_id` is that id, and one shared recording answering two calls would collapse two spend rows into one
+
+Two different tests never share a recording, and a provider call made outside any test (fixtures, session setup) is never cached, because the identity has no test node id to bind to
+
+A client that varies its own request between runs defeats that identity without breaking any rule, and the Claude Code compat cells did. The CLI sends a device id and a session id in `metadata.user_id`, and its system prompt names both its memory directory and its working directory, adding the branch and recent commits when that directory is a git repository. Driven with a fresh HOME and the checkout as its working directory, every cell sent different bytes every build. The fix belongs in the driver rather than here: `claude_code/cli_driver.py` pins the config directory, the working directory and both identifiers, which is why the cache needs no rule for any of it. Normalizing them instead would have hidden a real defect class, since a rule cannot tell a client's own churn from a value a test means to assert on
+
+Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies
 
 An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure
 
+## Bedrock
+
+Bedrock could not be mounted before because SigV4 signs the `Host` header, so a rewritten `api_base` failed signature verification at the provider. The edge now re-signs: it drops the proxy's signature headers, signs the upstream request with the run pod's own AWS identity from its EKS Pod Identity association, and forwards that. The signature headers are excluded from the key, since `x-amz-date` is a timestamp and keying on it would make every Bedrock call a permanent miss
+
+Almost every Bedrock deployment in the suite declares its region as `os.environ/AWS_REGION`, which only the proxy can resolve, and the run pod does not share that environment. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in.
+
+Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove
+
+Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here
+
+Vertex and Gemini are not mounted, for different reasons. litellm grafts the default Vertex path onto an `api_base` only when that `api_base` has no path of its own, so a path-prefixed Vertex mount instead becomes `{api_base}:{endpoint}`, dropping project, location and model. Vertex needs a root-mounted edge on its own port, or a change in litellm
+
+Gemini reaches a path-prefixed mount perfectly well and was mounted for one build, then backed out, because litellm's two Gemini endpoints disagree about what `api_base` means. Chat composes `{api_base}/models/{model}:{endpoint}` and defaults `api_base` to `https://generativelanguage.googleapis.com/v1beta`, so the version has to be inside it. File upload composes `{api_base}/upload/v1beta/files` and defaults to the host root, so the version has to be outside it. One `api_base` cannot satisfy both, and a deployment gives no signal at registration time about which it will be used for, so mounting Gemini turned `TestGeminiFiles::test_gemini_file_upload` red in build 227. Anyone pointing litellm's Gemini provider at an AI gateway or a corporate proxy hits the same thing; it is a litellm bug rather than a cache limitation, and mounting Gemini is one line once it is fixed
+
 Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires
 
 ## Configuration
@@ -18,16 +48,18 @@ The trusted runner receives:
 - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision
 - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory
 
-Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits
+Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all of one kind is a different problem from one whose rejections are nearly all of another, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits
 
-Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay
+Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read.
+
+One more class needs it, and it is the cost of normalizing the marker. A test that mints a fresh marker, sends it, and then asserts the provider's answer contains that exact value is asserting on the marker rather than using it as a salt. The key treats two such requests as the same identity, so a stale recording matches and answers with the marker from the run that recorded it. `TestOpenAIMessagesToolContinuation` is the one in the suite today: it sends a freshly minted receipt through a tool result and asserts the model echoes it back verbatim. If you add a test that asserts a provider echoed your own unique value, it belongs on the live path. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay
 
 ## Recorded response semantics
 
-Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Existing spend reconciliation requests use distinct prompt markers and retain their distinct-ID and row-count assertions; accounting tests are not automatically excluded from caching
+Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Spend reconciliation keeps its distinct-ID and row-count assertions: its prompts differ by an index as well as a marker, so they stay distinct once markers are normalized, and calls that are canonically equal within one test take separate FIFO slots and separate recordings anyway. Accounting tests are not automatically excluded from caching
 
 Provider remaining-quota headers describe the captured response. Metrics derived from them are historical on a cache hit, not a measurement of current provider capacity. Gateway-generated API-key quota headers are a separate contract. A test of fresh provider quota or timing must use the live-provider policy; replay can still exercise how the proxy processes the recorded headers
 
 ## Qualification
 
-`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence
+`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis, including the marker-canonical key, the FIFO slot index, per-test isolation, SigV4 re-signing against a local upstream, and each endpoint's completeness rule. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence

tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py

@@ -0,0 +1,162 @@
+"""The CLI must send the same request bytes from one build to the next.
+
+Markerless harness test: it drives the real `claude` binary against a local
+stub instead of a proxy, so it carries no `e2e` marker. The binary is a
+prerequisite of this whole suite, so a missing one is a failure rather than a
+skip.
+
+Two builds differ in ways the driver does not control: a fresh pod, so no CLI
+state survives, and a different candidate checked out at a different commit.
+Both used to reach the request body, through the memory path the system prompt
+names and through the git block the CLI adds for its working directory, so the
+shared provider cache missed on every Claude Code cell. This replays those two
+differences across a pair of invocations and holds the bytes equal.
+
+A pinned session id is what makes the second test necessary. The matrix runs
+its cells across xdist workers, and the CLI refuses to start a session id that
+another live process already holds, so pinning one without also opting out of
+session persistence turns most of a parallel run red.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+import threading
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from typing import List, Tuple
+
+import pytest
+
+from claude_code.cli_driver import _FIXED_CLI_USER_ID, _seed_cli_identity, _stable_cli_state, run_claude
+from claude_code.rate_limiter import RateLimiter
+
+pytestmark = pytest.mark.cli_determinism
+
+_STUB_REPLY = {
+    "id": "msg_stub",
+    "type": "message",
+    "role": "assistant",
+    "model": "claude-haiku-4-5",
+    "content": [{"type": "text", "text": "ok"}],
+    "stop_reason": "end_turn",
+    "usage": {"input_tokens": 10, "output_tokens": 2},
+}
+
+
+def _make_repo(root: Path, subject: str) -> Path:
+    root.mkdir(parents=True, exist_ok=True)
+    identity = {"NAME": "t", "EMAIL": "t@e2e"}
+    env = dict(
+        os.environ,
+        **{f"GIT_{role}_{key}": value for role in ("AUTHOR", "COMMITTER") for key, value in identity.items()},
+    )
+    (root / "file.txt").write_text(subject, encoding="utf-8")
+    for args in (["init", "-q"], ["add", "."], ["commit", "-q", "-m", subject]):
+        subprocess.run(["git", *args], cwd=root, env=env, check=True, capture_output=True)
+    return root
+
+
+@pytest.fixture(name="captured")
+def _captured() -> Tuple[str, List[bytes]]:
+    bodies: List[bytes] = []
+    lock = threading.Lock()
+
+    class Handler(BaseHTTPRequestHandler):
+        protocol_version = "HTTP/1.1"
+
+        def do_POST(self) -> None:
+            raw = self.rfile.read(int(self.headers.get("content-length") or 0))
+            if "count_tokens" not in self.path:
+                with lock:
+                    bodies.append(raw)
+            payload = json.dumps({"input_tokens": 10} if "count_tokens" in self.path else _STUB_REPLY).encode()
+            self.send_response(200)
+            self.send_header("content-type", "application/json")
+            self.send_header("content-length", str(len(payload)))
+            self.end_headers()
+            self.wfile.write(payload)
+
+        def log_message(self, *_args: object) -> None:
+            return
+
+    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+    threading.Thread(target=server.serve_forever, daemon=True).start()
+    try:
+        yield f"http://127.0.0.1:{server.server_address[1]}", bodies
+    finally:
+        server.shutdown()
+
+
+def test_two_builds_send_the_same_request_bytes(captured: Tuple[str, List[bytes]], tmp_path: Path) -> None:
+    base_url, bodies = captured
+    limiter = RateLimiter(state_dir=tmp_path / "limiter")
+    checkouts = (_make_repo(tmp_path / "build-1", "first"), _make_repo(tmp_path / "build-2", "second"))
+    origin = Path.cwd()
+
+    sent = []
+    for checkout in checkouts:
+        shutil.rmtree(Path(_stable_cli_state()[0]).parent, ignore_errors=True)
+        os.chdir(checkout)
+        try:
+            before = len(bodies)
+            run_claude(
+                prompt="say ok",
+                model="claude-haiku-4-5",
+                base_url=base_url,
+                api_key="stub",
+                extra_env={"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"},
+                rate_limiter=limiter,
+            )
+            sent.append(bodies[before:])
+        finally:
+            os.chdir(origin)
+
+    assert sent[0], "the CLI sent no request to the stub, so there is nothing to compare"
+    assert sent[0] == sent[1]
+
+
+def test_concurrent_cells_do_not_collide_on_the_pinned_session(
+    captured: Tuple[str, List[bytes]], tmp_path: Path
+) -> None:
+    base_url, bodies = captured
+    limiter = RateLimiter(state_dir=tmp_path / "limiter")
+
+    def one(_index: int) -> int:
+        return run_claude(
+            prompt="say ok",
+            model="claude-haiku-4-5",
+            base_url=base_url,
+            api_key="stub",
+            extra_env={"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"},
+            rate_limiter=limiter,
+        ).exit_code
+
+    with ThreadPoolExecutor(max_workers=4) as pool:
+        codes = list(pool.map(one, range(4)))
+
+    assert codes == [0, 0, 0, 0]
+    assert bodies, "the CLI sent no request to the stub, so there is nothing to compare"
+    assert set(Counter(bodies).values()) == {4}
+
+
+def test_seeding_the_device_id_survives_threads_racing_on_the_same_directory(tmp_path: Path) -> None:
+    """`run_claude_models_parallel` drives several models from one process, so the
+    seed's staged file has to be unique per thread and not merely per process."""
+    config_dir = tmp_path / "config"
+    config_dir.mkdir()
+    seeded = config_dir / ".claude.json"
+
+    for _round in range(20):
+        seeded.unlink(missing_ok=True)
+        with ThreadPoolExecutor(max_workers=16) as pool:
+            for outcome in [pool.submit(_seed_cli_identity, str(config_dir)) for _ in range(16)]:
+                outcome.result()
+
+    assert json.loads(seeded.read_text(encoding="utf-8"))["userID"] == _FIXED_CLI_USER_ID
+    assert sorted(entry.name for entry in config_dir.iterdir()) == [".claude.json"]

tests/e2e/claude_code/cli_driver.py

@@ -132,6 +132,62 @@ def _make_isolated_home() -> str:
     return tempfile.mkdtemp(prefix="claude-cli-home-")
 
 
+_FIXED_CLI_USER_ID = "0" * 64
+_FIXED_CLI_SESSION_ID = "00000000-0000-4000-8000-000000000000"
+
+
+def _seed_cli_identity(config_dir: str) -> None:
+    """Pin the device id the CLI would otherwise mint per config directory.
+
+    It mints 32 random bytes on first run, writes them to `.claude.json` as
+    `userID`, and sends them in `metadata.user_id` forever after, so the value
+    is stable for exactly as long as that file lives. Pinning it, and the
+    session id passed beside it, costs nothing: both feed abuse detection
+    rather than quota, caching or continuity.
+
+    The staged name has to be unique per *thread*, not per process:
+    `run_claude_models_parallel` drives several models from one process, so a
+    pid-suffixed name lets one thread rename the file another is still
+    writing, and the loser dies on a missing path."""
+    path = os.path.join(config_dir, ".claude.json")
+    try:
+        with open(path, encoding="utf-8") as handle:
+            if json.load(handle).get("userID") == _FIXED_CLI_USER_ID:
+                return
+    except (OSError, ValueError):
+        pass
+    handle_fd, staged = tempfile.mkstemp(dir=config_dir, prefix=".claude.json.")
+    with os.fdopen(handle_fd, "w", encoding="utf-8") as handle:
+        json.dump({"userID": _FIXED_CLI_USER_ID}, handle)
+    os.replace(staged, path)
+
+
+def _stable_cli_state() -> Tuple[str, str]:
+    """Config directory and working directory for the CLI, at fixed paths.
+
+    Both reach the request body. The memory directory the system prompt
+    names is `$CLAUDE_CONFIG_DIR/projects/<cwd slug>/memory`, and a working
+    directory inside a git repository also contributes its branch and recent
+    commits. So a per-invocation config directory rewrites every body, and
+    inheriting the checkout rewrites every body once per candidate, which is
+    why the shared provider cache could never serve a Claude Code cell.
+    Pinning both makes the bodies repeatable across builds.
+
+    This narrows what survives rather than widening it: HOME stays fresh and
+    empty per invocation, so the isolation `_make_isolated_home` describes is
+    unchanged, and the CLI's own state no longer outlives the pod either. The
+    working directory is deliberately not the checkout, so a model-directed
+    `Read` sees an empty directory instead of the repository.
+    """
+    root = os.path.join(tempfile.gettempdir(), f"litellm-e2e-claude-{os.getuid()}")
+    config_dir = os.path.join(root, "config")
+    workspace = os.path.join(root, "workspace")
+    for path in (root, config_dir, workspace):
+        os.makedirs(path, mode=0o700, exist_ok=True)
+    _seed_cli_identity(config_dir)
+    return config_dir, workspace
+
+
 class ClaudeCLIError(RuntimeError):
     """Raised when the `claude` CLI cannot be invoked or returns a fatal error."""
 
@@ -222,6 +278,9 @@ def run_claude(
         "--verbose",
         "--model",
         model,
+        "--session-id",
+        _FIXED_CLI_SESSION_ID,
+        "--no-session-persistence",
     ]
     if extra_args:
         cmd.extend(extra_args)
@@ -244,6 +303,8 @@ def run_claude(
     # regardless of how the subprocess exits.
     isolated_home = _make_isolated_home()
     env["HOME"] = isolated_home
+    config_dir, workspace = _stable_cli_state()
+    env["CLAUDE_CONFIG_DIR"] = config_dir
     if extra_env:
         env.update(extra_env)
 
@@ -262,6 +323,7 @@ def run_claude(
             completed = run_fn(
                 cmd,
                 env=env,
+                cwd=workspace,
                 input=stdin_input,
                 capture_output=True,
                 text=True,

tests/e2e/conftest.py

@@ -23,6 +23,7 @@ from typing import Final
 import pytest
 import requests
 from e2e_config import (
+    CLI_DETERMINISM_OPT_IN_ENV,
     CONTROL_PLANE_BASE_URL,
     COST_MAP_OPT_IN_ENV,
     FIXTURE_DIR,
@@ -55,6 +56,7 @@ OPT_IN_MARKERS: Final = MappingProxyType(
         "prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV,
         "redis_chaos": REDIS_CHAOS_OPT_IN_ENV,
         "cost_map_stack": COST_MAP_OPT_IN_ENV,
+        "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV,
     }
 )
 
@@ -87,7 +89,11 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient)
 
 
 def pytest_configure(config: pytest.Config) -> None:
-    config.addinivalue_line("markers", "provider_live: requires actual provider timing, limits or state; bypass shared cache")
+    config.addinivalue_line(
+        "markers",
+        "provider_live: requires actual provider timing, limits, state, or a response that echoes this"
+        " run's own unique value; bypass shared cache",
+    )
     config.addinivalue_line(
         "markers",
         "e2e: live test that requires a running proxy and real provider keys",
@@ -118,6 +124,11 @@ def pytest_configure(config: pytest.Config) -> None:
         "prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including "
         "prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set",
     )
+    config.addinivalue_line(
+        "markers",
+        "cli_determinism: drives the real claude CLI for several seconds, which widens the window in which "
+        "another test's in-flight upstream call is attributed to it; deselected unless E2E_CLI_DETERMINISM is set",
+    )
     config.addinivalue_line(
         "markers",
         "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "

tests/e2e/e2e_config.py

@@ -145,6 +145,7 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY"
 MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK"
 PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK"
 REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS"
+<<<<<<< HEAD
 # The cost_calculation suite needs a proxy booted with LITELLM_MODEL_COST_MAP_URL
 # pointing at tests/e2e/cost_map.json (its whole map is test-owned rates) plus a
 # scripted-provider sidecar; deselected unless the opt-in env var is set.
@@ -161,6 +162,10 @@ SCRIPTED_PROVIDER_CONTROL_URL = os.environ.get(
 SCRIPTED_PROVIDER_PROXY_BASE = os.environ.get(
     "E2E_SCRIPTED_PROVIDER_PROXY_BASE", SCRIPTED_PROVIDER_CONTROL_URL
 ).rstrip("/")
+||||||| 930ec9643a
+=======
+CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM"
+CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM"
 ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
 ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
 ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))

tests/e2e/fixture_canonical.py

@@ -51,6 +51,9 @@ SECRET_FIELD_SUFFIXES: Final[tuple[str, ...]] = (
 )
 SECRET_PLACEHOLDER: Final = "<secret>"
 
+MARKER_PATTERN: Final = re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{12}(?![0-9a-fA-F])")
+MARKER_PLACEHOLDER: Final = "<marker>"
+
 PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = (
     (re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{64}(?![0-9a-fA-F])"), "<sha256>"),
     (
@@ -67,7 +70,7 @@ PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = (
         re.compile(r"\b(?:chatcmpl|msgbatch|msg|resp|batch|call|req|ftjob|gen|file)[-_][A-Za-z0-9]{8,}\b"),
         "<id>",
     ),
-    (re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{12}(?![0-9a-fA-F])"), "<marker>"),
+    (MARKER_PATTERN, MARKER_PLACEHOLDER),
 )
 
 

tests/e2e/llm_translation/test_messages_e2e.py

@@ -372,6 +372,7 @@ def _request_tool(
 
 
 class TestOpenAIMessagesToolContinuation:
+    @pytest.mark.provider_live
     @pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"])
     def test_required_tool_arguments_and_correlated_result(
         self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool

tests/e2e/models.py

@@ -951,6 +951,7 @@ class LiteLLMParamsBody(BaseModel):
     aws_access_key_id: str | None = None
     aws_secret_access_key: str | None = None
     aws_region_name: str | None = None
+    aws_bedrock_runtime_endpoint: str | None = None
     vertex_project: str | None = None
     vertex_location: str | None = None
     vertex_credentials: str | None = None

tests/e2e/provider_cache.py

@@ -4,14 +4,18 @@ import base64
 import hashlib
 import hmac
 import io
+import json
+import os
 import threading
 import time
 from collections.abc import Callable, Generator, Mapping
 from contextlib import closing
 from dataclasses import dataclass, field
+from types import MappingProxyType
 from typing import Final, Literal, Protocol
 from urllib.parse import urlsplit
 
+from botocore.eventstream import EventStreamBuffer, ParserError
 from e2e_http import (
     NetworkError,
     StreamChunk,
@@ -23,12 +27,36 @@ from e2e_http import (
     prepare_forward,
     primed_steps,
 )
+from fixture_canonical import MARKER_PATTERN, MARKER_PLACEHOLDER
+from fixture_mode import SESSION_TEST_KEY, current_test_key
 from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
 
 LIFETIME_SECONDS: Final = 86_400
 MAX_REQUEST_BYTES: Final = 256 * 1024
 MAX_RESPONSE_BYTES: Final = 8 * 1024 * 1024
 UNRECORDED_RESPONSE_HEADERS: Final = frozenset({"set-cookie"})
+SIGNATURE_HEADERS: Final = frozenset(
+    {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"}
+)
+BEDROCK_MOUNT_PREFIX: Final = "bedrock"
+BEDROCK_CONVERSE_SUFFIX: Final = "/converse"
+BEDROCK_INVOKE_SUFFIX: Final = "/invoke"
+BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream"
+BEDROCK_INVOKE_STREAM_SUFFIX: Final = "/invoke-with-response-stream"
+BEDROCK_SUFFIXES: Final = (
+    BEDROCK_CONVERSE_SUFFIX,
+    BEDROCK_INVOKE_SUFFIX,
+    BEDROCK_CONVERSE_STREAM_SUFFIX,
+    BEDROCK_INVOKE_STREAM_SUFFIX,
+)
+EVENTSTREAM_PRELUDE_BYTES: Final = 4
+CUT_SHORT: Final = "cut_short"
+INCOMPLETE: Final = "incomplete"
+UNREACHABLE: Final = "unreachable"
+ERROR_STATUS: Final = "error_status"
+EVENT_TYPE_HEADER: Final = ":event-type"
+EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str])
+OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"})
 JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
 
 
@@ -56,6 +84,24 @@ class CacheUnavailable:
 
 
 type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable
+type RequestSigner = Callable[[str, str, Mapping[str, str], bytes | None], dict[str, str]]
+
+
+@dataclass(frozen=True, slots=True)
+class MountPolicy:
+    """What a mount needs beyond plain forwarding.
+
+    ``sign`` mints a fresh credential over the upstream URL, for providers whose
+    auth covers the Host the edge rewrote. ``unkeyed_headers`` names headers that
+    must stay out of the cache key because they change on every call and would
+    otherwise make the mount a permanent miss: a minted signature, or an OAuth
+    token the provider rotates. Naming one costs the guarantee that a recording
+    can never cross credentials, so a mount with a rotating token relies on the
+    environment holding one identity for that provider. Mounts with a static API
+    key name nothing here and keep the guarantee whole."""
+
+    sign: RequestSigner | None = None
+    unkeyed_headers: frozenset[str] = frozenset()
 
 
 class ResponseStore(Protocol):
@@ -83,28 +129,51 @@ class SignedResponse(BaseModel):
     signature: str
 
 
-def exact_key(secret: bytes, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> str:
+def canonical_text(value: str) -> str:
+    return MARKER_PATTERN.sub(MARKER_PLACEHOLDER, value)
+
+
+def canonical_body(body: bytes) -> bytes:
+    try:
+        return canonical_text(body.decode("utf-8")).encode("utf-8")
+    except UnicodeDecodeError:
+        return body
+
+
+def request_identity(
+    secret: bytes, test_key: str, method: str, url: str, headers: Mapping[str, str], body: bytes | None,
+) -> str:
     fields: Final = (
-        b"provider-cache-exact-v1", method.encode(), url.encode(),
+        b"provider-cache-canonical-v2", test_key.encode(), method.encode(), canonical_text(url).encode(),
         *(part.encode() for pair in sorted(headers.items()) for part in pair),
-        b"no-body" if body is None else b"body", b"" if body is None else body,
+        b"no-body" if body is None else b"body", b"" if body is None else canonical_body(body),
     )
     encoded: Final = b"".join(len(part).to_bytes(8, "big") + part for part in fields)
     return hmac.new(secret, encoded, hashlib.sha256).hexdigest()
 
 
-def cacheable_endpoint(method: str, url: str, body: bytes | None) -> bool:
-    return (
-        method == "POST"
-        and urlsplit(url).path in {"/v1/chat/completions", "/v1/messages"}
-        and body is not None
-        and len(body) <= MAX_REQUEST_BYTES
-    )
+def slotted_key(secret: bytes, identity: str, slot: int) -> str:
+    return hmac.new(secret, f"{identity}:{slot}".encode(), hashlib.sha256).hexdigest()
 
 
-def successful_response(url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool:
+def is_bedrock(mount: str) -> bool:
+    return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX
+
+
+def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool:
+    if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES:
+        return False
+    path: Final = urlsplit(url).path
+    if is_bedrock(mount):
+        return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES)
+    return path in OPENAI_JSON_PATHS
+
+
+def successful_response(mount: str, url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool:
     if not 200 <= status < 300 or len(body) > MAX_RESPONSE_BYTES:
         return False
+    if is_bedrock(mount):
+        return complete_bedrock_response(url, body)
     streaming: Final = "text/event-stream" in headers.get("content-type", "").lower()
     if streaming:
         try:
@@ -118,28 +187,33 @@ def successful_response(url: str, status: int, headers: Mapping[str, str], body:
             values: Final = tuple(JSON_VALUE.validate_json(event) for event in events if event != "[DONE]")
         except (UnicodeDecodeError, ValidationError):
             return False
-        if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values):
+        if not values or any(
+            not isinstance(value, dict) or value.get("error") is not None or value.get("type") == "error"
+            for value in values
+        ):
             return False
+        if urlsplit(url).path == "/v1/responses":
+            return complete_responses_stream(values)
         if urlsplit(url).path == "/v1/chat/completions":
             return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values)
-        return (
-            "[DONE]" not in events
-            and isinstance(values[0], dict) and values[0].get("type") == "message_start"
-            and isinstance(values[-1], dict) and values[-1].get("type") == "message_stop"
-            and any(
-                isinstance(value, dict) and value.get("type") == "message_delta"
-                and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str)
-                for value in values
-            )
-        )
+        return "[DONE]" not in events and complete_anthropic_stream(values)
     try:
         value: Final = JSON_VALUE.validate_json(body)
     except ValidationError:
         return False
-    if not isinstance(value, dict) or "error" in value:
+    if not isinstance(value, dict) or value.get("error") is not None:
         return False
-    if urlsplit(url).path == "/v1/messages":
+    path: Final = urlsplit(url).path
+    if path == "/v1/messages":
         return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str)
+    if path == "/v1/embeddings":
+        data: Final = value.get("data")
+        return isinstance(data, list) and bool(data) and isinstance(value.get("usage"), dict) and all(
+            isinstance(item, dict) and isinstance(item.get("embedding"), list) and bool(item["embedding"])
+            for item in data
+        )
+    if path == "/v1/responses":
+        return value.get("object") == "response" and value.get("status") == "completed"
     choices: Final = value.get("choices")
     return isinstance(choices, list) and bool(choices) and all(
         isinstance(choice, dict) and isinstance(choice.get("message"), dict) and isinstance(choice.get("finish_reason"), str)
@@ -147,6 +221,144 @@ def successful_response(url: str, status: int, headers: Mapping[str, str], body:
     )
 
 
+def complete_bedrock_response(url: str, body: bytes) -> bool:
+    """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an
+    Anthropic model answers the Anthropic message shape. Either way a truncated
+    or error body is missing the terminator field, which is what makes it safe to
+    record."""
+    path: Final = urlsplit(url).path
+    if path.endswith(BEDROCK_CONVERSE_STREAM_SUFFIX):
+        return complete_converse_stream(body)
+    if path.endswith(BEDROCK_INVOKE_STREAM_SUFFIX):
+        return complete_invoke_stream(body)
+    try:
+        value: Final = JSON_VALUE.validate_json(body)
+    except ValidationError:
+        return False
+    if not isinstance(value, dict) or "message" in value:
+        return False
+    if path.endswith(BEDROCK_CONVERSE_SUFFIX):
+        return isinstance(value.get("output"), dict) and isinstance(value.get("stopReason"), str)
+    return (
+        value.get("type") == "message"
+        and isinstance(value.get("content"), list)
+        and isinstance(value.get("stop_reason"), str)
+    )
+
+
+def whole_eventstream_messages(body: bytes) -> bool:
+    """Whether the body is exactly a whole number of eventstream messages.
+
+    A dropped connection is the failure this catches, and it has to be caught
+    here: botocore yields the messages it did receive and silently discards a
+    trailing partial one, so a stream cut a single byte short parses clean. Each
+    message declares its own total length in its first four bytes, so walking
+    those is enough to tell a complete body from a cut one."""
+    offset = 0  # rebind-ok: a cursor walking the declared frame lengths
+    while offset + EVENTSTREAM_PRELUDE_BYTES <= len(body):
+        total: int = int.from_bytes(body[offset : offset + EVENTSTREAM_PRELUDE_BYTES], "big")
+        if total <= 0 or offset + total > len(body):
+            return False
+        offset += total
+    return offset == len(body)
+
+
+def eventstream_events(body: bytes) -> tuple[tuple[str, JsonValue], ...] | None:
+    """The stream's (event type, decoded payload) pairs, or None if it is not a
+    complete, uncorrupted stream.
+
+    botocore validates both CRCs and raises ``ParserError`` rather than decoding
+    corruption into something plausible. A failure that began after Bedrock had
+    already answered 200 arrives as an ``exception`` frame in place of the
+    terminator, so it is the terminator rules below that reject it and this does
+    not need to inspect ``:message-type`` as well."""
+    if not body or not whole_eventstream_messages(body):
+        return None
+    buffer: Final = EventStreamBuffer()
+    buffer.add_data(body)
+    try:
+        return tuple(
+            (event_type(event.headers), JSON_VALUE.validate_json(event.payload))
+            for event in buffer
+        )
+    except (ParserError, ValidationError, ValueError):
+        return None
+
+
+def event_type(headers: object) -> str:
+    """botocore's eventstream headers come back untyped, so the one header this
+    reads is validated into a string rather than trusted."""
+    parsed: Final = EVENTSTREAM_HEADERS.validate_python(headers)
+    return parsed.get(EVENT_TYPE_HEADER, "")
+
+
+def complete_converse_stream(body: bytes) -> bool:
+    """ConverseStream ends with ``metadata``, not with ``messageStop``.
+
+    Requiring the metadata frame rather than the stop frame is deliberate: it
+    carries the token usage litellm prices the call from, so a stream cut between
+    the two still names a stop reason but would replay as a free call."""
+    events: Final = eventstream_events(body)
+    if not events or events[-1][0] != "metadata":
+        return False
+    return any(
+        event_type == "messageStop" and isinstance(payload, dict) and isinstance(payload.get("stopReason"), str)
+        for event_type, payload in events
+    )
+
+
+def complete_invoke_stream(body: bytes) -> bool:
+    """InvokeModelWithResponseStream wraps the ordinary Anthropic event grammar
+    in ``chunk`` frames, one base64 payload each, so it is held to the same
+    terminator rule as the Anthropic SSE path. A frame Bedrock sends instead of a
+    chunk, an exception among them, carries no such payload and fails the rule
+    without the frame type needing to be read."""
+    events: Final = eventstream_events(body)
+    if not events:
+        return False
+    values: Final = tuple(invoke_chunk_value(payload) for _, payload in events)
+    return all(value is not None for value in values) and complete_anthropic_stream(values)
+
+
+def invoke_chunk_value(payload: JsonValue) -> JsonValue | None:
+    """The Anthropic event inside one ``chunk`` frame, or None for a frame that
+    carries no readable one."""
+    if not isinstance(payload, dict) or not isinstance(encoded := payload.get("bytes"), str):
+        return None
+    try:
+        return JSON_VALUE.validate_json(base64.b64decode(encoded, validate=True))
+    except (ValidationError, ValueError):
+        return None
+
+
+def complete_anthropic_stream(values: tuple[JsonValue, ...]) -> bool:
+    """The Anthropic event grammar, shared by the SSE mounts and by Bedrock's
+    invoke stream, which carries the same events inside eventstream frames. A
+    ``message_delta`` naming a stop reason is what separates a finished turn from
+    one the connection cut short."""
+    if not values:
+        return False
+    first: Final = values[0]
+    last: Final = values[-1]
+    return (
+        isinstance(first, dict) and first.get("type") == "message_start"
+        and isinstance(last, dict) and last.get("type") == "message_stop"
+        and any(
+            isinstance(value, dict) and value.get("type") == "message_delta"
+            and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str)
+            for value in values
+        )
+    )
+
+
+def complete_responses_stream(values: tuple[JsonValue, ...]) -> bool:
+    """The Responses API streams typed events and ends with ``response.completed``.
+    A run that failed, was cancelled, or ran out of tokens ends with a different
+    terminal event, so requiring that one keeps a half-finished response out."""
+    last: Final = values[-1]
+    return isinstance(last, dict) and last.get("type") == "response.completed"
+
+
 def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool:
     if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values):
         return False
@@ -172,7 +384,7 @@ def encode_response(secret: bytes, response: CachedResponse) -> bytes:
     return SignedResponse(response=raw, signature=hmac.new(secret, raw.encode(), hashlib.sha256).hexdigest()).model_dump_json().encode()
 
 
-def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> CachedResponse | None:
+def decode_response(secret: bytes, key: str, payload: bytes, mount: str, url: str) -> CachedResponse | None:
     if len(payload) > 2 * MAX_RESPONSE_BYTES:
         return None
     try:
@@ -183,11 +395,58 @@ def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> Cached
         chunks: Final = tuple(base64.b64decode(chunk, validate=True) for chunk in response.chunks)
     except (ValidationError, ValueError):
         return None
-    if response.request_key != key or not successful_response(url, response.status_code, response.headers, b"".join(chunks)):
+    if response.request_key != key or not successful_response(
+        mount, url, response.status_code, response.headers, b"".join(chunks)
+    ):
         return None
     return response
 
 
+def component_digests(
+    test_key: str, method: str, url: str, headers: Mapping[str, str], body: bytes | None,
+) -> dict[str, str]:
+    """Per-component digests of everything the key covers.
+
+    A mount whose corpus never converges is a mount where one of these moves
+    between builds, and the flat key cannot say which. Values are digested, so
+    no payload or credential is written, and a JSON body contributes one digest
+    per top-level field so the field that moved can be named."""
+    parts: dict[str, str] = {  # rebind-ok: a report assembled from three differently shaped sources
+        "test_key": test_key,
+        "method": method,
+        "url": short_digest(canonical_text(url).encode()),
+    }
+    for name, value in sorted(headers.items()):
+        parts[f"header:{name.lower()}"] = short_digest(value.encode())
+    canonical: Final = b"" if body is None else canonical_body(body)
+    parts["body"] = short_digest(canonical)
+    try:
+        parsed: Final = JSON_VALUE.validate_json(canonical)
+    except ValidationError:
+        return parts
+    if isinstance(parsed, dict):
+        for name, value in sorted(parsed.items()):
+            parts[f"body:{name}"] = short_digest(json.dumps(value, sort_keys=True).encode())
+    return parts
+
+
+def short_digest(value: bytes) -> str:
+    return hashlib.sha256(value).hexdigest()[:16]
+
+
+@dataclass(slots=True)
+class KeyProbe:
+    """Every keyed request's components, when a metrics directory is configured."""
+
+    rows: tuple[tuple[tuple[str, str], ...], ...] = ()
+    lock: threading.Lock = field(default_factory=threading.Lock)
+
+    def observe(self, mount: str, outcome: str, parts: Mapping[str, str]) -> None:
+        row: Final = tuple({"mount": mount, "outcome": outcome, **parts}.items())
+        with self.lock:
+            self.rows = (*self.rows, row)
+
+
 @dataclass(slots=True)
 class CacheCounters:
     counts: tuple[tuple[str, int], ...] = ()
@@ -199,6 +458,24 @@ class CacheCounters:
             self.counts = tuple((current | {name: current.get(name, 0) + 1}).items())
 
 
+@dataclass(slots=True)
+class SlotCounter:
+    """FIFO position of a request among the canonically identical ones its test
+    has already sent. Two calls in one test that differ only by ``unique_marker``
+    canonicalize the same, so without this they would share one recording and the
+    second would replay the first's provider response id."""
+
+    counts: tuple[tuple[str, int], ...] = ()
+    lock: threading.Lock = field(default_factory=threading.Lock)
+
+    def take(self, identity: str) -> int:
+        with self.lock:
+            current: Final = dict(self.counts)
+            taken: Final = current.get(identity, 0)
+            self.counts = tuple((current | {identity: taken + 1}).items())
+            return taken
+
+
 @dataclass(slots=True)
 class ResponseCapture:
     buffer: io.BytesIO = field(default_factory=io.BytesIO)
@@ -226,14 +503,21 @@ def response_steps(response: CachedResponse) -> Generator[StreamStep, None, None
         yield StreamChunk(base64.b64decode(chunk, validate=True))
 
 
+NO_POLICIES: Final[Mapping[str, MountPolicy]] = MappingProxyType({})
+
+
 @dataclass(frozen=True, slots=True)
 class CacheEdge:
     store: ResponseStore
     secret: bytes = field(repr=False)
     counters: CacheCounters = field(default_factory=CacheCounters)
+    probe: KeyProbe = field(default_factory=KeyProbe)
+    slots: SlotCounter = field(default_factory=SlotCounter)
+    policies: Mapping[str, MountPolicy] = NO_POLICIES
     wait_seconds: float = 2.0
     clock: Callable[[], float] = time.monotonic
     sleep: Callable[[float], None] = time.sleep
+    test_key: Callable[[], str] = current_test_key
 
     def lookup(self, key: str) -> CacheLookup:
         deadline: Final = self.clock() + self.wait_seconds
@@ -241,59 +525,122 @@ class CacheEdge:
             self.sleep(min(0.05, max(0, deadline - self.clock())))
         return result
 
-    def forward(self, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float) -> StreamHead | NetworkError:
-        if not cacheable_endpoint(method, url, body):
-            self.counters.increment("bypass")
-            self.counters.increment("upstream_attempts")
-            return forward_stream(method, url, headers=headers, body=body, timeout=timeout)
-        prepared: Final = prepare_forward(method, url, headers, body)
+    def count(self, mount: str, name: str) -> None:
+        self.counters.increment(name)
+        self.counters.increment(f"mount:{mount}:{name}")
+
+    def record_key(
+        self, mount: str, outcome: str, test_key: str, method: str, url: str,
+        headers: Mapping[str, str], body: bytes | None,
+    ) -> None:
+        if not os.environ.get("E2E_PROVIDER_CACHE_METRICS_DIR"):
+            return
+        self.probe.observe(mount, outcome, component_digests(test_key, method, url, headers, body))
+
+    def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]:
+        """The headers actually sent upstream. A signing mount gets a signature
+        minted over the upstream URL, because the edge rewrote the Host the proxy
+        signed and the provider verifies it."""
+        signer: Final = self.policies.get(mount, MountPolicy()).sign
+        return headers if signer is None else signer(method, url, headers, body)
+
+    def keyed(self, mount: str, headers: Mapping[str, str]) -> Mapping[str, str]:
+        """Headers the cache key is built from. A mount keeps its credentials in
+        the key unless its policy names them unkeyed, so by default one account
+        can never read another's recording."""
+        unkeyed: Final = self.policies.get(mount, MountPolicy()).unkeyed_headers
+        if not unkeyed:
+            return headers
+        return {name: value for name, value in headers.items() if name.lower() not in unkeyed}
+
+    def forward(
+        self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float,
+    ) -> StreamHead | NetworkError:
+        test_key: Final = self.test_key()
+        if test_key == SESSION_TEST_KEY or not cacheable_endpoint(mount, method, url, body):
+            self.count(mount, "bypass")
+            self.count(mount, "upstream_attempts")
+            return forward_stream(
+                method, url, headers=self.outbound(mount, method, url, headers, body), body=body, timeout=timeout,
+            )
+        prepared: Final = prepare_forward(method, url, self.outbound(mount, method, url, headers, body), body)
         if isinstance(prepared, NetworkError):
-            self.counters.increment("rejected")
+            self.reject(mount, UNREACHABLE)
             return prepared
-        key: Final = exact_key(self.secret, method, url, prepared.headers, body)
+        keyed_headers: Final = self.keyed(mount, prepared.headers)
+        identity: Final = request_identity(self.secret, test_key, method, url, keyed_headers, body)
+        key: Final = slotted_key(self.secret, identity, self.slots.take(identity))
         found: Final = self.lookup(key)
         if isinstance(found, CacheHit):
-            response: Final = decode_response(self.secret, key, found.payload, url)
+            response: Final = decode_response(self.secret, key, found.payload, mount, url)
             if response is not None and self.clock() < found.valid_until:
-                self.counters.increment("hits")
+                self.count(mount, "hits")
+                self.record_key(mount, "hit", test_key, method, url, keyed_headers, body)
                 return StreamHead(response.status_code, response.headers, response_steps(response))
-            self.counters.increment("corrupt" if response is None else "expired")
+            self.count(mount, "corrupt" if response is None else "expired")
             self.store.discard(key, found.payload)
         capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found
-        self.counters.increment("misses")
+        self.count(mount, "misses")
+        self.record_key(mount, "miss", test_key, method, url, keyed_headers, body)
         if isinstance(capture_slot, CacheUnavailable):
-            self.counters.increment("cache_errors")
-        self.counters.increment("upstream_attempts")
+            self.count(mount, "cache_errors")
+        self.count(mount, "upstream_attempts")
         head: Final = forward_prepared_stream(prepared, timeout)
         if not isinstance(capture_slot, CaptureLease):
             return head
         if isinstance(head, NetworkError):
             self.store.release(key, capture_slot)
-            self.counters.increment("rejected")
+            self.reject(mount, UNREACHABLE)
             return head
-        return StreamHead(head.status_code, head.headers, primed_steps(self.capture(key, capture_slot, url, head)))
+        return StreamHead(
+            head.status_code, head.headers, primed_steps(self.capture(mount, key, capture_slot, url, head)),
+        )
 
-    def capture(self, key: str, lease: CaptureLease, url: str, head: StreamHead) -> Generator[StreamStep, None, None]:
+    def capture(
+        self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead,
+    ) -> Generator[StreamStep, None, None]:
         capture: Final = ResponseCapture()
+        reason = CUT_SHORT  # rebind-ok: a consumer that walks away never reaches the settle call below
         try:
             with closing(head.steps):
                 yield StreamChunk(b"")
                 for step in head.steps:
                     yield step
                     capture.observe(step)
-            chunks: Final = capture.chunks() if capture.eligible else ()
-            headers: Final = {
-                name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS
-            }
-            if not capture.eligible or not successful_response(url, head.status_code, headers, b"".join(chunks)):
-                self.counters.increment("rejected")
-                return
-            response: Final = CachedResponse(
-                request_key=key, status_code=head.status_code, headers=headers,
-                chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks),
-            )
-            published: Final = self.store.publish(key, lease, encode_response(self.secret, response))
-            self.counters.increment("writes" if published else "write_failures")
+            reason = self.settle(mount, key, lease, url, head, capture)
         finally:
+            self.reject(mount, reason)
             self.store.release(key, lease)
             capture.buffer.close()
+
+    def settle(
+        self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, capture: ResponseCapture,
+    ) -> str | None:
+        """None once the response is stored, otherwise the reason it was not."""
+        if not capture.eligible:
+            return CUT_SHORT
+        headers: Final = {
+            name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS
+        }
+        if not 200 <= head.status_code < 300:
+            return ERROR_STATUS
+        chunks: Final = capture.chunks()
+        if not successful_response(mount, url, head.status_code, headers, b"".join(chunks)):
+            return INCOMPLETE
+        response: Final = CachedResponse(
+            request_key=key, status_code=head.status_code, headers=headers,
+            chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks),
+        )
+        published: Final = self.store.publish(key, lease, encode_response(self.secret, response))
+        self.count(mount, "writes" if published else "write_failures")
+        return None
+
+    def reject(self, mount: str, reason: str | None) -> None:
+        """A flat rejection count cannot separate a connection that went away from
+        a body the provider finished sending and the rules turned down, and the two
+        have opposite fixes. A mount whose rejections are nearly all one or the
+        other is a different problem, so the report has to be able to say which."""
+        if reason is None:
+            return
+        self.count(mount, "rejected")
+        self.count(mount, f"rejected_{reason}")

tests/e2e/provider_cache_redis.py

@@ -134,6 +134,10 @@ def write_metrics(cache: CacheEdge) -> None:
             root: Final = Path(directory)
             root.mkdir(parents=True, exist_ok=True)
             (root / f"{os.getpid()}.json").write_text(report + "\n")
+            if cache.probe.rows:
+                (root / f"keys-{os.getpid()}.json").write_text(
+                    json.dumps([dict(row) for row in cache.probe.rows]) + "\n"
+                )
         except OSError:
             logging.getLogger(__name__).warning("provider cache metrics artifact unavailable")
     logging.getLogger(__name__).info("%s", report)

tests/e2e/provider_cache_routing.py

@@ -8,14 +8,80 @@ from models import LiteLLMParamsBody, ModelMode
 
 LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False)
 
+DEFAULT_BEDROCK_REGION: Final = "us-east-1"
+BEDROCK_CROSS_REGION_PREFIX: Final = "us."
+BEDROCK_EDGE_MODELS: Final = frozenset(
+    {
+        "us.anthropic.claude-haiku-4-5-20251001-v1:0",
+        "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
+        "us.anthropic.claude-sonnet-5",
+        "us.anthropic.claude-opus-4-7",
+    }
+)
+ENV_REFERENCE_PREFIX: Final = "os.environ/"
+
+
+def bedrock_region(declared: str | None) -> str:
+    """The region whose edge mount a deployment belongs to.
+
+    Most Bedrock deployments declare `os.environ/AWS_REGION`, which only the
+    proxy can resolve from its own environment; the run pod does not share it.
+    Answering those with the default mount is correct because every model on the
+    edge allowlist is a `us.` inference profile, which fans out across the US
+    regions and is reachable from any of them. That invariant is enforced on the
+    allowlist itself rather than re-checked per call."""
+    if declared is None or declared.startswith(ENV_REFERENCE_PREFIX):
+        return DEFAULT_BEDROCK_REGION
+    return declared
+
+
+def bedrock_mount(params: LiteLLMParamsBody) -> str | None:
+    """The edge mount a Bedrock deployment belongs to, or None.
+
+    The allowlist mirrors the runner role's IAM policy, which names its models
+    one by one. A model outside it would be re-signed with an identity that
+    cannot invoke it and come back 403 from Bedrock, so an unlisted model keeps
+    its direct path and loses only caching. Adding a model is a policy edit in
+    litellm-ops and a line here."""
+    route: Final = params.model.partition("/")[2]
+    model: Final = route.partition("/")[2] or route
+    if model not in BEDROCK_EDGE_MODELS:
+        return None
+    return f"bedrock/{bedrock_region(params.aws_region_name)}"
+
+
+def route_bedrock(
+    params: LiteLLMParamsBody, base_for: Callable[[str], str | None], mode: ModelMode | None,
+) -> LiteLLMParamsBody:
+    """Deployments that carry their own AWS identity stay off the edge. The edge
+    re-signs with the run pod's role, so routing an `aws_role_name` deployment
+    would quietly replace the very assume-role chain that test exists to prove."""
+    if mode is not None or params.aws_role_name is not None or params.aws_access_key_id is not None:
+        return params
+    if params.api_base is not None or params.aws_bedrock_runtime_endpoint is not None:
+        return params
+    mount: Final = bedrock_mount(params)
+    if mount is None:
+        return params
+    base: Final = base_for(mount)
+    if base is None:
+        return params
+    return params.model_copy(update={"aws_bedrock_runtime_endpoint": base})
+
 
 def route_cache_model(
     params: LiteLLMParamsBody, base_for: Callable[[str], str | None], *, enabled: bool, mode: ModelMode | None = None,
 ) -> LiteLLMParamsBody:
-    if not enabled or mode == "realtime" or LIVE_PROVIDER_REQUIRED.get() or params.api_base is not None or params.mock_response is not None:
+    if not enabled or LIVE_PROVIDER_REQUIRED.get() or params.mock_response is not None:
+        return params
+    if params.litellm_credential_name is not None:
         return params
     provider: Final = params.model.partition("/")[0]
-    if provider not in {"openai", "anthropic"} or params.litellm_credential_name is not None:
+    if provider == "bedrock":
+        return route_bedrock(params, base_for, mode)
+    if mode == "realtime" or params.api_base is not None:
+        return params
+    if provider not in {"openai", "anthropic"}:
         return params
     base: Final = base_for(provider)
     if base is None:

tests/e2e/provider_edge.py

@@ -48,7 +48,7 @@ import threading
 from collections import deque
 from collections.abc import Generator, Mapping, Sequence
 from contextlib import closing, contextmanager
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from itertools import islice
 from pathlib import Path
@@ -94,17 +94,41 @@ from fixture_mode import (
     parse_fixture_mode,
 )
 from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity
-from provider_cache import CacheEdge
+from provider_cache import SIGNATURE_HEADERS, CacheEdge, MountPolicy, is_bedrock
 from provider_cache_routing import LIVE_PROVIDER_REQUIRED
 from pydantic import JsonValue, TypeAdapter
 
+BEDROCK_REGIONS: Final[tuple[str, ...]] = ("us-east-1",)
+
 EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType(
     {
         "openai": "https://api.openai.com",
         "anthropic": "https://api.anthropic.com",
+        **{
+            f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com"
+            for region in BEDROCK_REGIONS
+        },
     }
 )
 
+
+@dataclass(frozen=True, slots=True)
+class ResolvedMount:
+    mount: str
+    upstream_base: str
+    upstream_path: str
+
+
+def resolve_mount(path: str, mounts: Mapping[str, str]) -> ResolvedMount | None:
+    """Longest mount prefix wins, so a region-qualified mount such as
+    ``bedrock/us-east-1`` resolves whole instead of leaving the region as the
+    first segment of the upstream path."""
+    trimmed: Final = path.lstrip("/")
+    for mount in sorted(mounts, key=len, reverse=True):
+        if trimmed == mount or trimmed.startswith(f"{mount}/"):
+            return ResolvedMount(mount, mounts[mount], trimmed[len(mount):].lstrip("/"))
+    return None
+
 REPLAY_MISS_STATUS: Final = 599
 
 _HOP_BY_HOP_HEADERS: Final[frozenset[str]] = frozenset(
@@ -754,14 +778,14 @@ def _handle_record(
 
 def _handle_live(
     method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float,
-    cache: CacheEdge | None = None,
+    cache: CacheEdge | None = None, mount: str = "",
 ) -> EdgeOutcome:
     forwarded: Final = {
         name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS
     }
     head: Final = (
         forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
-        if cache is None else cache.forward(method, url, forwarded, body, timeout)
+        if cache is None else cache.forward(mount, method, url, forwarded, body, timeout)
     )
     match head:
         case NetworkError(message=message):
@@ -796,10 +820,13 @@ def handle_edge_request(
     prefix, then record (forward + persist) or replay (serve from the bundle).
     Socket-free so unit tests exercise every branch without a server."""
     split: Final = urlsplit(raw_path)
-    mount, _, upstream_path = split.path.lstrip("/").partition("/")
-    upstream_base: Final = mounts.get(mount)
-    if upstream_base is None:
-        return _text_reply(404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}")
+    resolved: Final = resolve_mount(split.path, mounts)
+    if resolved is None:
+        unknown: Final = split.path.lstrip("/").partition("/")[0]
+        return _text_reply(404, f"unknown provider mount {unknown!r}; known mounts: {', '.join(sorted(mounts))}")
+    mount: Final = resolved.mount
+    upstream_base: Final = resolved.upstream_base
+    upstream_path: Final = resolved.upstream_path
     profile: Final = (
         backend.recorder.profile
         if isinstance(backend, RecordEdge)
@@ -830,7 +857,8 @@ def handle_edge_request(
     match backend:
         case CacheEdge():
             return _handle_live(
-                method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend,
+                method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout,
+                backend, mount,
             )
         case LiveEdge():
             return _handle_live(
@@ -891,7 +919,7 @@ class _EdgeHandler(BaseHTTPRequestHandler):
         )
         if isinstance(edge_server.backend, CacheEdge) and duplicate_headers:
             edge_server.backend.counters.increment("duplicate_header_bypass")
-            if urlsplit(self.path).path.lstrip("/").partition("/")[0] in edge_server.mounts:
+            if resolve_mount(urlsplit(self.path).path, edge_server.mounts) is not None:
                 edge_server.backend.counters.increment("upstream_attempts")
         outcome: Final = handle_edge_request(
             selected_backend,
@@ -1079,6 +1107,8 @@ def provider_edge_api_base(
                 return _shared_cache_edge(bind_host, advertise_host, forward_timeout).api_base(mount)
             return None
         case "record" | "replay":
+            if is_bedrock(mount):
+                return None
             if mount not in EDGE_MOUNTS:
                 raise ValueError(f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(EDGE_MOUNTS))}")
             return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout, match_profile()).api_base(
@@ -1108,7 +1138,22 @@ def configured_cache_backend() -> CacheEdge | None:
         return None
     from provider_cache_redis import configured_cache
 
-    return configured_cache()
+    cache: Final = configured_cache()
+    return None if cache is None else replace(cache, policies=bedrock_policies())
+
+
+@functools.lru_cache(maxsize=1)
+def bedrock_policies() -> Mapping[str, MountPolicy]:
+    """One policy per mounted Bedrock region, built lazily so a run that never
+    mounts Bedrock neither imports botocore nor resolves an AWS identity."""
+    from provider_edge_bedrock import bedrock_signer
+
+    return MappingProxyType(
+        {
+            f"bedrock/{region}": MountPolicy(sign=bedrock_signer(region), unkeyed_headers=SIGNATURE_HEADERS)
+            for region in BEDROCK_REGIONS
+        }
+    )
 
 
 @functools.lru_cache(maxsize=8)

tests/e2e/provider_edge_bedrock.py

@@ -0,0 +1,72 @@
+"""SigV4 re-signing for Bedrock traffic routed through the provider edge.
+
+Bedrock is the one provider the edge could never mount. SigV4 signs the Host
+header, so rewriting ``api_base`` to point at the edge invalidates the proxy's
+signature and Bedrock rejects the call before it reaches a model. The edge
+therefore has to drop the proxy's signature and mint its own over the upstream
+URL it is actually about to call.
+
+The identity it signs with is the run pod's own, from the EKS Pod Identity
+association on ServiceAccount ``buildkite-e2e-run``. That role carries Bedrock
+invoke and converse on an allowlist of the Anthropic models the suite registers
+and nothing else, so a re-signed call can reach exactly the models the suite
+already uses. The proxy's own Bedrock credentials are not involved in a routed
+deployment, which is why ``aws_role_name`` deployments stay off the edge: their
+whole point is to prove the product's assume-role chain.
+
+Signature headers are excluded from the cache key by the caller, and they have
+to be: ``x-amz-date`` is a timestamp, so keying on it would make every Bedrock
+request a permanent miss.
+"""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Final
+
+from botocore.auth import SigV4Auth
+from botocore.awsrequest import AWSRequest
+from botocore.credentials import Credentials
+from botocore.session import Session
+from provider_cache import SIGNATURE_HEADERS
+
+BEDROCK_SERVICE: Final = "bedrock"
+
+
+class MissingAwsCredentials(RuntimeError):
+    """No AWS identity is resolvable, so the edge cannot sign for Bedrock."""
+
+
+@dataclass(frozen=True, slots=True)
+class BedrockSigner:
+    region: str
+    credentials: Callable[[], Credentials]
+
+    def __call__(self, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> dict[str, str]:
+        unsigned: Final = {
+            name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS
+        }
+        request: Final = AWSRequest(method=method, url=url, headers=unsigned, data=body or b"")
+        SigV4Auth(self.credentials(), BEDROCK_SERVICE, self.region).add_auth(request)
+        return dict(request.headers)
+
+
+@functools.lru_cache(maxsize=1)
+def pod_credentials() -> Credentials:
+    """The run pod's own identity, resolved once per process through botocore's
+    ordinary chain, which reaches Pod Identity at the ``container-role`` link."""
+    resolved: Final = Session().get_credentials()
+    if resolved is None:  # pyright: ignore[reportUnnecessaryComparison]  # stubs miss the empty-chain None
+        raise MissingAwsCredentials(
+            "the provider edge is mounted for Bedrock but no AWS credentials resolve; "
+            "the run pod gets them from the Pod Identity association on buildkite-e2e-run"
+        )
+    return resolved
+
+
+def bedrock_signer(region: str, credentials: Callable[[], Credentials] = pod_credentials) -> BedrockSigner:
+    """Credentials are resolved on the first signed request, not here, so a run
+    that mounts Bedrock but never calls it needs no AWS identity at all."""
+    return BedrockSigner(region, credentials)

tests/e2e/pytest.ini

@@ -10,5 +10,6 @@ markers =
     weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set
     managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set
     prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set
+    cli_determinism: drives the real claude CLI for several seconds, which widens the window in which another test's in-flight upstream call is attributed to it; deselected unless E2E_CLI_DETERMINISM is set
     redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set
     cost_map_stack: needs a proxy whose whole cost map is tests/e2e/cost_map.json (LITELLM_MODEL_COST_MAP_URL) plus a scripted-provider sidecar; deselected unless E2E_COST_MAP_STACK is set

tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md

@@ -21,7 +21,7 @@ on the shared lifecycle (every entity it creates is deleted on teardown).
 
 | Entity | Unit | Pre-existing live | This suite (live) | Status |
 |--------|------|-------------------|-------------------|--------|
-| API key | `test_budget_reservation.py`, `test_max_budget_limiter.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** |
+| API key | `test_budget_reservation.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** |
 | Team | `test_team_budget_limits.py` | `otel_tests` | (org test builds a team) | **covered** |
 | Internal user | auth unit tests | - | `test_internal_user_budget_blocks` | **covered (new)** |
 | Team member | `test_team_member_budget.py` | - | `test_team_member_budget_blocks` | **covered (new)** |

tests/e2e/test_provider_edge.py

@@ -1279,15 +1279,30 @@ class TestApiBaseSeam:
             )
 
     def test_unknown_mount_raises_naming_the_known_mounts(self, tmp_path: Path) -> None:
-        with pytest.raises(ValueError, match="unknown provider mount 'bedrock'"):
+        with pytest.raises(ValueError, match="unknown provider mount 'cohere'"):
             provider_edge_api_base(
-                "bedrock",
+                "cohere",
                 mode_raw="record",
                 bundle_dir=tmp_path / "bundle",
                 bind_host="127.0.0.1",
                 advertise_host="127.0.0.1",
             )
 
+    @pytest.mark.parametrize("mode_raw", ["record", "replay"])
+    def test_bedrock_never_wires_a_bundle_because_the_edge_cannot_sign_into_one(
+        self, tmp_path: Path, mode_raw: str,
+    ) -> None:
+        """Record and replay serve from a bundle without re-signing, so a Bedrock
+        deployment pointed at that edge would send the proxy's signature over a
+        rewritten Host. It keeps its direct route in both modes."""
+        assert provider_edge_api_base(
+            "bedrock/us-east-1",
+            mode_raw=mode_raw,
+            bundle_dir=tmp_path / "bundle",
+            bind_host="127.0.0.1",
+            advertise_host="127.0.0.1",
+        ) is None
+
     def test_record_mode_boots_one_shared_edge_and_prepares_the_bundle(self, tmp_path: Path) -> None:
         root = tmp_path / "bundle"
         first = provider_edge_api_base(

tests/llm_translation/test_together_ai.py

@@ -3,6 +3,7 @@ Test TogetherAI LLM
 """
 
 from base_llm_unit_tests import BaseLLMChatTest
+from tests._live_test_helpers import cheapest_together_chat_model
 import json
 import os
 from datetime import datetime
@@ -16,7 +17,11 @@ import pytest
 class TestTogetherAI(BaseLLMChatTest):
     def get_base_completion_call_args(self) -> dict:
         litellm.set_verbose = True
-        return {"model": "together_ai/openai/gpt-oss-20b"}
+        return {
+            "model": cheapest_together_chat_model(
+                function_calling=True, response_schema=True
+            )
+        }
 
     def test_tool_call_no_arguments(self, tool_call_no_arguments):
         """Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833"""

tests/local_testing/test_completion.py

@@ -57,23 +57,6 @@ def test_response_model_none():
     assert isinstance(x, litellm.ModelResponse)
 
 
-def test_completion_custom_provider_model_name():
-    try:
-        litellm.cache = None
-        response = completion(
-            model="together_ai/openai/gpt-oss-20b",
-            messages=messages,
-            logger_fn=logger_fn,
-        )
-        # Add assertions here to check the-response
-        print(response)
-        print(response["choices"][0]["finish_reason"])
-    except litellm.Timeout as e:
-        pass
-    except Exception as e:
-        pytest.fail(f"Error occurred: {e}")
-
-
 def _openai_mock_response(*args, **kwargs) -> litellm.ModelResponse:
     new_response = MagicMock()
     new_response.headers = {"hello": "world"}
@@ -2803,41 +2786,6 @@ def test_completion_together_ai_llama():
 
 
 # test_completion_together_ai()
-def test_customprompt_together_ai():
-    try:
-        litellm.set_verbose = False
-        litellm.num_retries = 0
-        print("in test_customprompt_together_ai")
-        print(litellm.success_callback)
-        print(litellm._async_success_callback)
-        response = completion(
-            model="together_ai/openai/gpt-oss-20b",
-            messages=messages,
-            roles={
-                "system": {
-                    "pre_message": "<|im_start|>system\n",
-                    "post_message": "<|im_end|>",
-                },
-                "assistant": {
-                    "pre_message": "<|im_start|>assistant\n",
-                    "post_message": "<|im_end|>",
-                },
-                "user": {
-                    "pre_message": "<|im_start|>user\n",
-                    "post_message": "<|im_end|>",
-                },
-            },
-        )
-        print(response)
-    except litellm.exceptions.Timeout as e:
-        print(f"Timeout Error")
-        pass
-    except Exception as e:
-        print(f"ERROR TYPE {type(e)}")
-        pytest.fail(f"Error occurred: {e}")
-
-
-# test_customprompt_together_ai()
 
 
 def response_format_tests(response: litellm.ModelResponse):
@@ -3644,28 +3592,6 @@ async def test_acompletion_stream_watsonx():
 # test_maritalk()
 
 
-def test_completion_together_ai_stream():
-    litellm.set_verbose = True
-    user_message = "Write 1pg about YC & litellm"
-    messages = [{"content": user_message, "role": "user"}]
-    try:
-        response = completion(
-            model="together_ai/openai/gpt-oss-20b",
-            messages=messages,
-            stream=True,
-            max_tokens=5,
-        )
-        print(response)
-        for chunk in response:
-            print(chunk)
-        # print(string_response)
-    except Exception as e:
-        pytest.fail(f"Error occurred: {e}")
-
-
-# test_completion_together_ai_stream()
-
-
 def test_moderation():
     response = litellm.moderation(input="i'm ishaan cto of litellm")
     print(response)

tests/local_testing/test_text_completion.py

@@ -12,6 +12,7 @@ from unittest.mock import MagicMock, patch
 import pytest
 
 import litellm
+from tests._live_test_helpers import cheapest_together_chat_model
 from litellm import (
     RateLimitError,
     TextCompletionResponse,
@@ -4030,7 +4031,7 @@ def test_async_text_completion_together_ai():
     async def test_get_response():
         try:
             response = await litellm.atext_completion(
-                model="together_ai/openai/gpt-oss-20b",
+                model=cheapest_together_chat_model(),
                 prompt="good morning",
                 max_tokens=10,
             )

tests/logging_callback_tests/test_generic_api_callback.py

@@ -10,6 +10,7 @@ import httpx
 import json
 import logging
 import time
+from typing import Final
 from unittest.mock import AsyncMock, patch
 
 import pytest
@@ -98,8 +99,15 @@ async def test_generic_api_callback():
     assert isinstance(actual_request, list), "Request body should be a list"
     assert len(actual_request) > 0, "Request body list should not be empty"
 
-    # Validate the first payload item
-    payload_item: StandardLoggingPayload = StandardLoggingPayload(**actual_request[0])
+    this_test_messages: Final = [{"role": "user", "content": "Hello, world!"}]
+    mine: Final = [
+        item for item in actual_request if item.get("messages") == this_test_messages
+    ]
+    assert (
+        len(mine) == 1
+    ), f"Expected this test's single call in the batch, got {len(mine)} of {len(actual_request)}"
+
+    payload_item: StandardLoggingPayload = StandardLoggingPayload(**mine[0])
     print("##########\n")
     print(json.dumps(payload_item, indent=4))
     print("##########\n")
@@ -448,11 +456,17 @@ async def test_generic_api_callback_sumologic_uses_ndjson():
     assert isinstance(ndjson_data, str), "Data should be a string for NDJSON"
 
     lines = ndjson_data.strip().split("\n")
-    assert len(lines) == 2, f"Expected 2 lines of NDJSON, got {len(lines)}"
+    records: Final = [json.loads(line) for line in lines]
 
-    # Each line should be valid JSON
-    for line in lines:
-        json.loads(line)  # Will raise if invalid JSON
+    this_test_messages: Final = [
+        [{"role": "user", "content": f"Test {i}"}] for i in range(2)
+    ]
+    mine: Final = [
+        record for record in records if record.get("messages") in this_test_messages
+    ]
+    assert (
+        len(mine) == 2
+    ), f"Expected this test's 2 calls as NDJSON lines, got {len(mine)} of {len(records)}"
 
 
 @pytest.mark.asyncio

tests/logging_callback_tests/test_unit_test_litellm_logging.py

@@ -8,8 +8,8 @@ from typing import Literal
 import pytest
 import litellm
 from litellm.litellm_core_utils.litellm_logging import Logging
-from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
 from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
+from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler
 from litellm._service_logger import ServiceLogging
 import asyncio
 
@@ -58,11 +58,11 @@ def test_is_internal_litellm_proxy_callback():
     """
     Ensure we can determine if a callback is an internal litellm proxy callback
 
-    eg. `_PROXY_MaxBudgetLimiter`, `_PROXY_CacheControlCheck`
+    eg. `_PROXY_MaxIterationsHandler`, `_PROXY_CacheControlCheck`
     """
     logging = setup_logging()
 
-    assert logging._is_internal_litellm_proxy_callback(_PROXY_MaxBudgetLimiter) == True
+    assert logging._is_internal_litellm_proxy_callback(_PROXY_MaxIterationsHandler) == True
 
     # Test non-internal callbacks
     def regular_callback():
@@ -95,7 +95,7 @@ def test_should_run_sync_callbacks_for_async_calls():
     assert logging._should_run_sync_callbacks_for_async_calls() == True
 
     # Test with internal callback only
-    litellm.success_callback = [_PROXY_MaxBudgetLimiter]
+    litellm.success_callback = [_PROXY_MaxIterationsHandler]
     assert logging._should_run_sync_callbacks_for_async_calls() == False
 
 
@@ -107,7 +107,7 @@ def test_remove_internal_litellm_callbacks():
 
     callbacks = [
         regular_callback,
-        _PROXY_MaxBudgetLimiter,
+        _PROXY_MaxIterationsHandler,
         _PROXY_CacheControlCheck,
         "string_callback",
     ]
@@ -116,5 +116,5 @@ def test_remove_internal_litellm_callbacks():
     assert len(filtered) == 2  # Should only keep regular_callback and string_callback
     assert regular_callback in filtered
     assert "string_callback" in filtered
-    assert _PROXY_MaxBudgetLimiter not in filtered
+    assert _PROXY_MaxIterationsHandler not in filtered
     assert _PROXY_CacheControlCheck not in filtered

tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py

@@ -587,6 +587,8 @@ def _success_kwargs(
     response_cost=0.5,
     key_hash=None,
     key_model_max_budget=None,
+    team_id=None,
+    team_model_max_budget=None,
     user_id=None,
     user_model_max_budget=None,
     end_user_id=None,
@@ -600,6 +602,7 @@ def _success_kwargs(
             "end_user": end_user_id,
             "metadata": {
                 "user_api_key_hash": key_hash,
+                "user_api_key_team_id": team_id,
                 "user_api_key_user_id": user_id,
                 "user_api_key_end_user_id": end_user_id,
             },
@@ -607,6 +610,7 @@ def _success_kwargs(
         "litellm_params": {
             "metadata": {
                 "user_api_key_model_max_budget": key_model_max_budget,
+                "user_api_key_team_model_max_budget": team_model_max_budget,
                 "user_api_key_user_model_max_budget": user_model_max_budget,
                 "user_api_key_end_user_model_max_budget": end_user_model_max_budget,
             },
@@ -1417,3 +1421,266 @@ async def test_spend_logged_on_one_replica_is_enforced_and_reported_on_another()
     replica_c = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis))
     with pytest.raises(litellm.BudgetExceededError):
         await replica_c.is_key_within_model_budget(user_api_key, "gpt-4")
+
+
+def _log_success(limiter, **kwargs):
+    return limiter.async_log_success_event(
+        _success_kwargs(**kwargs), response_obj=None, start_time=None, end_time=None
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "request_model",
+    ["gpt-4", "openai/gpt-4"],
+    ids=["bare_model", "provider_prefixed_model"],
+)
+async def test_team_model_budget_is_shared_by_every_key_without_an_override(request_model):
+    """
+    Two keys on the same team, neither carrying a matching key-level entry,
+    charge one team counter and are both refused once it is spent.
+    """
+    dual_cache = DualCache()
+    limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+    team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}}
+    check = lambda: limiter.is_team_within_model_budget(
+        team_id="team-1",
+        team_model_max_budget=team_model_max_budget,
+        key_model_max_budget=None,
+        model=request_model,
+    )
+
+    assert await check() is True
+    await _log_success(
+        limiter,
+        model_group=request_model,
+        response_cost=0.6,
+        key_hash="vk-a",
+        team_id="team-1",
+        team_model_max_budget=team_model_max_budget,
+    )
+    assert await check() is True
+    await _log_success(
+        limiter,
+        model_group=request_model,
+        response_cost=0.6,
+        key_hash="vk-b",
+        team_id="team-1",
+        team_model_max_budget=team_model_max_budget,
+    )
+
+    assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == pytest.approx(1.2)
+    with pytest.raises(litellm.BudgetExceededError) as exc:
+        await check()
+    assert exc.value.entity_type == Litellm_EntityType.TEAM.value
+    assert await build_model_max_budget_usage(
+        entity_type=Litellm_EntityType.TEAM,
+        entity_id="team-1",
+        model_max_budget=team_model_max_budget,
+        cache=dual_cache,
+    ) == {"gpt-4": {"current_spend": pytest.approx(1.2), "budget_limit": 1.0, "time_period": "1d"}}
+
+
+@pytest.mark.asyncio
+async def test_key_override_replaces_the_team_cap_for_that_model():
+    """
+    A key with its own entry for the model is gated on the key counter alone:
+    the exhausted team counter does not block it, and its spend never lands on
+    the team counter.
+    """
+    dual_cache = DualCache()
+    limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+    team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}}
+    key_model_max_budget = {"gpt-4": {"budget_limit": 5.0, "time_period": "1d"}}
+    await dual_cache.async_set_cache(key="team_model_spend:team-1:gpt-4:1d", value=9.0)
+
+    assert (
+        await limiter.is_team_within_model_budget(
+            team_id="team-1",
+            team_model_max_budget=team_model_max_budget,
+            key_model_max_budget=key_model_max_budget,
+            model="openai/gpt-4",
+        )
+        is True
+    )
+
+    await _log_success(
+        limiter,
+        model_group="openai/gpt-4",
+        response_cost=2.0,
+        key_hash="vk-override",
+        key_model_max_budget=key_model_max_budget,
+        team_id="team-1",
+        team_model_max_budget=team_model_max_budget,
+    )
+
+    assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 9.0
+    assert await dual_cache.async_get_cache(key="virtual_key_spend:vk-override:gpt-4:1d") == 2.0
+
+
+@pytest.mark.asyncio
+async def test_key_entry_for_another_model_does_not_lift_the_team_cap():
+    """A key override only covers the model it names; other models stay on the team counter."""
+    dual_cache = DualCache()
+    limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+    team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}}
+    key_model_max_budget = {"claude-3": {"budget_limit": 5.0, "time_period": "1d"}}
+
+    await _log_success(
+        limiter,
+        model_group="gpt-4",
+        response_cost=1.5,
+        key_hash="vk-other",
+        key_model_max_budget=key_model_max_budget,
+        team_id="team-1",
+        team_model_max_budget=team_model_max_budget,
+    )
+
+    assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5
+    with pytest.raises(litellm.BudgetExceededError):
+        await limiter.is_team_within_model_budget(
+            team_id="team-1",
+            team_model_max_budget=team_model_max_budget,
+            key_model_max_budget=key_model_max_budget,
+            model="gpt-4",
+        )
+
+
+@pytest.mark.asyncio
+async def test_team_budget_leaves_unconfigured_models_alone():
+    dual_cache = DualCache()
+    limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+    team_model_max_budget = {"gpt-4": {"budget_limit": 0.0, "time_period": "1d"}}
+
+    assert (
+        await limiter.is_team_within_model_budget(
+            team_id="team-1",
+            team_model_max_budget=team_model_max_budget,
+            key_model_max_budget=None,
+            model="claude-3",
+        )
+        is True
+    )
+    with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment:
+        await _log_success(
+            limiter,
+            model_group="claude-3",
+            team_id="team-1",
+            team_model_max_budget=team_model_max_budget,
+        )
+    mock_increment.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_team_counters_are_isolated_by_team_model_and_window():
+    """Same model on two teams, and two models with different windows on one team, never share a counter."""
+    dual_cache = DualCache()
+    limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+    team_model_max_budget = {
+        "gpt-4": {"budget_limit": 10.0, "time_period": "1d"},
+        "claude-3": {"budget_limit": 10.0, "time_period": "30d"},
+    }
+
+    for team_id, model in (("team-1", "gpt-4"), ("team-2", "gpt-4"), ("team-1", "claude-3")):
+        await _log_success(
+            limiter,
+            model_group=model,
+            response_cost=1.0,
+            team_id=team_id,
+            team_model_max_budget=team_model_max_budget,
+        )
+
+    assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.0
+    assert await dual_cache.async_get_cache(key="team_model_spend:team-2:gpt-4:1d") == 1.0
+    assert await dual_cache.async_get_cache(key="team_model_spend:team-1:claude-3:30d") == 1.0
+    assert await dual_cache.async_get_cache(key="team_model_budget_start_time:team-1:claude-3:30d") is not None
+
+
+@pytest.mark.asyncio
+async def test_malformed_team_entry_is_skipped_and_its_sibling_still_enforced():
+    limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
+    team_model_max_budget = {
+        "gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"},
+        "claude-3": {"budget_limit": 0.0, "time_period": "1d"},
+    }
+
+    assert (
+        await limiter.is_team_within_model_budget(
+            team_id="team-1",
+            team_model_max_budget=team_model_max_budget,
+            key_model_max_budget=None,
+            model="gpt-4",
+        )
+        is True
+    )
+    with pytest.raises(litellm.BudgetExceededError):
+        await limiter.is_team_within_model_budget(
+            team_id="team-1",
+            team_model_max_budget=team_model_max_budget,
+            key_model_max_budget=None,
+            model="claude-3",
+        )
+
+
+@pytest.mark.asyncio
+async def test_malformed_key_entry_does_not_count_as_an_override():
+    """A key entry the limiter cannot enforce must not also switch the team cap off."""
+    dual_cache = DualCache()
+    limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+    team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}}
+    key_model_max_budget = {"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}}
+
+    await _log_success(
+        limiter,
+        model_group="gpt-4",
+        response_cost=1.5,
+        key_hash="vk-bad",
+        key_model_max_budget=key_model_max_budget,
+        team_id="team-1",
+        team_model_max_budget=team_model_max_budget,
+    )
+
+    assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5
+    with pytest.raises(litellm.BudgetExceededError):
+        await limiter.is_team_within_model_budget(
+            team_id="team-1",
+            team_model_max_budget=team_model_max_budget,
+            key_model_max_budget=key_model_max_budget,
+            model="gpt-4",
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "key_entry",
+    [
+        {"time_period": "1d", "tpm_limit": 100},
+        {"time_period": "1d", "rpm_limit": 10},
+        {"budget_limit": -1.0, "time_period": "1d"},
+    ],
+)
+async def test_key_entry_without_a_spend_cap_does_not_lift_the_team_cap(key_entry):
+    """A key row that only rate-limits the model, or has no enforceable cap, leaves the team cap in force."""
+    dual_cache = DualCache()
+    limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+    team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}}
+    key_model_max_budget = {"gpt-4": key_entry}
+
+    await _log_success(
+        limiter,
+        model_group="openai/gpt-4",
+        response_cost=1.5,
+        key_hash="vk-rate-limited",
+        key_model_max_budget=key_model_max_budget,
+        team_id="team-1",
+        team_model_max_budget=team_model_max_budget,
+    )
+
+    assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5
+    with pytest.raises(litellm.BudgetExceededError):
+        await limiter.is_team_within_model_budget(
+            team_id="team-1",
+            team_model_max_budget=team_model_max_budget,
+            key_model_max_budget=key_model_max_budget,
+            model="openai/gpt-4",
+        )

tests/router_unit_tests/test_router_helper_utils.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import os
 import traceback
@@ -10,9 +11,14 @@ import pytest
 import litellm
 from unittest.mock import patch, MagicMock, AsyncMock
 from create_mock_standard_logging_payload import create_standard_logging_payload
-from litellm.types.utils import StandardLoggingPayload
+from litellm.types.utils import ModelResponse, StandardLoggingPayload
+from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+from litellm.caching.dual_cache import DualCache
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.types.caching import RedisPipelineIncrementOperation
+from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute
 from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
-from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
+from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
 
 
 @pytest.fixture
@@ -928,18 +934,7 @@ async def test_set_response_headers(model_list):
 
 
 @pytest.mark.asyncio
-async def test_set_response_headers_subtracts_in_flight_delta(model_list):
-    """
-    LIT-2719: router-derived `x-ratelimit-remaining-*` headers must be
-    post-decrement (match OpenAI/Anthropic vendor semantics) so the proxy's
-    HTTP response headers and the prometheus gauges that read them stay
-    comparable across providers.
-
-    Router's TPM/RPM counter is incremented post-response by
-    `deployment_callback_on_success`, so `get_remaining_model_group_usage`
-    sees pre-decrement values. `set_response_headers` must replay the
-    in-flight increment before writing the headers.
-    """
+async def test_set_response_headers_passes_through_post_increment_counters(model_list):
     from pydantic import BaseModel
 
     class _Usage(BaseModel):
@@ -952,49 +947,10 @@ async def test_set_response_headers_subtracts_in_flight_delta(model_list):
     router = Router(model_list=model_list)
     router.get_remaining_model_group_usage = AsyncMock(
         return_value={
-            "x-ratelimit-remaining-tokens": 1000,
+            "x-ratelimit-remaining-tokens": 958,
             "x-ratelimit-limit-tokens": 1000,
-            "x-ratelimit-remaining-requests": 100,
+            "x-ratelimit-remaining-requests": 99,
             "x-ratelimit-limit-requests": 100,
-        }
-    )
-
-    resp = _Resp()
-    resp._hidden_params = {}
-    await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
-
-    headers = resp._hidden_params["additional_headers"]
-    assert headers["x-ratelimit-remaining-tokens"] == 958
-    assert headers["x-ratelimit-remaining-requests"] == 99
-    # Limit headers pass through unmodified.
-    assert headers["x-ratelimit-limit-tokens"] == 1000
-    assert headers["x-ratelimit-limit-requests"] == 100
-
-
-@pytest.mark.asyncio
-async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_list):
-    """
-    The in-flight replay applies only to the post-incremented TPM/RPM counters
-    (`x-ratelimit-remaining-tokens` / `-requests`). The ITPM/OTPM counters are
-    incremented at reservation time (pre-call), so the input/output token
-    headers already reflect this request and must pass through untouched.
-    """
-    from pydantic import BaseModel
-
-    class _Usage(BaseModel):
-        total_tokens: int = 30
-        prompt_tokens: int = 20
-        completion_tokens: int = 10
-
-    class _Resp(BaseModel):
-        usage: _Usage = _Usage()
-        _hidden_params: dict = {}
-
-    router = Router(model_list=model_list)
-    router.get_remaining_model_group_usage = AsyncMock(
-        return_value={
-            "x-ratelimit-remaining-tokens": 1000,
-            "x-ratelimit-remaining-requests": 100,
             "x-ratelimit-remaining-input-tokens": 1000,
             "x-ratelimit-remaining-output-tokens": 500,
         }
@@ -1005,14 +961,336 @@ async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_l
     await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
 
     headers = resp._hidden_params["additional_headers"]
-    # TPM/RPM headers replay the in-flight increment...
-    assert headers["x-ratelimit-remaining-tokens"] == 970
+    assert headers["x-ratelimit-remaining-tokens"] == 958
     assert headers["x-ratelimit-remaining-requests"] == 99
-    # ...but the reservation-based input/output headers pass through unchanged.
+    assert headers["x-ratelimit-limit-tokens"] == 1000
+    assert headers["x-ratelimit-limit-requests"] == 100
     assert headers["x-ratelimit-remaining-input-tokens"] == 1000
     assert headers["x-ratelimit-remaining-output-tokens"] == 500
 
 
+def _rpm_tpm_router(model_id: str) -> Router:
+    return Router(
+        model_list=[
+            {
+                "model_name": "gpt-5-mini",
+                "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake", "tpm": 1000, "rpm": 100},
+                "model_info": {"id": model_id},
+            }
+        ]
+    )
+
+
+def _ratelimit_headers(response: ModelResponse | CustomStreamWrapper) -> dict[str, int]:
+    return {k: v for k, v in response._hidden_params["additional_headers"].items() if k.startswith("x-ratelimit-")}
+
+
+@pytest.mark.asyncio
+async def test_acompletion_headers_read_post_increment_counter_and_count_once():
+    router = _rpm_tpm_router("lit-3058-async")
+
+    response = await router.acompletion(
+        model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong"
+    )
+    total_tokens = response.usage.total_tokens
+    assert total_tokens > 0
+
+    headers = _ratelimit_headers(response)
+    assert headers["x-ratelimit-remaining-tokens"] == 1000 - total_tokens
+    assert headers["x-ratelimit-remaining-requests"] == 99
+    assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1)
+
+    await asyncio.sleep(0.5)
+    assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1)
+
+
+@pytest.mark.asyncio
+async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deployment_name():
+    router = Router(
+        model_list=[
+            {
+                "model_name": "openai/*",
+                "litellm_params": {"model": "openai/*", "api_key": "sk-fake", "tpm": 1000, "rpm": 100},
+                "model_info": {"id": "lit-3058-wildcard"},
+            }
+        ]
+    )
+
+    response = await router.acompletion(
+        model="openai/gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong"
+    )
+    total_tokens = response.usage.total_tokens
+
+    headers = _ratelimit_headers(response)
+    assert headers["x-ratelimit-remaining-tokens"] == 1000 - total_tokens
+    assert headers["x-ratelimit-remaining-requests"] == 99
+    assert await router.get_model_group_usage("openai/gpt-5-mini") == (total_tokens, 1)
+
+
+@pytest.mark.asyncio
+async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_on_completion():
+    router = _rpm_tpm_router("lit-3058-stream")
+
+    stream = await router.acompletion(
+        model="gpt-5-mini",
+        messages=[{"role": "user", "content": "hi"}],
+        mock_response="pong pong pong",
+        stream=True,
+        stream_options={"include_usage": True},
+    )
+    headers = _ratelimit_headers(stream)
+    assert headers["x-ratelimit-remaining-tokens"] == 1000
+    assert headers["x-ratelimit-remaining-requests"] == 99
+    assert await router.get_model_group_usage("gpt-5-mini") == (0, 1)
+
+    chunks = [chunk async for chunk in stream]
+    total_tokens = chunks[-1].usage.total_tokens
+    assert total_tokens > 0
+
+    await asyncio.sleep(0.5)
+    assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1)
+
+
+@pytest.mark.asyncio
+async def test_deployment_callback_on_success_adds_only_uncounted_tokens():
+    import time
+
+    router = _rpm_tpm_router("lit-3058-callback")
+    standard_logging_payload = create_standard_logging_payload()
+    standard_logging_payload["total_tokens"] = 100
+    kwargs = {
+        "litellm_params": {
+            "metadata": {
+                "deployment": "gpt-5-mini",
+                "model_group": "gpt-5-mini",
+                ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 60,
+            },
+            "model_info": {"id": "lit-3058-callback"},
+        },
+        "standard_logging_object": standard_logging_payload,
+    }
+
+    tpm_key = await router.deployment_callback_on_success(
+        kwargs=kwargs,
+        completion_response=litellm.ModelResponse(model="gpt-5-mini", usage={"total_tokens": 100}),
+        start_time=time.time(),
+        end_time=time.time(),
+    )
+
+    assert tpm_key is not None
+    assert await router.get_model_group_usage("gpt-5-mini") == (40, 0)
+
+
+class _GatedIncrementCache(DualCache):
+    def __init__(self) -> None:
+        super().__init__(in_memory_cache=InMemoryCache())
+        self.first_increment_started = asyncio.Event()
+        self.release_first_increment = asyncio.Event()
+        self.increment_calls = 0
+
+    async def async_increment_cache_pipeline(
+        self,
+        increment_list: list[RedisPipelineIncrementOperation],
+        local_only: bool = False,
+        parent_otel_span: object = None,
+        **kwargs: object,
+    ) -> list[float] | None:
+        self.increment_calls += 1
+        if self.increment_calls == 1:
+            self.first_increment_started.set()
+            await self.release_first_increment.wait()
+        return await super().async_increment_cache_pipeline(
+            increment_list, local_only=local_only, parent_otel_span=parent_otel_span, **kwargs
+        )
+
+
+@pytest.mark.asyncio
+async def test_success_callback_running_during_pre_header_increment_does_not_double_count():
+    router = _rpm_tpm_router("lit-3058-race")
+    cache = _GatedIncrementCache()
+    router.cache = cache
+
+    request = asyncio.ensure_future(
+        router.acompletion(model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong")
+    )
+    await asyncio.wait_for(cache.first_increment_started.wait(), timeout=5)
+    for _ in range(50):
+        if get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1:
+            break
+        await asyncio.sleep(0.1)
+    assert get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1
+    assert cache.increment_calls == 1
+
+    cache.release_first_increment.set()
+    response = await request
+
+    assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1)
+
+
+class _UnavailableIncrementCache(DualCache):
+    def __init__(self) -> None:
+        super().__init__(in_memory_cache=InMemoryCache())
+        self.first_increment_started = asyncio.Event()
+        self.release_first_increment = asyncio.Event()
+        self.increment_calls = 0
+
+    async def async_increment_cache_pipeline(
+        self,
+        increment_list: list[RedisPipelineIncrementOperation],
+        local_only: bool = False,
+        parent_otel_span: object = None,
+        **kwargs: object,
+    ) -> list[float] | None:
+        self.increment_calls += 1
+        if self.increment_calls == 1:
+            self.first_increment_started.set()
+            await self.release_first_increment.wait()
+        raise RuntimeError("cache unavailable")
+
+
+@pytest.mark.asyncio
+async def test_callback_observing_stamp_before_pre_header_increment_fails_leaves_no_stamp_behind():
+    router = _rpm_tpm_router("lit-3058-fail")
+    cache = _UnavailableIncrementCache()
+    router.cache = cache
+    metadata: dict[str, object] = {}
+
+    request = asyncio.ensure_future(
+        router.acompletion(
+            model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong", metadata=metadata
+        )
+    )
+    await asyncio.wait_for(cache.first_increment_started.wait(), timeout=5)
+    assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 30
+    for _ in range(50):
+        if get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1:
+            break
+        await asyncio.sleep(0.1)
+    assert get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1
+    assert cache.increment_calls == 1
+
+    cache.release_first_increment.set()
+    response = await request
+
+    assert response.usage.total_tokens == 30
+    assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in metadata
+    assert _ratelimit_headers(response)["x-ratelimit-remaining-requests"] == 100
+    assert await router.get_model_group_usage("gpt-5-mini") == (None, None)
+
+
+@pytest.mark.asyncio
+async def test_increment_deployment_usage_for_response_skips_session_wrappers():
+    router = _rpm_tpm_router("lit-3058-ws")
+    request_kwargs = {
+        "model": "gpt-5-mini",
+        "litellm_metadata": {"model_group": "gpt-5-mini", "model_info": {"id": "lit-3058-ws"}},
+    }
+
+    await router.increment_deployment_usage_for_response(response=None, request_kwargs=request_kwargs)
+
+    assert await router.get_model_group_usage("gpt-5-mini") == (None, None)
+    assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in request_kwargs["litellm_metadata"]
+
+
+@pytest.mark.asyncio
+async def test_increment_deployment_usage_writes_only_positive_deltas_for_limited_deployments():
+    router = _rpm_tpm_router("lit-3058-delta")
+    unlimited = Router(
+        model_list=[
+            {
+                "model_name": "gpt-5-mini",
+                "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"},
+                "model_info": {"id": "lit-3058-unlimited"},
+            }
+        ]
+    )
+
+    tpm_key = await router._increment_deployment_usage(
+        deployment_id="lit-3058-delta",
+        deployment_name="gpt-5-mini",
+        model_group="gpt-5-mini",
+        total_tokens=25,
+        rpm_increment=1,
+        parent_otel_span=None,
+    )
+    assert tpm_key is not None
+    assert await router.get_model_group_usage("gpt-5-mini") == (25, 1)
+
+    assert (
+        await router._increment_deployment_usage(
+            deployment_id="lit-3058-delta",
+            deployment_name="gpt-5-mini",
+            model_group="gpt-5-mini",
+            total_tokens=0,
+            rpm_increment=0,
+            parent_otel_span=None,
+        )
+        is None
+    )
+    assert await router.get_model_group_usage("gpt-5-mini") == (25, 1)
+
+    assert (
+        await unlimited._increment_deployment_usage(
+            deployment_id="lit-3058-unlimited",
+            deployment_name="gpt-5-mini",
+            model_group="gpt-5-mini",
+            total_tokens=25,
+            rpm_increment=1,
+            parent_otel_span=None,
+        )
+        is None
+    )
+    assert await unlimited.get_model_group_usage("gpt-5-mini") == (None, None)
+
+
+def _shared_redis_stub(store: dict) -> MagicMock:
+    from litellm.caching.redis_cache import RedisCache
+
+    async def increment_pipeline(increment_list, **kwargs):
+        for op in increment_list:
+            store[op["key"]] = store.get(op["key"], 0.0) + op["increment_value"]
+        return [store[op["key"]] for op in increment_list]
+
+    async def batch_get(keys, **kwargs):
+        return {key: store.get(key) for key in keys}
+
+    redis_stub = MagicMock(spec=RedisCache)
+    redis_stub.async_increment_pipeline = increment_pipeline
+    redis_stub.async_batch_get_cache = batch_get
+    return redis_stub
+
+
+@pytest.mark.asyncio
+async def test_headers_on_fresh_worker_reflect_shared_redis_usage():
+    store: dict = {}
+    worker_a = _rpm_tpm_router("lit-3058-workers")
+    worker_b = _rpm_tpm_router("lit-3058-workers")
+    worker_a.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache())
+    worker_b.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache())
+
+    messages = [{"role": "user", "content": "hi"}]
+    tokens_on_a = 0
+    for _ in range(3):
+        response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
+        tokens_on_a += response.usage.total_tokens
+
+    response = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
+    headers = _ratelimit_headers(response)
+    assert headers["x-ratelimit-remaining-requests"] == 96
+    assert headers["x-ratelimit-remaining-tokens"] == 1000 - tokens_on_a - response.usage.total_tokens
+
+    counted_tokens = tokens_on_a + response.usage.total_tokens
+    for _ in range(2):
+        response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
+        counted_tokens += response.usage.total_tokens
+
+    stream = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong", stream=True)
+    stream_headers = _ratelimit_headers(stream)
+    assert stream_headers["x-ratelimit-remaining-requests"] == 93
+    assert stream_headers["x-ratelimit-remaining-tokens"] == 1000 - counted_tokens
+    assert [chunk async for chunk in stream]
+
+
 @pytest.mark.asyncio
 async def test_get_model_group_io_token_usage_sums_across_deployments():
     """
@@ -1154,8 +1432,8 @@ async def test_set_response_headers_native_input_token_header_does_not_suppress_
     await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
 
     headers = resp._hidden_params["additional_headers"]
-    assert headers["x-ratelimit-remaining-tokens"] == 958
-    assert headers["x-ratelimit-remaining-requests"] == 99
+    assert headers["x-ratelimit-remaining-tokens"] == 1000
+    assert headers["x-ratelimit-remaining-requests"] == 100
     # the provider's native header is left untouched
     assert headers["x-ratelimit-remaining-input-tokens"] == 5
 
@@ -1187,7 +1465,7 @@ async def test_set_response_headers_native_token_header_does_not_suppress_io_hea
 
     headers = resp._hidden_params["additional_headers"]
     assert headers["x-ratelimit-remaining-tokens"] == 5
-    assert headers["x-ratelimit-remaining-requests"] == 99
+    assert headers["x-ratelimit-remaining-requests"] == 100
     assert headers["x-ratelimit-remaining-input-tokens"] == 900
     assert headers["x-ratelimit-remaining-output-tokens"] == 450
 
@@ -1196,8 +1474,7 @@ async def test_set_response_headers_native_token_header_does_not_suppress_io_hea
 async def test_set_response_headers_handles_missing_usage(model_list):
     """
     Streaming chunks and some response shapes may lack a `usage` attribute or
-    populated `total_tokens`. The in-flight subtraction must default to 0
-    tokens (still subtract 1 from requests) and never raise.
+    populated `total_tokens`. Header composition must not depend on usage and never raise.
     """
     from pydantic import BaseModel
 
@@ -1218,7 +1495,7 @@ async def test_set_response_headers_handles_missing_usage(model_list):
 
     headers = resp._hidden_params["additional_headers"]
     assert headers["x-ratelimit-remaining-tokens"] == 1000
-    assert headers["x-ratelimit-remaining-requests"] == 99
+    assert headers["x-ratelimit-remaining-requests"] == 100
 
 
 @pytest.mark.asyncio

tests/test_litellm/integrations/otel/test_otel_v2_components.py

@@ -5,6 +5,7 @@ builders, and the registry validator's failure paths. Needs the OTel SDK."""
 import json
 import threading
 from collections.abc import Iterator
+from contextvars import Context as ContextVarContext
 from dataclasses import replace
 from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer
 
@@ -15,6 +16,8 @@ pytest.importorskip("opentelemetry")
 from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (  # noqa: E402
     ExportTraceServiceRequest,
 )
+from opentelemetry import baggage  # noqa: E402
+from opentelemetry.context import attach, detach  # noqa: E402
 from opentelemetry.sdk.metrics import MeterProvider  # noqa: E402
 from opentelemetry.sdk.metrics.export import InMemoryMetricReader  # noqa: E402
 from opentelemetry.sdk.trace import TracerProvider  # noqa: E402
@@ -26,7 +29,10 @@ from opentelemetry.sdk.trace.export import (  # noqa: E402
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import (  # noqa: E402
     InMemorySpanExporter,
 )
-from opentelemetry.trace import SpanKind  # noqa: E402
+from opentelemetry.trace import SpanKind, get_current_span  # noqa: E402
+from opentelemetry.trace.propagation.tracecontext import (  # noqa: E402
+    TraceContextTextMapPropagator,
+)
 
 from litellm.integrations.otel.plumbing import context as ctx_mod  # noqa: E402
 from litellm.integrations.otel.plumbing import providers  # noqa: E402
@@ -464,6 +470,115 @@ def test_extract_traceparent():
     assert ctx_mod.extract_traceparent({"x": "y"}) is None
 
 
+def _test_tracer():
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    return provider.get_tracer("test")
+
+
+def test_inject_trace_context_prefers_request_root_span():
+    def run():
+        tracer = _test_tracer()
+        with tracer.start_as_current_span("root") as root:
+            ctx_mod.set_request_root_span(root)
+            result = ctx_mod.inject_trace_context(
+                {"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"}
+            )
+            propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
+            return result, root, propagated
+
+    result, root, propagated = ContextVarContext().run(run)
+    assert result["traceparent"] != "00-11111111111111111111111111111111-2222222222222222-01"
+    assert propagated.get_span_context().trace_id == root.get_span_context().trace_id
+    assert propagated.get_span_context().span_id == root.get_span_context().span_id
+
+
+def test_inject_trace_context_uses_ambient_span_without_request_root():
+    def run():
+        tracer = _test_tracer()
+        with tracer.start_as_current_span("ambient") as ambient:
+            result = ctx_mod.inject_trace_context({})
+            propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
+            return ambient, propagated
+
+    ambient, propagated = ContextVarContext().run(run)
+    assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id
+    assert propagated.get_span_context().span_id == ambient.get_span_context().span_id
+
+
+def test_inject_trace_context_replaces_stale_trace_headers():
+    def run():
+        tracer = _test_tracer()
+        with tracer.start_as_current_span("ambient") as ambient:
+            headers = {
+                "Traceparent": "00-" + "a" * 32 + "-" + "b" * 16 + "-01",
+                "Tracestate": "vendor=old",
+                "x-keep": "1",
+            }
+            result = ctx_mod.inject_trace_context(headers)
+            propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
+            return result, ambient, propagated
+
+    result, ambient, propagated = ContextVarContext().run(run)
+    assert sum(key.lower() == "traceparent" for key in result) == 1
+    assert not any(key.lower() == "tracestate" for key in result)
+    assert result["x-keep"] == "1"
+    assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id
+
+
+def test_inject_trace_context_prefers_explicit_parent_span_over_root_and_ambient():
+    def run():
+        tracer = _test_tracer()
+        parent = tracer.start_span("litellm_request")
+        with tracer.start_as_current_span("ambient") as ambient:
+            ctx_mod.set_request_root_span(ambient)
+            result = ctx_mod.inject_trace_context({}, parent_span=parent)
+            propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
+            return parent, ambient, propagated
+
+    parent, ambient, propagated = ContextVarContext().run(run)
+    assert propagated.get_span_context().trace_id == parent.get_span_context().trace_id
+    assert propagated.get_span_context().span_id == parent.get_span_context().span_id
+    assert propagated.get_span_context().span_id != ambient.get_span_context().span_id
+
+
+def test_inject_trace_context_skips_unusable_parent_span():
+    def run():
+        tracer = _test_tracer()
+        with tracer.start_as_current_span("ambient") as ambient:
+            result = ctx_mod.inject_trace_context({}, parent_span=object())
+            propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
+            return ambient, propagated
+
+    ambient, propagated = ContextVarContext().run(run)
+    assert propagated.get_span_context().span_id == ambient.get_span_context().span_id
+
+
+def test_inject_trace_context_returns_headers_unchanged_without_context():
+    headers = {"x-custom": "value"}
+
+    result = ContextVarContext().run(lambda: ctx_mod.inject_trace_context(headers))
+
+    assert result == headers
+    assert "traceparent" not in result
+    assert result is not headers
+
+
+def test_inject_trace_context_does_not_forward_baggage():
+    def run():
+        tracer = _test_tracer()
+        with tracer.start_as_current_span("ambient"):
+            token = attach(baggage.set_baggage("litellm.team.id", "team"))
+            try:
+                return ctx_mod.inject_trace_context({})
+            finally:
+                detach(token)
+
+    result = ContextVarContext().run(run)
+    assert "baggage" not in result
+
+
 def test_set_request_baggage_empty_returns_context():
     assert ctx_mod.set_request_baggage({}) is not None
 

tests/test_litellm/integrations/test_opentelemetry.py

@@ -5424,6 +5424,65 @@ class TestGetSpanContextLitellmMetadataFallback(unittest.TestCase):
         self.assertIsNone(detected_span)
 
 
+class TestInboundTraceContextKeepsCallerTracestate(unittest.TestCase):
+    """The request span built from inbound W3C headers must carry the caller's
+    tracestate so outbound propagation (passthrough) re-emits it instead of
+    dropping it alongside the stripped stale header."""
+
+    CALLER_TRACEPARENT = "00-" + "a" * 32 + "-" + "b" * 16 + "-01"
+    CALLER_TRACESTATE = "vendor=abc,other=xyz"
+
+    def _otel(self):
+        provider = TracerProvider()
+        provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
+        otel = OpenTelemetry()
+        otel.tracer = provider.get_tracer(__name__)
+        return otel
+
+    def test_request_span_propagates_caller_tracestate_downstream(self):
+        from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+        from litellm.integrations.otel.plumbing.context import inject_trace_context
+
+        inbound = {"traceparent": self.CALLER_TRACEPARENT, "tracestate": self.CALLER_TRACESTATE}
+        span = self._otel().create_litellm_proxy_request_started_span(
+            start_time=datetime.now(timezone.utc), headers=inbound
+        )
+        outbound = inject_trace_context(inbound, parent_span=span)
+        span.end()
+
+        propagated = trace.get_current_span(TraceContextTextMapPropagator().extract(outbound)).get_span_context()
+        self.assertEqual(outbound["tracestate"], self.CALLER_TRACESTATE)
+        self.assertEqual(propagated.trace_id, span.get_span_context().trace_id)
+        self.assertEqual(propagated.span_id, span.get_span_context().span_id)
+        self.assertNotEqual(outbound["traceparent"], self.CALLER_TRACEPARENT)
+
+    def test_request_span_without_caller_tracestate_emits_none(self):
+        from litellm.integrations.otel.plumbing.context import inject_trace_context
+
+        inbound = {"traceparent": self.CALLER_TRACEPARENT}
+        span = self._otel().create_litellm_proxy_request_started_span(
+            start_time=datetime.now(timezone.utc), headers=inbound
+        )
+        outbound = inject_trace_context(inbound, parent_span=span)
+        span.end()
+
+        self.assertNotIn("tracestate", outbound)
+        self.assertNotEqual(outbound["traceparent"], self.CALLER_TRACEPARENT)
+
+    def test_span_context_from_header_keeps_caller_tracestate(self):
+        kwargs = {
+            "litellm_params": {
+                "proxy_server_request": {
+                    "headers": {"traceparent": self.CALLER_TRACEPARENT, "tracestate": self.CALLER_TRACESTATE}
+                }
+            }
+        }
+        ctx, detected_span = self._otel()._get_span_context(kwargs)
+        self.assertIsNone(detected_span)
+        self.assertEqual(trace.get_current_span(ctx).get_span_context().trace_state.to_header(), self.CALLER_TRACESTATE)
+
+
 class TestEndProxySpanLitellmMetadataFallback(unittest.TestCase):
     """
     Tests for _end_proxy_span_from_kwargs() falling back to litellm_metadata.

tests/test_litellm/integrations/test_s3.py

@@ -1,16 +1,24 @@
+import copy
+import json
 from datetime import datetime
 from unittest.mock import MagicMock, patch
 
+import pytest
+
 import litellm
 from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES, MAX_S3_OBJECT_KEY_BYTES
-from litellm.integrations.s3 import S3Logger
+from litellm.integrations.s3 import S3Logger, prompts_only_payload, resolve_s3_log_prompts_only
 
 TEST_KMS_KEY_ARN = "arn:aws:kms:us-east-1:111122223333:key/test-key-id"
+TEST_MESSAGES = [{"role": "user", "content": "Reply with exactly the word PINEAPPLE."}]
+TEST_RESPONSE = {"choices": [{"message": {"role": "assistant", "content": "PINEAPPLE"}}]}
 
 
 def _standard_logging_payload(response_id: str = "chatcmpl-test-id") -> dict:
     return {
         "id": response_id,
+        "messages": copy.deepcopy(TEST_MESSAGES),
+        "response": copy.deepcopy(TEST_RESPONSE),
         "metadata": {"user_api_key_team_alias": None},
     }
 
@@ -22,7 +30,9 @@ def _log_event_kwargs(response_id: str = "chatcmpl-test-id") -> dict:
     }
 
 
-def _run_log_event(callback_params: dict, response_id: str = "chatcmpl-test-id") -> MagicMock:
+def _run_log_event(
+    callback_params: dict, response_id: str = "chatcmpl-test-id", log_kwargs: dict[str, object] | None = None
+) -> MagicMock:
     original = litellm.s3_callback_params
     litellm.s3_callback_params = callback_params
     try:
@@ -31,7 +41,7 @@ def _run_log_event(callback_params: dict, response_id: str = "chatcmpl-test-id")
             mock_boto3_client.return_value = mock_s3_client
             logger = S3Logger()
             logger.log_event(
-                kwargs=_log_event_kwargs(response_id),
+                kwargs=_log_event_kwargs(response_id) if log_kwargs is None else log_kwargs,
                 response_obj={"id": response_id},
                 start_time=datetime(2026, 7, 30, 12, 0, 0),
                 end_time=datetime(2026, 7, 30, 12, 0, 1),
@@ -182,3 +192,123 @@ def test_put_object_keeps_the_configured_path_intact_when_only_the_id_has_to_shr
     key = mock_s3_client.put_object.call_args.kwargs["Key"]
     assert key.startswith(long_path + "/2026-07-30/")
     assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES
+
+
+def _uploaded_body(mock_s3_client: MagicMock) -> dict[str, object]:
+    return json.loads(mock_s3_client.put_object.call_args.kwargs["Body"])
+
+
+def test_log_event_prompts_only_drops_response_and_keeps_messages(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False)
+    log_kwargs = _log_event_kwargs()
+    original_payload = copy.deepcopy(log_kwargs["standard_logging_object"])
+
+    mock_s3_client = _run_log_event(
+        {"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1", "s3_log_prompts_only": True},
+        log_kwargs=log_kwargs,
+    )
+
+    body = _uploaded_body(mock_s3_client)
+    assert body["messages"] == TEST_MESSAGES
+    assert body["response"] is None
+    assert body["id"] == "chatcmpl-test-id"
+    assert log_kwargs["standard_logging_object"] == original_payload
+
+
+def test_log_event_default_keeps_response(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False)
+
+    mock_s3_client = _run_log_event({"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1"})
+
+    body = _uploaded_body(mock_s3_client)
+    assert body["response"] == TEST_RESPONSE
+    assert body["messages"] == TEST_MESSAGES
+
+
+def test_log_event_reads_prompts_only_env_var_at_log_time(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False)
+    original = litellm.s3_callback_params
+    litellm.s3_callback_params = {"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1"}
+    try:
+        with patch("boto3.client") as mock_boto3_client:
+            mock_s3_client = MagicMock()
+            mock_boto3_client.return_value = mock_s3_client
+            logger = S3Logger()
+            monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true")
+            logger.log_event(
+                kwargs=_log_event_kwargs(),
+                response_obj={"id": "chatcmpl-test-id"},
+                start_time=datetime(2026, 7, 30, 12, 0, 0),
+                end_time=datetime(2026, 7, 30, 12, 0, 1),
+                print_verbose=lambda *args, **kwargs: None,
+            )
+    finally:
+        litellm.s3_callback_params = original
+
+    body = _uploaded_body(mock_s3_client)
+    assert body["response"] is None
+    assert body["messages"] == TEST_MESSAGES
+
+
+def test_log_event_explicit_false_param_beats_env_var(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true")
+
+    mock_s3_client = _run_log_event(
+        {"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1", "s3_log_prompts_only": False}
+    )
+
+    assert _uploaded_body(mock_s3_client)["response"] == TEST_RESPONSE
+
+
+def test_s3_logger_init_does_not_mutate_global_callback_params(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("MY_S3_BUCKET", "resolved-bucket")
+    callback_params = {"s3_bucket_name": "os.environ/MY_S3_BUCKET", "s3_region_name": "us-east-1"}
+    snapshot = copy.deepcopy(callback_params)
+    original = litellm.s3_callback_params
+    litellm.s3_callback_params = callback_params
+    try:
+        with patch("boto3.client"):
+            logger = S3Logger()
+    finally:
+        litellm.s3_callback_params = original
+
+    assert logger.bucket_name == "resolved-bucket"
+    assert callback_params == snapshot
+
+
+@pytest.mark.parametrize(
+    "configured,env_value,expected",
+    [
+        (True, None, True),
+        (False, "true", False),
+        ("true", None, True),
+        ("False", "true", False),
+        ("1", None, True),
+        ("0", None, False),
+        (" yes ", None, True),
+        (None, None, False),
+        (None, "true", True),
+        (None, "false", False),
+        (None, "", False),
+        ("", "true", False),
+    ],
+)
+def test_resolve_s3_log_prompts_only(configured: object, env_value: str | None, expected: bool):
+    environ = {} if env_value is None else {"S3_LOG_PROMPTS_ONLY": env_value}
+    assert resolve_s3_log_prompts_only(configured, environ) is expected
+
+
+def test_resolve_s3_log_prompts_only_unparseable_value_fails_toward_prompts_only():
+    assert resolve_s3_log_prompts_only("enabled", {}) is True
+
+
+def test_prompts_only_payload_returns_copy_with_response_cleared():
+    payload = _standard_logging_payload()
+    snapshot = copy.deepcopy(payload)
+
+    stripped = prompts_only_payload(payload)
+
+    assert stripped["response"] is None
+    assert stripped["messages"] == TEST_MESSAGES
+    assert stripped is not payload
+    assert payload == snapshot