Merge branch 'main' into litellm_isolate_generic_api_ndjson_test

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

@@ -21,6 +21,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import (
     OTELSemconvCategory,
     parse_semconv_opt_in,
 )
+from litellm.integrations.otel.mappers.utils import drop_none
 from litellm.integrations.otel.model.baggage import promoted_metadata
 from litellm.integrations.otel.model.db_endpoint import db_span_attributes
 from litellm.integrations.otel.model.metadata import flatten_metadata
@@ -208,6 +209,20 @@ def _resolve_metric_attribute_filter(
     )
 
 
+def _provider_label(custom_llm_provider: object) -> str | None:
+    """The provider label for one call's metrics and events, or None when the
+    call carries no provider.
+
+    Every attribute set drops None before export, so the label is simply absent
+    in that case: the OTLP encoder rejects a None attribute value outright, and a
+    placeholder would mint a permanent metric series that no operator can act
+    on. Mirrors the v2 integration's ``_provider_attributes``.
+    """
+    if not isinstance(custom_llm_provider, str) or not custom_llm_provider:
+        return None
+    return custom_llm_provider
+
+
 def _normalize_team_metadata_keys(value: str | Iterable[object] | None) -> list[str]:
     """Coerce a team-metadata allowlist from a list or comma-separated string.
 
@@ -1616,19 +1631,22 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
         ) = _resolve_metric_attribute_filter(attributes)
         self._metric_attr_filter_resolved = True
 
-    def _filter_metric_attributes(self, attrs: dict[str, str]) -> dict[str, str]:
+    def _filter_metric_attributes(self, attrs: Mapping[str, str | None]) -> dict[str, str]:
         if not self._metric_attr_filter_resolved:
             self._ensure_metric_attribute_filter()
+        return {k: v for k, v in attrs.items() if v is not None and self._metric_attribute_allowed(k)}
+
+    def _metric_attribute_allowed(self, key: str) -> bool:
         if self._metric_attr_include is not None:
-            return {k: v for k, v in attrs.items() if k in self._metric_attr_include}
+            return key in self._metric_attr_include
         if self._metric_attr_exclude is not None:
-            return {k: v for k, v in attrs.items() if k not in self._metric_attr_exclude}
-        return attrs
+            return key not in self._metric_attr_exclude
+        return True
 
     def _record_metrics(self, kwargs, response_obj, start_time, end_time):
         duration_s: Final = (end_time - start_time).total_seconds()
         params: Final = kwargs.get("litellm_params") or {}
-        provider: Final = params.get("custom_llm_provider", "Unknown")
+        provider: Final = _provider_label(params.get("custom_llm_provider"))
 
         common_attrs = {
             "gen_ai.operation.name": (
@@ -1872,7 +1890,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
         otel_logger: Final = self._logger_provider.get_logger(LITELLM_LOGGER_NAME)
 
         parent_ctx: Final = span.get_span_context()
-        provider: Final = (kwargs.get("litellm_params") or {}).get("custom_llm_provider", "Unknown")
+        provider: Final = _provider_label((kwargs.get("litellm_params") or {}).get("custom_llm_provider"))
 
         if self._gen_ai_semconv_latest_experimental:
             self._emit_inference_details_event(
@@ -1909,7 +1927,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
                 severity_number=SeverityNumber.INFO,
                 severity_text="INFO",
                 body=body,
-                attributes=attrs,
+                attributes=drop_none(attrs),
             )
             otel_logger.emit(log_record)
 
@@ -1941,7 +1959,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
                 severity_number=SeverityNumber.INFO,
                 severity_text="INFO",
                 body=body,
-                attributes=attrs,
+                attributes=drop_none(attrs),
             )
             otel_logger.emit(log_record)
 

litellm/integrations/opentelemetry_utils/gen_ai_semconv.py

@@ -33,6 +33,7 @@ from datetime import datetime
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Final
 
+from litellm.integrations.otel.mappers.utils import drop_none
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
 if TYPE_CHECKING:
@@ -195,13 +196,16 @@ class OTELGenAISemconvMixin:
             if value:
                 self.safe_set_attribute(span=span, key=semconv_key, value=value)
 
-    def _build_inference_details_attrs(self, kwargs: dict, response_obj: dict, provider: str) -> dict[str, str]:
+    def _build_inference_details_attrs(
+        self, kwargs: dict, response_obj: dict, provider: str | None
+    ) -> dict[str, str | None]:
         """Build the attribute payload for the inference-details event.
 
-        Always includes provider/operation; input/output messages are added
+        Always includes operation and provider (None when the call carries none,
+        dropped before the event is emitted); input/output messages are added
         only when content capture is enabled and non-empty. Mixin-internal.
         """
-        attrs: Final[dict[str, str]] = {
+        attrs: Final[dict[str, str | None]] = {
             "event_name": _INFERENCE_DETAILS_EVENT_NAME,
             "gen_ai.provider.name": provider,
             "gen_ai.operation.name": self._gen_ai_operation_name(kwargs),
@@ -221,7 +225,7 @@ class OTELGenAISemconvMixin:
         self,
         kwargs: dict,
         response_obj: dict,
-        provider: str,
+        provider: str | None,
         otel_logger,
         parent_ctx,
     ) -> None:
@@ -239,6 +243,6 @@ class OTELGenAISemconvMixin:
             severity_number=SeverityNumber.INFO,
             severity_text="INFO",
             body=None,
-            attributes=self._build_inference_details_attrs(kwargs, response_obj, provider),
+            attributes=drop_none(self._build_inference_details_attrs(kwargs, response_obj, provider)),
         )
         otel_logger.emit(log_record)

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/chat/handler.py

@@ -632,6 +632,7 @@ class ModelResponseIterator:
         self.tool_name_reverse_map: dict[str, str] = tool_name_reverse_map or {}
         # Generate response ID once per stream to match OpenAI-compatible behavior
         self.response_id = _generate_id()
+        self.served_model: str | None = None
 
         # Track if we're currently streaming a response_format tool
         self.is_response_format_tool: bool = False
@@ -1067,6 +1068,9 @@ class ModelResponseIterator:
                 }
                 """
                 message_start_block: Final = MessageStartBlock(**chunk)
+                start_message: Final = message_start_block["message"]
+                if "model" in start_message:
+                    self.served_model = start_message["model"]
                 if "usage" in message_start_block["message"]:
                     usage = self._handle_usage(anthropic_usage_chunk=message_start_block["message"]["usage"])
             elif type_chunk == "error":
@@ -1098,6 +1102,7 @@ class ModelResponseIterator:
                 ],
                 usage=usage,
                 id=self.response_id,
+                model=self.served_model,
             )
 
             return returned_chunk

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/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/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(

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

@@ -4340,6 +4340,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 +4780,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/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/test_opentelemetry.py

@@ -15,10 +15,11 @@ from unittest.mock import MagicMock, patch
 
 # Adds the grandparent directory to sys.path to allow importing project modules
 from opentelemetry import trace
+from opentelemetry.sdk._logs import LogData
 from opentelemetry.sdk._logs import LoggerProvider as OTLoggerProvider
 from opentelemetry.sdk._logs.export import InMemoryLogExporter, SimpleLogRecordProcessor
 from opentelemetry.sdk.metrics import MeterProvider
-from opentelemetry.sdk.metrics.export import InMemoryMetricReader
+from opentelemetry.sdk.metrics.export import InMemoryMetricReader, MetricsData
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
@@ -5921,13 +5922,11 @@ class TestOpenTelemetryMetricAttributeFiltering(unittest.TestCase):
                 }
             )
 
-    def test_no_filter_returns_attrs_object_unchanged(self):
-        """The no-config path is a hot-path no-op: it returns the same dict
-        object, so default emission pays zero copy cost. Locking identity makes
-        a future refactor that always copies/filters trip here."""
+    def test_no_filter_keeps_every_attribute(self):
+        """The no-config path drops nothing: every attribute the caller set reaches the meter."""
         otel = OpenTelemetry(config=OpenTelemetryConfig(exporter="console"))
         attrs = {"gen_ai.request.model": "m", "hidden_params": "{}"}
-        self.assertIs(otel._filter_metric_attributes(attrs), attrs)
+        self.assertEqual(otel._filter_metric_attributes(attrs), attrs)
 
     def test_token_type_discriminator_rejected_from_either_list(self):
         """gen_ai.token.type is a structural discriminator stamped onto the
@@ -6068,6 +6067,118 @@ class TestOTELServiceTierAttributes(unittest.TestCase):
         self.assertEqual(attributes[self.RESPONSE_KEY], "tier-added-by-provider-later")
 
 
+class TestOpenTelemetryProviderlessCallAttributes(unittest.TestCase):
+    """Regression for the OTLP exporter rejecting a None gen_ai.system or gen_ai.request.model
+    attribute on every export cycle."""
+
+    HERE = os.path.dirname(__file__)
+    POLL_INTERVAL = 0.05
+    POLL_TIMEOUT = 2.0
+
+    def _providerless_kwargs(self) -> tuple[dict[str, object], dict[str, object]]:
+        with open(os.path.join(self.HERE, "open_telemetry", "data", "captured_kwargs.json")) as f:
+            kwargs = json.load(f)
+        with open(os.path.join(self.HERE, "open_telemetry", "data", "captured_response.json")) as f:
+            response_obj = json.load(f)
+        kwargs["litellm_params"]["custom_llm_provider"] = None
+        return kwargs, response_obj
+
+    def _modelless_kwargs(self) -> tuple[dict[str, object], dict[str, object]]:
+        kwargs, response_obj = self._providerless_kwargs()
+        kwargs["model"] = None
+        return kwargs, response_obj
+
+    def _recorded_metrics(self, kwargs: dict[str, object], response_obj: dict[str, object]) -> MetricsData | None:
+        metric_reader = InMemoryMetricReader()
+        meter_provider = MeterProvider(metric_readers=[metric_reader])
+        tracer_provider = TracerProvider()
+        tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
+        otel = OpenTelemetry(
+            config=OpenTelemetryConfig(exporter="console", enable_metrics=True),
+            tracer_provider=tracer_provider,
+            meter_provider=meter_provider,
+        )
+        otel.tracer = tracer_provider.get_tracer(__name__)
+
+        start = datetime.utcnow()
+        otel._handle_success(kwargs, response_obj, start, start + timedelta(seconds=1))
+
+        deadline = time.time() + self.POLL_TIMEOUT
+        while time.time() < deadline:
+            data = metric_reader.get_metrics_data()
+            if data and getattr(data, "resource_metrics", None):
+                return data
+            time.sleep(self.POLL_INTERVAL)
+        return None
+
+    def _emitted_log_records(self, semconv_opt_in: str) -> tuple[LogData, ...]:
+        log_exporter = InMemoryLogExporter()
+        logger_provider = OTLoggerProvider()
+        logger_provider.add_log_record_processor(SimpleLogRecordProcessor(log_exporter))
+        with patch.dict(os.environ, {"OTEL_SEMCONV_STABILITY_OPT_IN": semconv_opt_in}):
+            handler = OpenTelemetry(
+                config=OpenTelemetryConfig(exporter="console", enable_events=True),
+                logger_provider=logger_provider,
+            )
+        handler.message_logging = True
+
+        kwargs, response_obj = self._providerless_kwargs()
+        span = handler.tracer.start_span("test")
+        with self.assertNoLogs("opentelemetry.attributes", level="WARNING"):
+            handler._emit_semantic_logs(kwargs, response_obj, span)
+        span.end()
+        handler._logger_provider.force_flush(2000)
+        return log_exporter.get_finished_logs()
+
+    def _assert_every_attribute_encodes(self, attrs: dict[str, object]) -> None:
+        from opentelemetry.exporter.otlp.proto.common._internal import _encode_attributes
+
+        self.assertEqual(len(_encode_attributes(attrs) or []), len(attrs))
+
+    def _recorded_data_points(self, kwargs: dict[str, object], response_obj: dict[str, object]) -> list[object]:
+        data = self._recorded_metrics(kwargs, response_obj)
+        self.assertIsNotNone(data, "no metrics were recorded")
+        data_points = [
+            dp
+            for rm in data.resource_metrics
+            for sm in rm.scope_metrics
+            for m in sm.metrics
+            for dp in m.data.data_points
+        ]
+        self.assertTrue(data_points, "no metric data points were recorded")
+        return data_points
+
+    def test_metrics_are_encodable_and_carry_no_provider_label(self):
+        kwargs, response_obj = self._providerless_kwargs()
+        for dp in self._recorded_data_points(kwargs, response_obj):
+            self.assertNotIn("gen_ai.system", dp.attributes)
+            self.assertEqual(dp.attributes["gen_ai.request.model"], kwargs["model"])
+            self._assert_every_attribute_encodes(dict(dp.attributes))
+
+    def test_metrics_are_encodable_and_carry_no_model_label_when_the_call_has_none(self):
+        for dp in self._recorded_data_points(*self._modelless_kwargs()):
+            self.assertNotIn("gen_ai.request.model", dp.attributes)
+            self._assert_every_attribute_encodes(dict(dp.attributes))
+
+    def test_legacy_content_events_are_encodable_and_carry_no_provider_label(self):
+        logs = self._emitted_log_records("")
+        self.assertTrue(logs, "no content events were emitted")
+        for log in logs:
+            attrs = dict(log.log_record.attributes or {})
+            self.assertNotIn("gen_ai.system", attrs)
+            self.assertNotIn(None, attrs.values())
+            self._assert_every_attribute_encodes(attrs)
+
+    def test_inference_details_event_is_encodable_and_carries_no_provider_label(self):
+        logs = self._emitted_log_records("gen_ai_latest_experimental")
+        self.assertEqual(len(logs), 1)
+        attrs = dict(logs[0].log_record.attributes or {})
+        self.assertEqual(attrs["event_name"], "gen_ai.client.inference.operation.details")
+        self.assertNotIn("gen_ai.provider.name", attrs)
+        self.assertNotIn(None, attrs.values())
+        self._assert_every_attribute_encodes(attrs)
+
+
 class TestDynamicTracerProviderCache(unittest.TestCase):
     """Every credential-scoped TracerProvider that owns its exporter also owns a
     BatchSpanProcessor worker thread that only stops on shutdown, so the cache holding them

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

tests/test_litellm/integrations/test_s3_v2.py

@@ -1,8 +1,11 @@
 import asyncio
+import copy
+import json
 import re
 import sys
 import textwrap
 import uuid
+from collections.abc import Awaitable, Callable
 from contextlib import asynccontextmanager
 from datetime import datetime
 from pathlib import Path
@@ -10,6 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, call, patch
 
 import httpx
 import pytest
+import respx
 
 from litellm.integrations.s3_v2 import S3Logger
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@@ -2310,3 +2314,137 @@ def _s3_logger_for_region(region_name: str) -> S3Logger:
 )
 def test_build_object_url_uses_partition_dns_suffix(region_name: str, expected_url: str) -> None:
     assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url
+
+
+def _prompts_only_logger(s3_log_prompts_only: bool | None = None) -> S3Logger:
+    return S3Logger(
+        s3_bucket_name="test-bucket",
+        s3_aws_access_key_id="test-key",
+        s3_aws_secret_access_key="test-secret",
+        s3_region_name="us-east-1",
+        s3_log_prompts_only=s3_log_prompts_only,
+    )
+
+
+def _chat_payload() -> StandardLoggingPayload:
+    return StandardLoggingPayload(
+        id="chatcmpl-prompts-only",
+        messages=[{"role": "user", "content": "Reply with exactly the word PINEAPPLE."}],
+        response={"choices": [{"message": {"role": "assistant", "content": "PINEAPPLE"}}]},
+        metadata={"user_api_key_team_alias": None},
+    )
+
+
+async def _queued_body_via_async_upload(
+    logger: S3Logger, log_event: Callable[..., Awaitable[None]]
+) -> dict[str, object]:
+    payload = _chat_payload()
+    original = copy.deepcopy(payload)
+    await log_event(
+        kwargs={"standard_logging_object": payload},
+        response_obj=None,
+        start_time=datetime(2026, 7, 30, 12, 0, 0),
+        end_time=datetime(2026, 7, 30, 12, 0, 1),
+    )
+    assert payload == original, "the caller's standard_logging_object must not be mutated"
+    (element,) = logger.log_queue
+
+    response = MagicMock()
+    response.status_code = 200
+    response.raise_for_status = MagicMock()
+    logger.async_httpx_client = AsyncMock()
+    logger.async_httpx_client.put.return_value = response
+    await logger.async_upload_data_to_s3(element)
+    return json.loads(logger.async_httpx_client.put.call_args.kwargs["data"])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("event_name", ["async_log_success_event", "async_log_failure_event"])
+async def test_prompts_only_drops_response_but_keeps_messages_in_uploaded_object(
+    monkeypatch: pytest.MonkeyPatch, event_name: str
+):
+    import litellm
+
+    monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": True})
+    logger = _prompts_only_logger()
+
+    log_event: Callable[..., Awaitable[None]] = (
+        logger.async_log_success_event if event_name == "async_log_success_event" else logger.async_log_failure_event
+    )
+    body = await _queued_body_via_async_upload(logger, log_event)
+
+    assert body["messages"] == _chat_payload()["messages"]
+    assert body["response"] is None
+    assert body["id"] == "chatcmpl-prompts-only"
+
+
+@pytest.mark.asyncio
+async def test_prompts_only_default_off_keeps_response_in_uploaded_object(monkeypatch: pytest.MonkeyPatch):
+    import litellm
+
+    monkeypatch.setattr(litellm, "s3_callback_params", {})
+    monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False)
+    logger = _prompts_only_logger()
+
+    body = await _queued_body_via_async_upload(logger, logger.async_log_success_event)
+
+    assert body["response"] == _chat_payload()["response"]
+    assert body["messages"] == _chat_payload()["messages"]
+
+
+@pytest.mark.asyncio
+async def test_prompts_only_explicit_false_in_params_beats_env_var(monkeypatch: pytest.MonkeyPatch):
+    import litellm
+
+    monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": False})
+    monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true")
+    logger = _prompts_only_logger()
+
+    body = await _queued_body_via_async_upload(logger, logger.async_log_success_event)
+
+    assert body["response"] == _chat_payload()["response"]
+
+
+@pytest.mark.asyncio
+async def test_prompts_only_env_var_applies_when_param_unset(monkeypatch: pytest.MonkeyPatch):
+    import litellm
+
+    monkeypatch.setattr(litellm, "s3_callback_params", {})
+    logger = _prompts_only_logger()
+    monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true")
+
+    body = await _queued_body_via_async_upload(logger, logger.async_log_success_event)
+
+    assert body["response"] is None
+    assert body["messages"] == _chat_payload()["messages"]
+
+
+@respx.mock
+def test_prompts_only_constructor_kwarg_applies_to_sync_upload(monkeypatch: pytest.MonkeyPatch):
+    import litellm
+
+    monkeypatch.setattr(litellm, "s3_callback_params", {})
+    monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False)
+    logger = _prompts_only_logger(s3_log_prompts_only=True)
+    payload = _chat_payload()
+
+    element = logger.create_s3_batch_logging_element(
+        start_time=datetime(2026, 7, 30, 12, 0, 0),
+        standard_logging_payload=payload,
+    )
+    assert element is not None
+    assert payload["response"] == _chat_payload()["response"]
+
+    put_route = respx.put(url__regex=r"https://test-bucket\.s3\..*").mock(return_value=httpx.Response(200))
+    logger.upload_data_to_s3(element)
+
+    body = json.loads(put_route.calls.last.request.content)
+    assert body["response"] is None
+    assert body["messages"] == _chat_payload()["messages"]
+
+
+@pytest.mark.parametrize("callback_name", ["s3", "s3_v2"])
+def test_prompts_only_toggle_is_exposed_to_admin_ui_for_both_s3_callbacks(callback_name: str):
+    from litellm.integrations.custom_logger import CustomLogger
+
+    assert "S3_LOG_PROMPTS_ONLY" in CustomLogger.get_callback_env_vars(callback_name)

tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py

@@ -2482,7 +2482,10 @@ class TestAnthropicMessagesHandlerStreamingScanKey:
         open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use])
         ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")])
         assert open_key == StreamingScanKey(texts=("hi",))
+        assert open_key.tool_calls_in_flight is True
+        assert handler.get_streaming_scan_key([self._text_delta("hi")]).tool_calls_in_flight is False
         assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
+        assert ended_key.tool_calls_in_flight is False
         assert ended_key != open_key
 
 

tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py

@@ -2719,3 +2719,74 @@ class TestRustChatCompletionsHook:
             "model": "m",
             "messages": [],
         }
+
+
+def _served_model_stream_chunks(model: str | None) -> list[dict[str, object]]:
+    return [
+        {
+            "type": "message_start",
+            "message": {
+                "id": "msg_served",
+                "type": "message",
+                "role": "assistant",
+                "content": [],
+                "usage": {"input_tokens": 10, "output_tokens": 1},
+                **({"model": model} if model is not None else {}),
+            },
+        },
+        {
+            "type": "content_block_start",
+            "index": 0,
+            "content_block": {"type": "text", "text": ""},
+        },
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "text_delta", "text": "Hello"},
+        },
+        {"type": "content_block_stop", "index": 0},
+        {
+            "type": "message_delta",
+            "delta": {"stop_reason": "end_turn"},
+            "usage": {"output_tokens": 2},
+        },
+        {"type": "message_stop"},
+    ]
+
+
+def test_message_start_model_is_carried_on_stream_chunks():
+    iterator: Final = ModelResponseIterator(None, sync_stream=True)
+
+    parsed: Final = [iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks("claude-served-1")]
+
+    assert all(chunk.model == "claude-served-1" for chunk in parsed)
+
+
+def test_message_start_without_model_leaves_chunk_model_unset():
+    iterator: Final = ModelResponseIterator(None, sync_stream=True)
+
+    parsed: Final = [iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks(None)]
+
+    assert all(chunk.model is None for chunk in parsed)
+
+
+def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper():
+    from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+
+    served_model: Final = "claude-served-1"
+    sse_lines: Final = [f"data: {json.dumps(chunk)}\n".encode() for chunk in _served_model_stream_chunks(served_model)]
+    iterator: Final = ModelResponseIterator(iter(sse_lines), sync_stream=True)
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=iter(iterator),
+        model="anthropic/claude-requested",
+        custom_llm_provider="anthropic",
+        logging_obj=MagicMock(),
+    )
+
+    chunks: Final = list(wrapper)
+
+    assert len(chunks) > 1
+    for chunk in chunks[1:]:
+        assert chunk._hidden_params["provider_response_model"] == served_model
+    assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=[{"role": "user", "content": "hi"}])
+    assert assembled._hidden_params["provider_response_model"] == served_model

tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py

@@ -1200,6 +1200,7 @@ def _fake_user_api_key_auth(
     team_models=None,
     team_id=None,
     model_max_budget=None,
+    team_model_max_budget=None,
     end_user_model_max_budget=None,
     end_user_id=None,
     user_model_max_budget=None,
@@ -1220,6 +1221,7 @@ def _fake_user_api_key_auth(
     auth.team_id = team_id
     auth.team_model_aliases = None
     auth.model_max_budget = model_max_budget
+    auth.team_model_max_budget = team_model_max_budget
     auth.end_user_model_max_budget = end_user_model_max_budget
     auth.end_user_id = end_user_id
     auth.user_model_max_budget = user_model_max_budget
@@ -1860,6 +1862,78 @@ async def test_summary_model_rate_limit_skipped_for_legacy_limiter():
     assert not result.applied_edits[0].get("error")
 
 
+async def test_summary_model_denied_when_team_over_model_budget():
+    """The team per-model budget gates the summary subrequest, whose spend is
+    charged to the team counter via the propagated `user_api_key_team_model_max_budget`.
+    The key's own `model_max_budget` is handed to the limiter so a key-level
+    override keeps taking precedence over the team cap here as it does in auth."""
+    import litellm
+
+    messages = _simple_messages()
+    mock_call = AsyncMock(return_value=_make_mock_response("<summary>x</summary>"))
+    key_budget = {"claude-opus-4-8": {"budget_limit": 1}}
+    team_budget = {"claude-haiku-4-5": {"budget_limit": 5, "time_period": "1d"}}
+
+    auth = _fake_user_api_key_auth(
+        key_models=["all-proxy-models"],
+        model_max_budget=key_budget,
+        team_model_max_budget=team_budget,
+        team_id="team-over-budget",
+        token="hashed-token",
+    )
+
+    limiter = MagicMock()
+    limiter.is_key_within_model_budget = AsyncMock(return_value=True)
+    limiter.is_team_within_model_budget = AsyncMock(
+        side_effect=litellm.BudgetExceededError(
+            message="over budget", current_cost=10, max_budget=5
+        )
+    )
+
+    with (
+        patch(  # test-quality-ok: apply_compact_20260112 reads the summary model setting as a module global, no seam
+            "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+            return_value="claude-haiku-4-5",
+        ),
+        patch("litellm.token_counter", return_value=200_000),  # test-quality-ok: forces the over-threshold branch
+        patch(  # test-quality-ok: the summary call is the observable that must NOT happen when the team is over budget
+            "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+            mock_call,
+        ),
+        patch(  # test-quality-ok: the limiter is a proxy_server module global the editor imports, no injection seam
+            "litellm.proxy.proxy_server.model_max_budget_limiter", limiter
+        ),
+    ):
+        result = await apply_compact_20260112(
+            model=MODEL,
+            messages=messages,
+            tools=None,
+            system=None,
+            edit_spec=_EDIT_SPEC_DEFAULT,
+            user_api_key_auth=auth,
+        )
+
+    mock_call.assert_not_awaited()
+    assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded"
+    limiter.is_team_within_model_budget.assert_awaited_once_with(
+        team_id="team-over-budget",
+        team_model_max_budget=team_budget,
+        key_model_max_budget=key_budget,
+        model="claude-haiku-4-5",
+    )
+    import inspect
+
+    from litellm.proxy.hooks.model_max_budget_limiter import (
+        _PROXY_VirtualKeyModelMaxBudgetLimiter,
+    )
+
+    real_params = inspect.signature(
+        _PROXY_VirtualKeyModelMaxBudgetLimiter.is_team_within_model_budget
+    ).parameters
+    for kwarg in ("team_id", "team_model_max_budget", "key_model_max_budget", "model"):
+        assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter does not accept"
+
+
 async def test_scoped_budget_metadata_propagated_to_summary_call():
     """The end-user/project scope identifiers and the end-user budget the post-call
     spend and rate-limit hooks key on are forwarded to the summary subrequest, and

tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py

@@ -5,12 +5,15 @@ from typing import NamedTuple
 import pytest
 
 import litellm
+from litellm.cost_calculator import completion_cost
 from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
 from litellm.llms.bedrock.common_utils import BedrockModelInfo
 from litellm.types.utils import (
     Choices,
     Message,
     ModelResponse,
+    PromptTokensDetailsWrapper,
+    Usage,
 )
 
 
@@ -152,3 +155,55 @@ def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(prof
     assert "reasoning_effort" in supported
     assert "thinking" not in supported
     assert "output_config" not in supported
+
+
+# Cache-read prices are the `*-cache-read-input-tokens` usagetype rows of the AWS Price List API, us-east-1,
+# https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json on 2026-09-15
+@pytest.mark.parametrize(
+    "model,expected_cache_read",
+    [
+        ("amazon.nova-lite-v1:0", 1.5e-8),
+        ("us.amazon.nova-lite-v1:0", 1.5e-8),
+        ("amazon.nova-micro-v1:0", 8.75e-9),
+        ("us.amazon.nova-micro-v1:0", 8.75e-9),
+        ("amazon.nova-pro-v1:0", 2e-7),
+        ("us.amazon.nova-pro-v1:0", 2e-7),
+        ("us.amazon.nova-premier-v1:0", 6.25e-7),
+    ],
+)
+def test_bedrock_nova_cache_read_prices(
+    model, expected_cache_read, local_model_cost_map
+):
+    model_info = litellm.model_cost[model]
+    assert model_info["cache_read_input_token_cost"] == expected_cache_read
+    usage = Usage(
+        prompt_tokens=1_000,
+        completion_tokens=100,
+        total_tokens=1_100,
+        prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=400),
+    )
+    response = _bedrock_response(model, usage)
+
+    cost = completion_cost(
+        completion_response=response,
+        model=model,
+        custom_llm_provider="bedrock",
+    )
+    expected_cost = (
+        600 * model_info["input_cost_per_token"]
+        + 400 * expected_cache_read
+        + 100 * model_info["output_cost_per_token"]
+    )
+    assert cost == pytest.approx(expected_cost)
+
+    uncached_usage = Usage(
+        prompt_tokens=1_000,
+        completion_tokens=100,
+        total_tokens=1_100,
+    )
+    uncached_cost = completion_cost(
+        completion_response=_bedrock_response(model, uncached_usage),
+        model=model,
+        custom_llm_provider="bedrock",
+    )
+    assert cost < uncached_cost

tests/test_litellm/llms/gemini/test_cost_calculator.py

@@ -1,8 +1,7 @@
-import os
-
 import pytest
 
 import litellm
+from litellm.cost_calculator import completion_cost
 from litellm.llms.gemini.cost_calculator import (
     cost_per_google_maps_grounding_request,
     cost_per_web_search_request,
@@ -18,6 +17,7 @@ from litellm.types.utils import (
     ImageResponse,
     ImageUsage,
     ImageUsageInputTokensDetails,
+    ModelResponse,
     PromptTokensDetailsWrapper,
     Usage,
 )
@@ -452,6 +452,42 @@ def test_map_traffic_type_to_service_tier(
     )
 
 
+# Alias targets are the `modelVersion` returned by
+# POST https://generativelanguage.googleapis.com/v1beta/models/<alias>:generateContent on 2026-09-15
+@pytest.mark.parametrize(
+    "alias,target",
+    [
+        ("gemini/gemini-flash-latest", "gemini/gemini-3.8-flash"),
+        ("gemini/gemini-flash-lite-latest", "gemini/gemini-3.5-flash-lite"),
+        ("gemini/gemini-pro-latest", "gemini/gemini-3.1-pro-preview"),
+    ],
+)
+def test_latest_aliases_cost_the_same_as_their_current_target(
+    monkeypatch, alias, target
+):
+    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
+    monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
+
+    usage = Usage(
+        prompt_tokens=1_000,
+        completion_tokens=500,
+        total_tokens=1_500,
+        prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=400),
+    )
+
+    def cost_of(model: str) -> float:
+        return completion_cost(
+            completion_response=ModelResponse(model=model, usage=usage),
+            model=model,
+            custom_llm_provider="gemini",
+        )
+
+    alias_cost = cost_of(alias)
+    target_cost = cost_of(target)
+    assert alias_cost == pytest.approx(target_cost)
+    assert alias_cost > 0
+
+
 @pytest.mark.parametrize(
     "prefixed,bare",
     [

tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py

@@ -2206,10 +2206,35 @@ class TestStreamingScanKey:
             [self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")]
         )
         assert open_key == StreamingScanKey(texts=("hi",))
+        assert open_key.tool_calls_in_flight is True
+        assert handler.get_streaming_scan_key([self._chunk("hi")]).tool_calls_in_flight is False
         assert ended_key.texts == ("hi",)
         assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
+        assert ended_key.tool_calls_in_flight is False
         assert ended_key != open_key
 
+    def test_legacy_function_call_delta_is_held_like_a_tool_call(self):
+        from litellm.types.utils import Delta, FunctionCall, ModelResponseStream, StreamingChoices
+
+        handler = OpenAIChatCompletionsHandler()
+        function_chunk = ModelResponseStream(
+            choices=[
+                StreamingChoices(
+                    index=0,
+                    delta=Delta(content=None, function_call=FunctionCall(name="run_shell", arguments='{"cmd": "rm"}')),
+                    finish_reason=None,
+                )
+            ]
+        )
+        open_key = handler.get_streaming_scan_key([self._chunk("hi"), function_chunk])
+        ended_key = handler.get_streaming_scan_key(
+            [self._chunk("hi"), function_chunk, self._chunk(None, finish_reason="function_call")]
+        )
+        assert open_key.tool_calls_in_flight is True
+        assert open_key.tool_calls == ()
+        assert len(ended_key.tool_calls) == 1 and "run_shell" in ended_key.tool_calls[0]
+        assert ended_key.tool_calls_in_flight is False
+
     def test_text_after_the_first_choice_finishes_still_changes_the_key(self):
         handler = OpenAIChatCompletionsHandler()
         first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)]

tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py

@@ -3211,3 +3211,42 @@ class TestOpenAIResponsesHandlerStreamingScanKey:
     def test_output_item_done_round_is_never_deduped(self):
         done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}}
         assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None
+
+    @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"])
+    def test_non_completed_terminal_envelopes_key_their_output_items(self, terminal_type):
+        handler = OpenAIResponsesHandler()
+        arguments_delta = {
+            "type": "response.function_call_arguments.delta",
+            "sequence_number": 1,
+            "item_id": "fc_1",
+            "delta": '{"city":',
+        }
+        function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city":'}
+        terminal = {"type": terminal_type, "sequence_number": 2, "response": {"id": "resp_1", "output": [function_call]}}
+        mid_stream_key = handler.get_streaming_scan_key([arguments_delta])
+        ended_key = handler.get_streaming_scan_key([arguments_delta, terminal])
+        assert ended_key.stream_ended is True
+        assert ended_key.tool_calls_in_flight is False
+        assert len(ended_key.tool_calls) == 1
+        assert ended_key != mid_stream_key
+
+    def test_streamed_tool_call_events_flag_tool_calls_in_flight_until_the_stream_ends(self):
+        handler = OpenAIResponsesHandler()
+        added = {
+            "type": "response.output_item.added",
+            "sequence_number": 1,
+            "item": {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "get_weather"},
+        }
+        arguments_delta = {
+            "type": "response.function_call_arguments.delta",
+            "sequence_number": 2,
+            "item_id": "fc_1",
+            "delta": '{"city":',
+        }
+        function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"}
+        assert handler.get_streaming_scan_key([self._delta(0, "hi")]).tool_calls_in_flight is False
+        assert handler.get_streaming_scan_key([self._delta(0, "hi"), added]).tool_calls_in_flight is True
+        assert handler.get_streaming_scan_key([self._delta(0, "hi"), arguments_delta]).tool_calls_in_flight is True
+        ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), added, self._completed(3, [function_call])])
+        assert ended_key.tool_calls_in_flight is False
+        assert len(ended_key.tool_calls) == 1

tests/test_litellm/models/test_models.py

@@ -362,6 +362,14 @@ class TestVerificationToken:
         assert deleted.deleted_at is not None
         assert deleted.token == "t1"
 
+    def test_total_spend_is_carried_separately_from_resettable_spend(self):
+        token = LiteLLM_VerificationToken(token="t1", spend=0.0, total_spend=12.5)
+        assert token.model_dump()["total_spend"] == 12.5
+        assert token.model_dump()["spend"] == 0.0
+
+        deleted = LiteLLM_DeletedVerificationToken.model_validate({**token.model_dump(), "deleted_by": "admin"})
+        assert deleted.total_spend == 12.5
+
 
 class TestConfigTable:
     def test_config_creation(self):

tests/test_litellm/ocr/test_ocr_file_input.py

@@ -12,32 +12,16 @@ Tests that:
 import base64
 import os
 import tempfile
-from collections.abc import Generator
 from io import BytesIO
 from pathlib import Path
 from typing import Final
-from unittest.mock import AsyncMock, MagicMock, Mock
+from unittest.mock import AsyncMock, MagicMock
 
 import orjson
 import pytest
 from starlette.datastructures import FormData
 
-from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type
-
-
-@pytest.fixture(autouse=True, params=["native", "disabled", "unavailable"])
-def document_runtime(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
-    from litellm.rust_bridge import bindings, configuration
-
-    configuration.reset_rust_configuration()
-    monkeypatch.delenv("LITELLM_RUST", raising=False)
-    if request.param == "disabled":
-        monkeypatch.setenv("LITELLM_RUST", "0")
-        monkeypatch.setattr(bindings, "get_native_bridge", Mock(side_effect=AssertionError("Rust is disabled")))
-    elif request.param == "unavailable":
-        monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
-    yield
-    configuration.reset_rust_configuration()
+from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type
 
 
 class TestGetMimeType:
@@ -503,10 +487,9 @@ class TestProxySecurityGuard:
 async def test_proxy_upload_stops_reading_at_size_limit() -> None:
     from starlette.datastructures import UploadFile
 
-    from litellm.ocr.input import get_max_file_bytes
-    from litellm.proxy.ocr_endpoints.endpoints import _parse_multipart_form
+    from litellm.proxy.ocr_endpoints.endpoints import _MAX_FILE_BYTES, _parse_multipart_form
 
-    limit: Final = get_max_file_bytes()
+    limit: Final = _MAX_FILE_BYTES
     with tempfile.TemporaryFile() as stream:
         stream.truncate(limit * 2)
         upload: Final = UploadFile(file=stream, filename="large.pdf")

tests/test_litellm/proxy/auth/test_fallback_budget.py

@@ -0,0 +1,202 @@
+import pytest
+
+from litellm import Router
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.auth.fallback_budget import (
+    RouterFallbackBudgetCheck,
+    is_token_within_budget_for_model,
+    router_fallback_budget_check,
+)
+
+FREE_MODEL = {
+    "model_name": "free-model",
+    "litellm_params": {
+        "model": "ollama/llama2",
+        "api_base": "http://localhost:11434",
+        "input_cost_per_token": 0.0,
+        "output_cost_per_token": 0.0,
+    },
+    "model_info": {
+        "id": "free-model-id",
+        "input_cost_per_token": 0.0,
+        "output_cost_per_token": 0.0,
+    },
+}
+
+PAID_MODEL = {
+    "model_name": "paid-model",
+    "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"},
+    "model_info": {"id": "paid-model-id"},
+}
+
+
+def _router() -> Router:
+    return Router(model_list=[FREE_MODEL, PAID_MODEL], fallbacks=[{"free-model": ["paid-model"]}])
+
+
+def _token(**overrides) -> UserAPIKeyAuth:
+    fields = {
+        "api_key": "hashed",
+        "token": "hashed",
+        "spend": 0.0,
+        "max_budget": None,
+        "user_id": "u1",
+        "user_spend": 0.0,
+        "user_max_budget": None,
+    }
+    fields.update(overrides)
+    return UserAPIKeyAuth(**fields)
+
+
+ENFORCED = RouterFallbackBudgetCheck(is_enforced=lambda: True)
+NOT_ENFORCED = RouterFallbackBudgetCheck(is_enforced=lambda: False)
+
+
+@pytest.mark.asyncio
+async def test_paid_target_allowed_when_under_budget():
+    token = _token(spend=1.0, max_budget=50.0, user_spend=1.0, user_max_budget=50.0)
+    assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True
+
+
+@pytest.mark.asyncio
+async def test_paid_target_refused_when_over_key_budget():
+    token = _token(spend=100.0, max_budget=50.0)
+    assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False
+
+
+@pytest.mark.asyncio
+async def test_paid_target_refused_when_over_user_budget():
+    token = _token(user_spend=1900.0, user_max_budget=50.0)
+    assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False
+
+
+@pytest.mark.asyncio
+async def test_zero_cost_target_allowed_even_when_over_budget():
+    """Refusing a free target would deny a request on spend some other model accrued."""
+    token = _token(user_spend=1900.0, user_max_budget=50.0)
+    assert await is_token_within_budget_for_model(model="free-model", valid_token=token, llm_router=_router()) is True
+
+
+@pytest.mark.asyncio
+async def test_no_budget_configured_is_always_within_budget():
+    token = _token(spend=9999.0, user_spend=9999.0)
+    assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True
+
+
+@pytest.mark.asyncio
+async def test_team_key_does_not_inherit_personal_budget_by_default(monkeypatch):
+    """Mirrors _PROXY_MaxBudgetLimiter: a team key ignores the owner's personal cap."""
+    from litellm.proxy import proxy_server
+
+    monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False)
+    token = _token(team_id="t1", user_spend=1900.0, user_max_budget=50.0)
+    assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True
+
+
+@pytest.mark.asyncio
+async def test_team_key_inherits_personal_budget_when_opted_in(monkeypatch):
+    from litellm.proxy import proxy_server
+
+    monkeypatch.setattr(proxy_server, "general_settings", {"apply_user_budget_to_team_keys": True}, raising=False)
+    token = _token(team_id="t1", user_spend=1900.0, user_max_budget=50.0)
+    assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False
+
+
+@pytest.mark.asyncio
+async def test_check_is_a_no_op_while_not_enforced():
+    request = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}}
+    assert await NOT_ENFORCED(model="paid-model", request_kwargs=request, llm_router=_router()) is True
+
+
+@pytest.mark.asyncio
+async def test_request_without_a_key_is_unrestricted():
+    assert await ENFORCED(model="paid-model", request_kwargs={}, llm_router=_router()) is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"])
+async def test_enforced_check_reads_the_key_from_request_metadata(metadata_field: str):
+    over = {metadata_field: {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}}
+    under = {metadata_field: {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}}
+
+    assert await ENFORCED(model="paid-model", request_kwargs=over, llm_router=_router()) is False
+    assert await ENFORCED(model="paid-model", request_kwargs=under, llm_router=_router()) is True
+
+
+@pytest.mark.asyncio
+async def test_a_stale_low_counter_still_refuses_a_paid_target(monkeypatch):
+    """
+    The counter can read low (e.g. restored from an older Redis snapshot). Passing the budget makes
+    `get_current_spend` verify against authoritative spend instead of trusting that read, so the
+    paid target is still refused.
+    """
+    from litellm.proxy import proxy_server
+
+    seen: list[dict] = []
+
+    async def _stale_counter(**kwargs):
+        seen.append(kwargs)
+        # a stale-low counter read; the authoritative spend is what the budget must be judged on
+        return 0.0 if kwargs.get("max_budget") is None else kwargs["fallback_spend"]
+
+    monkeypatch.setattr(proxy_server, "get_current_spend", _stale_counter, raising=False)
+    token = _token(user_spend=1900.0, user_max_budget=50.0)
+
+    assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False
+    assert [call["max_budget"] for call in seen] == [50.0]
+
+
+@pytest.mark.asyncio
+async def test_check_fails_closed_when_the_spend_lookup_breaks(monkeypatch):
+    from litellm.proxy import proxy_server
+
+    async def _boom(**kwargs):
+        raise RuntimeError("spend counter unavailable")
+
+    monkeypatch.setattr(proxy_server, "get_current_spend", _boom, raising=False)
+    request = {"metadata": {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}}
+
+    assert await ENFORCED(model="paid-model", request_kwargs=request, llm_router=_router()) is False
+
+
+@pytest.mark.asyncio
+async def test_router_skips_the_paid_fallback_target_when_over_budget():
+    from litellm.router_utils.fallback_event_handlers import _is_fallback_target_within_budget
+
+    router = Router(
+        model_list=[FREE_MODEL, PAID_MODEL],
+        fallbacks=[{"free-model": ["paid-model"]}],
+        fallback_budget_check=ENFORCED,
+    )
+    over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}}
+    under = {"metadata": {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}}
+
+    assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", over) is False
+    assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", under) is True
+
+
+@pytest.mark.asyncio
+async def test_router_without_a_budget_check_attempts_every_fallback():
+    from litellm.router_utils.fallback_event_handlers import _is_fallback_target_within_budget
+
+    router = _router()  # fallback_budget_check defaults to None
+    over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}}
+
+    assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", over) is True
+
+
+@pytest.mark.asyncio
+async def test_enforcement_is_on_by_default_and_opt_out_restores_the_leak(monkeypatch):
+    """
+    Leaving the paid fallback unguarded is the budget bypass this module exists to close, so an
+    unconfigured proxy has to enforce. `enforce_fallback_budget: false` is the deliberate opt-out.
+    """
+    from litellm.proxy import proxy_server
+
+    over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}}
+
+    monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False)
+    assert await router_fallback_budget_check(model="paid-model", request_kwargs=over, llm_router=_router()) is False
+
+    monkeypatch.setattr(proxy_server, "general_settings", {"enforce_fallback_budget": False}, raising=False)
+    assert await router_fallback_budget_check(model="paid-model", request_kwargs=over, llm_router=_router()) is True

tests/test_litellm/proxy/auth/test_team_grants.py

@@ -31,6 +31,7 @@ def _full_team(model_aliases=ALIASES) -> LiteLLM_TeamTable:
         max_budget=50.0,
         soft_budget=25.0,
         spend=12.5,
+        model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}},
         models=["gpt-4o", "gpt-4o-mini"],
         blocked=True,
         metadata={"tier": "gold"},
@@ -72,6 +73,7 @@ def test_team_grants_cover_every_team_field_the_key_path_gets():
     assert token.team_max_budget == 50.0
     assert token.team_soft_budget == 25.0
     assert token.team_spend == 12.5
+    assert token.team_model_max_budget == {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}
     assert token.team_models == ["gpt-4o", "gpt-4o-mini"]
     assert token.team_blocked is True
     assert token.team_metadata == {"tier": "gold"}

tests/test_litellm/proxy/auth/test_user_api_key_auth.py

@@ -4374,6 +4374,74 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t
     }
 
 
+class _RecordingTeamModelBudgetLimiter:
+    def __init__(self):
+        self.calls = []
+
+    async def is_team_within_model_budget(self, team_id, team_model_max_budget, key_model_max_budget, model):
+        self.calls.append((team_id, dict(team_model_max_budget), key_model_max_budget, model))
+        return True
+
+
+@pytest.mark.asyncio
+async def test_centralized_common_checks_enforces_team_model_max_budget_from_the_resolved_team():
+    """The team's model_max_budget is enforced at the single authz gate, off the
+    team object auth resolved (not the possibly stale token copy), and the key's
+    own model_max_budget is handed to the limiter so a matching key entry can
+    override the team cap."""
+    from fastapi import Request
+    from starlette.datastructures import URL
+
+    import litellm.proxy.proxy_server as _proxy_server_mod
+
+    team_caps = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}
+    key_caps = {"claude-sonnet-4-6": {"max_budget": 1.0, "budget_duration": "1d"}}
+    token = UserAPIKeyAuth(
+        api_key="sk-test",
+        token="hashed",
+        team_id="t1",
+        team_model_max_budget={"gpt-4o": {"max_budget": 999.0, "budget_duration": "30d"}},
+        model_max_budget=key_caps,
+    )
+    request = Request(scope={"type": "http"})
+    request._url = URL(url="/chat/completions")
+
+    user_api_key_cache = DualCache()
+    await user_api_key_cache.async_set_cache(
+        key="team_id:t1",
+        value=LiteLLM_TeamTableCachedObj(team_id="t1", model_max_budget=team_caps),
+    )
+    limiter = _RecordingTeamModelBudgetLimiter()
+    attrs = {
+        **_proxy_attrs_for_centralized_checks(user_custom_auth=None),
+        "prisma_client": MagicMock(),
+        "user_api_key_cache": user_api_key_cache,
+        "model_max_budget_limiter": limiter,
+    }
+    originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
+    try:
+        for k, v in attrs.items():
+            setattr(_proxy_server_mod, k, v)
+        with (
+            patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock),  # test-quality-ok: stubs the sibling check so only the team model-budget gate is under test
+            patch(  # test-quality-ok: stubs the budget reservation so only the team model-budget gate is under test
+                "litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks",
+                new_callable=AsyncMock,
+            ),
+        ):
+            await _run_centralized_common_checks(
+                user_api_key_auth_obj=token,
+                request=request,
+                request_data={"model": "gpt-4o"},
+                route="/chat/completions",
+            )
+    finally:
+        for k, v in originals.items():
+            setattr(_proxy_server_mod, k, v)
+
+    assert limiter.calls == [("t1", team_caps, key_caps, "gpt-4o")]
+
+
 @pytest.mark.asyncio
 async def test_centralized_common_checks_skipped_for_custom_auth_without_flag():
     """Existing RPS guarantee: custom-auth deployments without

tests/test_litellm/proxy/common_utils/test_reset_budget_job.py

@@ -291,6 +291,23 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client):
     assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
 
 
+def test_reset_budget_for_key_leaves_lifetime_total_spend_alone(reset_budget_job, mock_prisma_client):
+    """A period reset zeroes spend but must neither write nor touch the lifetime total_spend."""
+    now = datetime.now(timezone.utc)
+    key = LiteLLM_VerificationToken(
+        token="tok-key-1", spend=100.0, total_spend=340.0, budget_duration="30d", budget_reset_at=now
+    )
+    mock_prisma_client.data["key"] = [key]
+
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
+
+    (write,) = _batch_writes(mock_prisma_client, "key")
+    assert write["data"]["spend"] == {"decrement": 100.0}
+    assert "total_spend" not in write["data"]
+    assert key.spend == 0.0
+    assert key.total_spend == 340.0
+
+
 def test_reset_budget_for_key_honors_injected_reset_time(mock_prisma_client, mock_proxy_logging):
     """Injected BudgetResetSettings drives the written reset time end to end (DI, no globals).
 

tests/test_litellm/proxy/db/test_create_views.py

@@ -71,6 +71,7 @@ async def test_create_views_creates_view_on_does_not_exist():
     mock_db.execute_raw.assert_called_once()
     created_sql = mock_db.execute_raw.call_args[0][0]
     assert 'CREATE VIEW "LiteLLM_VerificationTokenView"' in created_sql
+    assert "t.model_max_budget AS team_model_max_budget" in created_sql
 
 
 @pytest.mark.asyncio

tests/test_litellm/proxy/db/test_db_spend_update_writer.py

@@ -1658,6 +1658,57 @@ async def test_commit_key_spend_updates_includes_last_active():
     assert before_call <= last_active <= after_call
 
 
+@pytest.mark.asyncio
+async def test_commit_spend_updates_to_db_increments_key_total_spend_alongside_spend():
+    """
+    The key table write must increment the lifetime total_spend by the same amount as the
+    resettable spend, in the same update so the two cannot drift.
+    """
+    db_writer = DBSpendUpdateWriter()
+
+    mock_batcher = MagicMock()
+    mock_batcher.litellm_verificationtoken = MagicMock()
+    mock_batcher.litellm_verificationtoken.update_many = MagicMock()
+
+    mock_transaction = AsyncMock()
+    mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
+    mock_transaction.__aexit__ = AsyncMock(return_value=False)
+    mock_transaction.batch_ = MagicMock(
+        return_value=AsyncMock(
+            __aenter__=AsyncMock(return_value=mock_batcher),
+            __aexit__=AsyncMock(return_value=False),
+        )
+    )
+
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db = MagicMock()
+    mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
+
+    db_spend_update_transactions = {
+        "user_list_transactions": {},
+        "end_user_list_transactions": {},
+        "key_list_transactions": {"hashed_token_abc": 0.05, "hashed_token_def": 1.25},
+        "team_list_transactions": {},
+        "team_member_list_transactions": {},
+        "org_list_transactions": {},
+        "tag_list_transactions": {},
+        "agent_list_transactions": {},
+    }
+
+    await db_writer._commit_spend_updates_to_db(
+        prisma_client=mock_prisma_client,
+        n_retry_times=0,
+        proxy_logging_obj=MagicMock(),
+        db_spend_update_transactions=db_spend_update_transactions,
+    )
+
+    calls = mock_batcher.litellm_verificationtoken.update_many.call_args_list
+    assert [c.kwargs["where"] for c in calls] == [{"token": "hashed_token_abc"}, {"token": "hashed_token_def"}]
+    for call, expected_cost in zip(calls, (0.05, 1.25)):
+        assert call.kwargs["data"]["spend"] == {"increment": expected_cost}
+        assert call.kwargs["data"]["total_spend"] == call.kwargs["data"]["spend"]
+
+
 @pytest.mark.asyncio
 async def test_update_database_creates_single_task():
     """
@@ -2813,7 +2864,7 @@ async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at
     mock_batcher.litellm_verificationtoken.update_many.assert_called_once()
     call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1]
     assert call_kwargs["where"] == {"token": token}
-    assert set(call_kwargs["data"]) == {"spend", "last_active"}
+    assert set(call_kwargs["data"]) == {"spend", "total_spend", "last_active"}
     assert call_kwargs["data"]["spend"] == {"increment": response_cost}
 
 

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py

@@ -5603,6 +5603,7 @@ def test_initialize_bedrock_wires_streaming_flags():
             streaming_buffer_until_moderated=False,
             streaming_sampling_rate=3,
             streaming_end_of_stream_only=True,
+            streaming_buffer_release_on_scan=True,
         ),
         {"guardrail_name": "bedrock-streaming"},
     )
@@ -5616,9 +5617,11 @@ def test_initialize_bedrock_wires_streaming_flags():
     assert configured.streaming_buffer_until_moderated is False
     assert configured.streaming_sampling_rate == 3
     assert configured.streaming_end_of_stream_only is True
+    assert configured.streaming_buffer_release_on_scan is True
     assert defaulted.streaming_buffer_until_moderated is True
     assert defaulted.streaming_sampling_rate == 5
     assert defaulted.streaming_end_of_stream_only is False
+    assert defaulted.streaming_buffer_release_on_scan is False
 
 
 def test_initialize_bedrock_rejects_non_positive_sampling_rate():
@@ -5721,6 +5724,44 @@ async def test_buffered_default_hook_scans_before_any_chunk():
     assert len([e for e in events if e != "scan"]) >= 1
 
 
+@pytest.mark.asyncio
+async def test_buffered_release_on_scan_hook_releases_each_window_after_its_scan():
+    guardrail = BedrockGuardrail(
+        guardrail_name="bedrock-release-on-scan",
+        guardrailIdentifier="test-id",
+        guardrailVersion="DRAFT",
+        event_hook=GuardrailEventHooks.post_call,
+        default_on=True,
+        streaming_buffer_release_on_scan=True,
+        streaming_sampling_rate=1,
+    )
+
+    assert guardrail._streams_incrementally() is True
+    events = await _run_streaming_hook_recording_order(guardrail)
+
+    assert events == ["scan", ("chunk", "Hello"), "scan", ("chunk", " world"), ("chunk", "")]
+
+
+@pytest.mark.asyncio
+async def test_buffered_release_on_scan_defers_to_end_of_stream_only():
+    guardrail = BedrockGuardrail(
+        guardrail_name="bedrock-release-on-scan-end-only",
+        guardrailIdentifier="test-id",
+        guardrailVersion="DRAFT",
+        event_hook=GuardrailEventHooks.post_call,
+        default_on=True,
+        streaming_buffer_release_on_scan=True,
+        streaming_end_of_stream_only=True,
+        streaming_sampling_rate=1,
+    )
+
+    assert guardrail._streams_incrementally() is False
+    events = await _run_streaming_hook_recording_order(guardrail)
+
+    assert events.count("scan") == 1
+    assert events[0] == "scan"
+
+
 @pytest.mark.asyncio
 async def test_masking_keeps_buffered_path_even_when_unbuffered_configured():
     guardrail = BedrockGuardrail(

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py

@@ -1622,10 +1622,23 @@ def test_initialize_guardrail_rejects_unsupported_mode_instead_of_running_other_
 def test_initialize_guardrail_defaults_streaming_params() -> None:
     handler = _initialize_from_config(mode="post_call")
 
+    assert handler.streaming_buffer_until_moderated is False
+    assert handler.streaming_buffer_release_on_scan is False
     assert handler.streaming_end_of_stream_only is False
     assert handler.streaming_sampling_rate == 5
 
 
+def test_initialize_guardrail_forwards_buffer_streaming_params() -> None:
+    handler = _initialize_from_config(
+        mode="post_call",
+        streaming_buffer_until_moderated=True,
+        streaming_buffer_release_on_scan=True,
+    )
+
+    assert handler.streaming_buffer_until_moderated is True
+    assert handler.streaming_buffer_release_on_scan is True
+
+
 @pytest.mark.parametrize(
     "configured",
     [

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py

@@ -1291,8 +1291,9 @@ class TestToolPermissionGuardrailAnthropicMessages:
     async def test_rewrite_mode_keeps_the_stream_identity_it_had_before_the_shared_helper(self):
         """Well-formed SSE must round-trip exactly as it did before the helpers were shared.
 
-        The shared module can stamp the upstream message id and model onto the assembled response
-        for callers that ask for it; this path never did, and a client reads those bytes.
+        The shared module can stamp the upstream message id onto the assembled response for
+        callers that ask for it; this path never did, and a client reads those bytes. The model,
+        though, is now the upstream's, matching what the untouched passthrough shows clients.
         """
         with patch.object(self.rewriting, "should_run_guardrail", return_value=True):
             out = await self._drain(self.rewriting, self._sse_chunks("Read"))
@@ -1304,7 +1305,7 @@ class TestToolPermissionGuardrailAnthropicMessages:
             if line.startswith("data: ") and json.loads(line[6:]).get("type") == "message_start"
         )["message"]
         assert message_start["id"].startswith("chatcmpl-"), "the rewritten stream must not adopt the upstream message id"
-        assert message_start["model"] == "unknown-model", "the rewritten stream must not adopt the upstream model"
+        assert message_start["model"] == "claude-sonnet-4-5", "the rewritten stream reports the model the upstream served"
 
     @pytest.mark.asyncio
     async def test_message_start_without_a_dict_message_fails_closed(self):

tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py

@@ -11,7 +11,7 @@ released unchanged after moderation passes.
 """
 
 import json
-from typing import Any, List, Literal, Optional
+from typing import Any, AsyncGenerator, List, Literal, Optional
 
 import pytest
 
@@ -19,14 +19,25 @@ from litellm.integrations.custom_guardrail import (
     CustomGuardrail,
     ModifyResponseException,
 )
+from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
     UnifiedLLMGuardrails,
+    _is_redundant_scan,
+)
+from litellm.types.utils import (
+    ChatCompletionDeltaToolCall,
+    Delta,
+    Function,
+    FunctionCall,
+    GenericGuardrailAPIInputs,
+    ModelResponseStream,
+    StreamingChoices,
 )
-from litellm.types.utils import GenericGuardrailAPIInputs
 
 BLOCK_MESSAGE = "Blocked by policy: this response was withheld."
 ORIGINAL_MARKER = "ORIGINAL-SECRET-ANSWER"
+TOOL_ARGUMENTS_MARKER = "TOOL-ARGS-SECRET"
 
 
 class _BlockingGuardrail(CustomGuardrail):
@@ -60,6 +71,85 @@ class _PassingGuardrail(CustomGuardrail):
         return inputs
 
 
+class _CountingPassingGuardrail(_PassingGuardrail):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.scan_count = 0
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.scan_count += 1
+        return inputs
+
+
+class _ToolCallRecordingGuardrail(_CountingPassingGuardrail):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.tool_call_scan_indexes: List[int] = []
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.scan_count += 1
+        if inputs.get("tool_calls"):
+            self.tool_call_scan_indexes.append(self.scan_count)
+        return inputs
+
+
+class _SecondScanBlockingGuardrail(_CountingPassingGuardrail):
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.scan_count += 1
+        if self.scan_count == 2:
+            raise ModifyResponseException(
+                message=BLOCK_MESSAGE,
+                model="gpt-4",
+                request_data=request_data,
+                guardrail_name=self.guardrail_name,
+            )
+        return inputs
+
+
+class _MarkerBlockingGuardrail(_CountingPassingGuardrail):
+    """Blocks as soon as the inspected input field (texts or tool_calls) carries the marker."""
+
+    def __init__(self, *args, marker: str, field: Literal["texts", "tool_calls"] = "texts", **kwargs):
+        super().__init__(*args, **kwargs)
+        self.marker = marker
+        self.field = field
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.scan_count += 1
+        if self.marker in json.dumps(inputs.get(self.field, [])):
+            raise ModifyResponseException(
+                message=BLOCK_MESSAGE,
+                model="gpt-4o",
+                request_data=request_data,
+                guardrail_name=self.guardrail_name,
+            )
+        return inputs
+
+
 def _sse_event(event_type: str, data: dict) -> bytes:
     return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
 
@@ -115,6 +205,212 @@ def _decode(chunks: List[Any]) -> str:
     return "".join(c.decode() if isinstance(c, bytes) else str(c) for c in chunks)
 
 
+def _chat_chunk(content: str = "", finish_reason: str | None = None) -> ModelResponseStream:
+    return ModelResponseStream(
+        id="chatcmpl-windowed",
+        created=1724900000,
+        model="gpt-4",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(role="assistant", content=content),
+                finish_reason=finish_reason,
+            )
+        ],
+    )
+
+
+def _tool_call_chunk(
+    arguments: str, finish_reason: str | None = None, legacy_function_call: bool = False
+) -> ModelResponseStream:
+    delta = (
+        Delta(role="assistant", content=None, function_call=FunctionCall(name="run_shell", arguments=arguments))
+        if legacy_function_call
+        else Delta(
+            role="assistant",
+            content=None,
+            tool_calls=[
+                ChatCompletionDeltaToolCall(
+                    id="call_1",
+                    type="function",
+                    index=0,
+                    function=Function(name="run_shell", arguments=arguments),
+                )
+            ],
+        )
+    )
+    return ModelResponseStream(
+        id="chatcmpl-windowed",
+        created=1724900000,
+        model="gpt-4",
+        choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)],
+    )
+
+
+async def _windowed_chat_stream(
+    yielded_count: List[int],
+    collected: List[Any],
+    content_chunks: List[str],
+    tool_argument_chunks: List[str] | None = None,
+    legacy_function_call: bool = False,
+) -> AsyncGenerator[ModelResponseStream, None]:
+    for content in content_chunks:
+        yielded_count.append(len(collected))
+        yield _chat_chunk(content)
+    for arguments in tool_argument_chunks or []:
+        yielded_count.append(len(collected))
+        yield _tool_call_chunk(arguments, legacy_function_call=legacy_function_call)
+    yielded_count.append(len(collected))
+    yield _chat_chunk(finish_reason="tool_calls" if tool_argument_chunks else "stop")
+
+
+def _tool_argument_text(chunks: List[Any]) -> str:
+    return "".join(
+        tool_call.function.arguments or ""
+        for chunk in chunks
+        if isinstance(chunk, ModelResponseStream)
+        for choice in chunk.choices
+        for tool_call in choice.delta.tool_calls or []
+    )
+
+
+def _function_call_argument_text(chunks: list[Any]) -> str:
+    return "".join(
+        choice.delta.function_call.arguments or ""
+        for chunk in chunks
+        if isinstance(chunk, ModelResponseStream)
+        for choice in chunk.choices
+        if choice.delta.function_call is not None
+    )
+
+
+async def _run_windowed(
+    guardrail: CustomGuardrail,
+    content_chunks: List[str],
+    end_of_stream_only: bool = False,
+    tool_argument_chunks: List[str] | None = None,
+    legacy_function_call: bool = False,
+) -> tuple[List[Any], List[int]]:
+    guardrail.streaming_buffer_until_moderated = True
+    guardrail.streaming_buffer_release_on_scan = True
+    guardrail.streaming_end_of_stream_only = end_of_stream_only
+    guardrail.streaming_sampling_rate = 2
+    unified = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions")
+    request_data = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "guardrail_to_apply": guardrail,
+        "metadata": {"guardrails": [guardrail.guardrail_name]},
+    }
+    collected: List[Any] = []
+    yielded_count: List[int] = []
+    async for chunk in unified.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=user_api_key_dict,
+        response=_windowed_chat_stream(
+            yielded_count, collected, content_chunks, tool_argument_chunks, legacy_function_call
+        ),
+        request_data=request_data,
+    ):
+        collected.append(chunk)
+    return collected, yielded_count
+
+
+def _responses_message_stream_events(text_chunks: List[str]) -> List[dict]:
+    message = {"type": "message", "id": "msg_1", "status": "completed", "role": "assistant"}
+    content = [{"type": "output_text", "text": "".join(text_chunks), "annotations": []}]
+    return [
+        {"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}},
+        *(
+            {
+                "type": "response.output_text.delta",
+                "item_id": "msg_1",
+                "output_index": 0,
+                "content_index": 0,
+                "delta": text,
+            }
+            for text in text_chunks
+        ),
+        {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": content}},
+        {
+            "type": "response.completed",
+            "response": {
+                "id": "resp_1",
+                "model": "gpt-4o",
+                "status": "completed",
+                "output": [{**message, "content": content}],
+            },
+        },
+    ]
+
+
+def _responses_truncated_function_call_events(text: str, argument_chunks: List[str]) -> List[dict]:
+    message = {"type": "message", "id": "msg_1", "status": "completed", "role": "assistant"}
+    content = [{"type": "output_text", "text": text, "annotations": []}]
+    function_call = {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "run_shell"}
+    return [
+        {"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}},
+        {
+            "type": "response.output_text.delta",
+            "item_id": "msg_1",
+            "output_index": 0,
+            "content_index": 0,
+            "delta": text,
+        },
+        {"type": "response.output_item.added", "output_index": 1, "item": {**function_call, "arguments": ""}},
+        *(
+            {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": arguments}
+            for arguments in argument_chunks
+        ),
+        {
+            "type": "response.incomplete",
+            "response": {
+                "id": "resp_1",
+                "model": "gpt-4o",
+                "status": "incomplete",
+                "output": [
+                    {**message, "content": content},
+                    {**function_call, "arguments": "".join(argument_chunks), "status": "incomplete"},
+                ],
+            },
+        },
+    ]
+
+
+async def _replay(events: List[dict]) -> AsyncGenerator[dict, None]:
+    for event in events:
+        yield event
+
+
+async def _run_windowed_responses(guardrail: CustomGuardrail, events: List[dict]) -> str:
+    guardrail.streaming_buffer_until_moderated = True
+    guardrail.streaming_buffer_release_on_scan = True
+    guardrail.streaming_sampling_rate = 2
+    unified = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses")
+    request_data = {
+        "input": "hi",
+        "guardrail_to_apply": guardrail,
+        "metadata": {"guardrails": [guardrail.guardrail_name]},
+    }
+    collected: List[Any] = []
+    async for chunk in unified.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=user_api_key_dict,
+        response=_replay(events),
+        request_data=request_data,
+    ):
+        collected.append(chunk)
+    return json.dumps([chunk if isinstance(chunk, dict) else str(chunk) for chunk in collected])
+
+
+def _chat_text(chunks: List[Any]) -> str:
+    return "".join(
+        choice.delta.content or ""
+        for chunk in chunks
+        if isinstance(chunk, ModelResponseStream)
+        for choice in chunk.choices
+    )
+
+
 async def _run(guardrail: CustomGuardrail) -> str:
     # Rubrik's real config: end-of-stream-only moderation. Without buffering
     # this releases every chunk before moderation runs (content leaks on
@@ -159,6 +455,110 @@ async def test_buffered_clean_releases_all_content():
     assert BLOCK_MESSAGE not in raw
 
 
+@pytest.mark.asyncio
+async def test_windowed_buffer_releases_after_each_passing_scan():
+    guardrail = _CountingPassingGuardrail(guardrail_name="windowed-pass", event_hook="post_call")
+    content_chunks = ["one ", "two ", "three ", "four ", "five ", "six "]
+
+    collected, yielded_count = await _run_windowed(guardrail, content_chunks)
+
+    assert yielded_count[2] >= 2
+    assert yielded_count == [0, 0, 2, 2, 4, 4, 6]
+    assert _chat_text(collected) == "".join(content_chunks)
+    assert guardrail.scan_count > 1
+
+
+@pytest.mark.asyncio
+async def test_windowed_buffer_drops_blocked_window():
+    guardrail = _SecondScanBlockingGuardrail(guardrail_name="windowed-block", event_hook="post_call")
+    content_chunks = ["one ", "two ", "MARKER ", "four ", "five ", "six "]
+
+    collected, _ = await _run_windowed(guardrail, content_chunks)
+    raw = _decode(collected)
+
+    assert _chat_text(collected) == "one two "
+    assert "MARKER" not in raw
+    assert BLOCK_MESSAGE in raw
+    assert '"error"' not in raw
+
+
+@pytest.mark.asyncio
+async def test_windowed_buffer_holds_tool_call_windows_until_end_of_stream_scan():
+    guardrail = _ToolCallRecordingGuardrail(guardrail_name="windowed-tools", event_hook="post_call")
+    content_chunks = ["one ", "two ", "three "]
+    tool_argument_chunks = ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}']
+
+    collected, yielded_count = await _run_windowed(guardrail, content_chunks, tool_argument_chunks=tool_argument_chunks)
+
+    assert yielded_count == [0, 0, 2, 2, 2, 2, 2]
+    assert _chat_text(collected) == "".join(content_chunks)
+    assert _tool_argument_text(collected) == "".join(tool_argument_chunks)
+    assert guardrail.tool_call_scan_indexes == [guardrail.scan_count]
+
+
+@pytest.mark.asyncio
+async def test_windowed_buffer_holds_legacy_function_call_windows_until_end_of_stream():
+    guardrail = _PassingGuardrail(guardrail_name="windowed-functions", event_hook="post_call")
+    content_chunks = ["one ", "two ", "three "]
+    function_argument_chunks = ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}']
+
+    collected, yielded_count = await _run_windowed(
+        guardrail, content_chunks, tool_argument_chunks=function_argument_chunks, legacy_function_call=True
+    )
+
+    assert yielded_count == [0, 0, 2, 2, 2, 2, 2]
+    assert _chat_text(collected) == "".join(content_chunks)
+    assert _function_call_argument_text(collected) == "".join(function_argument_chunks)
+
+
+def test_tool_call_only_scan_key_is_not_skipped_as_empty():
+    assert _is_redundant_scan(StreamingScanKey(texts=("",)), None) is True
+    assert _is_redundant_scan(StreamingScanKey(texts=("",), tool_calls=("run_shell:{}",)), None) is False
+
+
+@pytest.mark.asyncio
+async def test_windowed_responses_output_item_done_round_keeps_text_window_withheld():
+    guardrail = _MarkerBlockingGuardrail(
+        guardrail_name="windowed-responses", event_hook="post_call", marker=ORIGINAL_MARKER
+    )
+    events = _responses_message_stream_events(["one ", f"{ORIGINAL_MARKER} "])
+
+    raw = await _run_windowed_responses(guardrail, events)
+
+    assert ORIGINAL_MARKER not in raw, f"unscanned window leaked: {raw!r}"
+    assert BLOCK_MESSAGE in raw
+    assert guardrail.scan_count >= 1
+
+
+@pytest.mark.asyncio
+async def test_windowed_responses_incomplete_stream_scans_tool_call_before_release():
+    guardrail = _MarkerBlockingGuardrail(
+        guardrail_name="windowed-responses-tools",
+        event_hook="post_call",
+        marker=TOOL_ARGUMENTS_MARKER,
+        field="tool_calls",
+    )
+    events = _responses_truncated_function_call_events("hi ", ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'])
+
+    raw = await _run_windowed_responses(guardrail, events)
+
+    assert '"hi "' in raw
+    assert TOOL_ARGUMENTS_MARKER not in raw, f"unscanned tool call leaked: {raw!r}"
+    assert BLOCK_MESSAGE in raw
+
+
+@pytest.mark.asyncio
+async def test_windowed_buffer_with_explicit_end_of_stream_only_stays_fully_buffered():
+    guardrail = _CountingPassingGuardrail(guardrail_name="windowed-eos", event_hook="post_call")
+    content_chunks = ["one ", "two ", "three ", "four ", "five ", "six "]
+
+    collected, yielded_count = await _run_windowed(guardrail, content_chunks, end_of_stream_only=True)
+
+    assert yielded_count == [0, 0, 0, 0, 0, 0, 0]
+    assert _chat_text(collected) == "".join(content_chunks)
+    assert guardrail.scan_count == 1
+
+
 @pytest.mark.asyncio
 async def test_buffered_mode_disabled_for_content_rewriting_guardrail():
     """Buffered replay yields the withheld *original* chunks verbatim, which

tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py

@@ -682,6 +682,22 @@ async def test_provider_specific_params_includes_embedding_toggle():
         assert field["default_value"] is False
 
 
+@pytest.mark.asyncio
+async def test_provider_specific_params_exposes_bedrock_streaming_flags():
+    from litellm.proxy.guardrails.guardrail_endpoints import get_provider_specific_params
+
+    provider_params = await get_provider_specific_params()
+
+    bedrock = provider_params["bedrock"]
+    assert "guardrailIdentifier" in bedrock
+    assert "guardrailVersion" in bedrock
+    assert bedrock["streaming_buffer_release_on_scan"]["type"] == "boolean"
+    assert bedrock["streaming_buffer_release_on_scan"]["default_value"] is False
+    assert bedrock["streaming_buffer_until_moderated"]["default_value"] is True
+    assert bedrock["streaming_end_of_stream_only"]["type"] == "boolean"
+    assert bedrock["streaming_sampling_rate"]["type"] == "number"
+
+
 @pytest.mark.asyncio
 async def test_provider_specific_params_includes_hide_secrets():
     """hide-secrets lives in the enterprise package so it is not in

tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py

@@ -7,6 +7,7 @@ import logging
 import os
 import sys
 import time
+from collections.abc import Sequence
 from contextlib import contextmanager
 from datetime import datetime, timedelta
 from typing import Any, Dict, List, Optional
@@ -32,6 +33,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
 )
 from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
 from litellm.types.caching import RedisPipelineIncrementOperation
+from litellm.types.llms.openai import ResponsesAPIResponse
 from litellm.types.utils import (
     EmbeddingResponse,
     ModelResponse,
@@ -4054,6 +4056,125 @@ async def _seed_max_parallel_requests_slots(
     )
 
 
+@pytest.mark.asyncio
+async def test_completed_responses_post_call_releases_parallel_slot() -> None:
+    api_key = hash_token("sk-responses-post-call")
+    local_cache = DualCache()
+    handler = _PROXY_MaxParallelRequestsHandler(
+        internal_usage_cache=InternalUsageCache(local_cache)
+    )
+    user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
+    data = {
+        "model": "gpt-4o-mini",
+        "input": "hello",
+        "litellm_call_id": "responses-owner",
+    }
+    parallel_key = f"{{api_key:{api_key}}}:max_parallel_requests"
+
+    await handler.async_pre_call_hook(
+        user_api_key_dict=user_api_key_dict,
+        cache=local_cache,
+        data=data,
+        call_type="aresponses",
+    )
+    assert handler._gauge_in_flight_from_cache_value(
+        await local_cache.async_get_cache(key=parallel_key)
+    ) == 1
+
+    await handler.async_post_call_success_hook(
+        data=data,
+        user_api_key_dict=user_api_key_dict,
+        response=ResponsesAPIResponse(
+            id="resp_parallel_slot",
+            created_at=0,
+            model="gpt-4o-mini",
+            object="response",
+            output=[],
+            status="completed",
+        ),
+    )
+    assert handler._gauge_in_flight_from_cache_value(
+        await local_cache.async_get_cache(key=parallel_key)
+    ) == 0
+
+    await handler.async_log_success_event(
+        kwargs={"litellm_call_id": data["litellm_call_id"]},
+        response_obj=None,
+        start_time=None,
+        end_time=None,
+    )
+    assert handler._gauge_in_flight_from_cache_value(
+        await local_cache.async_get_cache(key=parallel_key)
+    ) == 0
+
+
+@pytest.mark.asyncio
+async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None:
+    from unittest.mock import AsyncMock
+
+    api_key = hash_token("sk-concurrent-release")
+    local_cache = DualCache()
+    handler = _PROXY_MaxParallelRequestsHandler(
+        internal_usage_cache=InternalUsageCache(local_cache)
+    )
+    user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2)
+    call_id = "concurrent-release-owner"
+    parallel_key = f"{{api_key:{api_key}}}:max_parallel_requests"
+    release_started = asyncio.Event()
+    allow_redis_failure = asyncio.Event()
+
+    async def failing_release(
+        keys: Sequence[str], args: Sequence[object]
+    ) -> list[int]:
+        release_started.set()
+        await allow_redis_failure.wait()
+        raise ConnectionError("redis unavailable")
+
+    release_script = AsyncMock(side_effect=failing_release)
+    handler.parallel_release_script = release_script
+    await local_cache.async_set_cache(key=parallel_key, value=2, local_only=True)
+    stash = get_or_create_request_stash()
+    stash.owner_litellm_call_id = call_id
+    stash.parallel_slot = ParallelSlotAcquisition(
+        slot_id="slot-concurrent-release",
+        counter_keys=[parallel_key],
+    )
+    data = {"litellm_call_id": call_id}
+
+    post_call_task = asyncio.create_task(
+        handler.async_post_call_success_hook(
+            data=data,
+            user_api_key_dict=user_api_key_dict,
+            response=ResponsesAPIResponse(
+                id="resp_concurrent_release",
+                created_at=0,
+                model="gpt-4o-mini",
+                object="response",
+                output=[],
+                status="completed",
+            ),
+        )
+    )
+    await asyncio.wait_for(release_started.wait(), timeout=5)
+    logging_task = asyncio.create_task(
+        handler.async_log_success_event(
+            kwargs=data,
+            response_obj=None,
+            start_time=None,
+            end_time=None,
+        )
+    )
+    allow_redis_failure.set()
+    await asyncio.wait_for(
+        asyncio.gather(post_call_task, logging_task),
+        timeout=5,
+    )
+
+    assert release_script.await_count == 1
+    assert await local_cache.async_get_cache(key=parallel_key) == 1
+    assert stash.parallel_slot is None
+
+
 async def _build_seeded_limiter():
     """Build a v3 limiter whose api-key slot registry already holds the pre-call slot."""
     api_key = hash_token("sk-disconnect")

tests/test_litellm/proxy/management_endpoints/test_common_utils.py

@@ -35,6 +35,7 @@ from litellm.proxy.management_endpoints.common_utils import (
     admin_can_invite_user,
 )
 from litellm.proxy.management_endpoints.common_utils import _has_non_empty_value
+from litellm.types.utils import BudgetConfig
 
 
 class TestUpdateMetadataFieldsEmptyCollections:
@@ -1162,3 +1163,54 @@ async def test_router_weights_validate_current_deployment_scope(
         assert exc.value.detail == error
     else:
         await validation
+
+
+@pytest.mark.parametrize(
+    "model_max_budget, error",
+    [
+        ({"gpt-4o": BudgetConfig(max_budget=-1.0, budget_duration="1d")}, "non-negative finite"),
+        ({"gpt-4o": BudgetConfig(max_budget=float("inf"), budget_duration="1d")}, "non-negative finite"),
+        ({"gpt-4o": BudgetConfig(max_budget=float("nan"), budget_duration="1d")}, "non-negative finite"),
+        ({"gpt-4o": BudgetConfig(budget_duration="1d")}, "non-negative finite"),
+        ({"gpt-4o": BudgetConfig(max_budget=5.0)}, "requires a budget_duration"),
+        ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="fortnight")}, "budget_duration"),
+        ({"  ": BudgetConfig(max_budget=5.0, budget_duration="1d")}, "non-empty model names"),
+        ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="1d", tpm_limit=1000)}, "not enforced on a team"),
+        ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="1d", rpm_limit=10)}, "not enforced on a team"),
+    ],
+    ids=["negative", "inf", "nan", "no_cap", "no_duration", "bad_duration", "blank_model", "tpm_limit", "rpm_limit"],
+)
+def test_validate_team_model_max_budget_rejects_unenforceable_entries(model_max_budget, error) -> None:
+    from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget
+
+    with pytest.raises(HTTPException) as exc:
+        validate_team_model_max_budget(model_max_budget=model_max_budget, premium_user=True)
+    assert exc.value.status_code == 400
+    assert error in exc.value.detail["error"]
+
+
+def test_validate_team_model_max_budget_accepts_a_zero_cap_and_prefixed_models() -> None:
+    from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget
+
+    assert (
+        validate_team_model_max_budget(
+            model_max_budget={
+                "gpt-4o": BudgetConfig(max_budget=0.0, budget_duration="1d"),
+                "openai/gpt-4o-mini": BudgetConfig(max_budget=2.5, budget_duration="30d"),
+            },
+            premium_user=True,
+        )
+        is None
+    )
+
+
+def test_validate_team_model_max_budget_is_license_gated_only_when_set() -> None:
+    from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget
+
+    validate_team_model_max_budget(model_max_budget=None, premium_user=False)
+    validate_team_model_max_budget(model_max_budget={}, premium_user=False)
+    with pytest.raises(HTTPException) as exc:
+        validate_team_model_max_budget(
+            model_max_budget={"gpt-4o": BudgetConfig(max_budget=1.0, budget_duration="1d")}, premium_user=False
+        )
+    assert exc.value.status_code == 403

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -1431,6 +1431,60 @@ async def test_key_info_returns_object_permission(monkeypatch):
     )
 
 
+def _stored_key_with_lifetime_spend(token: str, spend: float, total_spend: float) -> LiteLLM_VerificationToken:
+    return LiteLLM_VerificationToken.model_validate(
+        {"token": token, "user_id": "user123", "spend": spend, "total_spend": total_spend}
+    )
+
+
+@pytest.mark.asyncio
+async def test_key_info_returns_lifetime_total_spend_next_to_resettable_spend(monkeypatch):
+    """After a budget reset the period spend is 0 while total_spend keeps the lifetime figure."""
+    from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn
+
+    mock_prisma_client = AsyncMock()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
+        return_value=_stored_key_with_lifetime_spend(token="hashed_key", spend=0.0, total_spend=3.75)
+    )
+
+    result = await info_key_fn(
+        key="sk-test-key-456",
+        user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test-key-456"),
+    )
+
+    assert result["info"]["spend"] == 0.0
+    assert result["info"]["total_spend"] == 3.75
+
+
+@pytest.mark.asyncio
+async def test_list_keys_full_object_returns_lifetime_total_spend():
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
+        return_value=[_stored_key_with_lifetime_spend(token="hashed_key", spend=0.0, total_spend=3.75)]
+    )
+    mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=1)
+
+    result = await _list_key_helper(
+        prisma_client=mock_prisma_client,
+        page=1,
+        size=50,
+        user_id=None,
+        team_id=None,
+        organization_id=None,
+        key_alias=None,
+        key_hash=None,
+        exclude_team_id=None,
+        return_full_object=True,
+        admin_team_ids=None,
+    )
+
+    listed_key = result["keys"][0]
+    assert isinstance(listed_key, UserAPIKeyAuth)
+    assert listed_key.spend == 0.0
+    assert listed_key.total_spend == 3.75
+
+
 @pytest.mark.asyncio
 async def test_get_new_token_with_valid_key(monkeypatch):
     """Test get_new_token function when provided with a valid key that starts with 'sk-'"""
@@ -4923,6 +4977,23 @@ def test_transform_verification_tokens_to_deleted_records():
     assert json.loads(record2["budget_fallbacks"]) == {"gpt-4": ["gpt-4o-mini"]}
 
 
+def test_transform_verification_tokens_to_deleted_records_keeps_organization_id():
+    live_row = MagicMock()
+    live_row.model_dump.return_value = {
+        "token": "hashed-token-org",
+        "user_id": "user-123",
+        "team_id": None,
+        "organization_id": "org-finops",
+    }
+
+    records = _transform_verification_tokens_to_deleted_records(
+        keys=[live_row],
+        user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", api_key="sk-admin"),
+    )
+
+    assert records[0]["organization_id"] == "org-finops"
+
+
 def test_transform_verification_tokens_to_deleted_records_empty_list():
     user_api_key_dict = UserAPIKeyAuth(
         user_id="user-123",
@@ -6022,6 +6093,244 @@ async def test_list_keys_with_invalid_status():
         assert "deleted" in str(exc_info.value.message)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("status_filter", ["active", "expired", "revoked"])
+async def test_list_keys_accepts_live_status_filters(monkeypatch, status_filter):
+    from unittest.mock import Mock
+
+    from litellm.proxy.management_endpoints.key_management_endpoints import list_keys
+
+    live_row = MagicMock()
+    live_row.model_dump.return_value = {"token": "hashed_live_token", "object_permission_id": None}
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[live_row])
+    mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=1)
+    mock_prisma_client.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+
+    response = await list_keys(
+        request=Mock(),
+        user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+        page=1,
+        size=10,
+        user_id=None,
+        team_id=None,
+        organization_id=None,
+        key_hash=None,
+        key_alias=None,
+        search=None,
+        return_full_object=False,
+        include_team_keys=False,
+        include_created_by_keys=False,
+        sort_by=None,
+        sort_order="desc",
+        expand=None,
+        status=status_filter,
+        project_id=None,
+        access_group_id=None,
+        agent_id=None,
+        substring_matching=False,
+        expires=None,
+    )
+
+    assert response["keys"] == ["hashed_live_token"]
+    assert response["total_count"] == 1
+    mock_prisma_client.db.litellm_deletedverificationtoken.find_many.assert_not_called()
+
+
+def _status_filter_where(status_filter: str | None) -> Mapping[str, object]:
+    from litellm.proxy.management_endpoints.key_management_endpoints import _build_key_filter_conditions
+
+    return _build_key_filter_conditions(
+        user_id=None,
+        team_id=None,
+        organization_id=None,
+        key_alias=None,
+        key_hash=None,
+        exclude_team_id=None,
+        admin_team_ids=None,
+        status_filter=status_filter,
+    )
+
+
+def test_build_key_filter_conditions_status_filter_partitions_live_keys():
+    not_blocked = {"OR": [{"blocked": None}, {"blocked": False}]}
+
+    revoked_where = _status_filter_where("revoked")
+    assert {"blocked": True} in revoked_where["AND"]
+
+    expired_clause = next(clause for clause in _status_filter_where("expired")["AND"] if "AND" in clause)
+    assert expired_clause["AND"][0] == not_blocked
+    assert expired_clause["AND"][1]["AND"][0] == {"expires": {"not": None}}
+    assert "lt" in expired_clause["AND"][1]["AND"][1]["expires"]
+
+    active_clause = next(clause for clause in _status_filter_where("active")["AND"] if "AND" in clause)
+    assert active_clause["AND"][0] == not_blocked
+    assert active_clause["AND"][1]["OR"][0] == {"expires": None}
+    assert "gte" in active_clause["AND"][1]["OR"][1]["expires"]
+
+
+def test_build_key_filter_conditions_deleted_status_adds_no_live_clause():
+    assert _status_filter_where("deleted") == _status_filter_where(None)
+
+
+@pytest.mark.asyncio
+async def test_list_key_helper_revoked_status_filters_live_table_on_blocked():
+    mock_prisma_client = AsyncMock()
+    mock_find_many = AsyncMock(return_value=[])
+    mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many
+    mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0)
+    mock_prisma_client.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
+
+    await _list_key_helper(
+        prisma_client=mock_prisma_client,
+        page=1,
+        size=50,
+        user_id=None,
+        team_id=None,
+        organization_id=None,
+        key_alias=None,
+        key_hash=None,
+        exclude_team_id=None,
+        return_full_object=True,
+        admin_team_ids=None,
+        include_created_by_keys=False,
+        status="revoked",
+    )
+
+    mock_prisma_client.db.litellm_deletedverificationtoken.find_many.assert_not_called()
+    where = mock_find_many.call_args.kwargs["where"]
+    assert {"blocked": True} in where["AND"]
+
+
+def _archived_key_row(token: str, user_id: str) -> MagicMock:
+    row = MagicMock()
+    row.model_dump.return_value = {
+        "id": "archive-row-1",
+        "token": token,
+        "key_alias": "finops-2024",
+        "user_id": user_id,
+        "team_id": None,
+        "organization_id": "org-finops",
+        "blocked": None,
+        "deleted_at": datetime(2024, 11, 15, 10, 0, tzinfo=timezone.utc),
+        "deleted_by": "admin-1",
+    }
+    return row
+
+
+@pytest.mark.asyncio
+async def test_info_key_fn_serves_deleted_key_from_archive(monkeypatch):
+    from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn
+
+    hashed = "hashed_deleted_token"
+    mock_prisma_client = AsyncMock()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
+    mock_prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(
+        return_value=_archived_key_row(hashed, "user-x")
+    )
+
+    result = await info_key_fn(
+        key=hashed,
+        user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"),
+    )
+
+    mock_prisma_client.db.litellm_deletedverificationtoken.find_first.assert_awaited_once()
+    assert mock_prisma_client.db.litellm_deletedverificationtoken.find_first.await_args.kwargs["where"] == {
+        "token": hashed
+    }
+    info = result["info"]
+    assert info["status"] == "deleted"
+    assert info["key_alias"] == "finops-2024"
+    assert info["organization_id"] == "org-finops"
+    assert info["deleted_by"] == "admin-1"
+    assert info["deleted_at"] is not None
+    assert "token" not in info
+
+
+@pytest.mark.asyncio
+async def test_info_key_fn_archived_key_keeps_owner_authorization(monkeypatch):
+    from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn
+
+    hashed = "hashed_deleted_token"
+    mock_prisma_client = AsyncMock()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
+    mock_prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(
+        return_value=_archived_key_row(hashed, "owner-1")
+    )
+
+    with pytest.raises(ProxyException) as exc_info:
+        await info_key_fn(
+            key=hashed,
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role=LitellmUserRoles.INTERNAL_USER, user_id="someone-else", api_key="sk-other"
+            ),
+        )
+    assert exc_info.value.code == "403"
+
+    owner_result = await info_key_fn(
+        key=hashed,
+        user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="owner-1", api_key="sk-own"),
+    )
+    assert owner_result["info"]["status"] == "deleted"
+
+
+@pytest.mark.asyncio
+async def test_info_key_fn_unknown_key_still_404s(monkeypatch):
+    from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn
+
+    mock_prisma_client = AsyncMock()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
+    mock_prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=None)
+
+    with pytest.raises(ProxyException) as exc_info:
+        await info_key_fn(
+            key="hashed_missing",
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"),
+        )
+    assert exc_info.value.code == "404"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("blocked", "expires", "expected_status"),
+    [
+        (True, None, "revoked"),
+        (True, "2020-01-01T00:00:00Z", "revoked"),
+        (False, "2020-01-01T00:00:00Z", "expired"),
+        (None, datetime(2020, 1, 1, tzinfo=timezone.utc), "expired"),
+        (False, None, "active"),
+        (None, "2999-01-01T00:00:00Z", "active"),
+    ],
+)
+async def test_info_key_fn_reports_live_key_status(monkeypatch, blocked, expires, expected_status):
+    from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn
+
+    mock_prisma_client = AsyncMock()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    live_row = MagicMock(spec=LiteLLM_VerificationToken)
+    live_row.model_dump.return_value = {
+        "token": "hashed_live",
+        "user_id": "user-x",
+        "team_id": None,
+        "object_permission_id": None,
+        "blocked": blocked,
+        "expires": expires,
+    }
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=live_row)
+
+    result = await info_key_fn(
+        key="hashed_live",
+        user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"),
+    )
+
+    assert result["info"]["status"] == expected_status
+    mock_prisma_client.db.litellm_deletedverificationtoken.find_first.assert_not_called()
+
+
 @pytest.mark.asyncio
 async def test_list_keys_non_admin_user_id_auto_set():
     """

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -14651,3 +14651,246 @@ async def test_team_info_reports_parent_organization_models_only_to_team_manager
         )
 
     assert response["team_info"].organization_models == expected_models
+
+
+_EXISTING_TEAM_MODEL_CAPS: Final = {
+    "gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"},
+    "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "7d"},
+}
+
+
+@pytest.mark.parametrize(
+    "requested",
+    [
+        {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 20.0, "budget_duration": "1d"}},
+        {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 10.0, "budget_duration": "30d"}},
+        {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"budget_duration": "1d"}},
+        {"claude-sonnet-4-6": _EXISTING_TEAM_MODEL_CAPS["claude-sonnet-4-6"]},
+        {},
+        None,
+        {**_EXISTING_TEAM_MODEL_CAPS, "openai/gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}},
+        {**_EXISTING_TEAM_MODEL_CAPS, "openai/gpt-4o": {"max_budget": 10.0, "budget_duration": "30d"}},
+        {**_EXISTING_TEAM_MODEL_CAPS, "anthropic/claude-sonnet-4-6": {"budget_duration": "7d"}},
+    ],
+    ids=[
+        "raise",
+        "change_duration",
+        "drop_cap_value",
+        "remove_model",
+        "clear_all",
+        "clear_with_null",
+        "raise_via_provider_alias",
+        "rewindow_via_provider_alias",
+        "uncap_via_provider_alias",
+    ],
+)
+def test_team_admin_cannot_loosen_team_model_caps(requested) -> None:
+    from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority
+
+    with pytest.raises(HTTPException) as exc:
+        _check_team_model_budget_update_authority(
+            data=UpdateTeamRequest(team_id="t1", model_max_budget=requested),
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"),
+            existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS,
+        )
+    assert exc.value.status_code == 403
+    assert "proxy admin" in exc.value.detail["error"]
+
+
+@pytest.mark.parametrize(
+    "requested",
+    [
+        {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 2.0, "budget_duration": "1d"}},
+        {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o-mini": {"max_budget": 1.0, "budget_duration": "1d"}},
+        dict(_EXISTING_TEAM_MODEL_CAPS),
+        {**_EXISTING_TEAM_MODEL_CAPS, "openai/gpt-4o": {"max_budget": 2.0, "budget_duration": "1d"}},
+    ],
+    ids=["lower", "add_model", "unchanged", "tighten_via_provider_alias"],
+)
+def test_team_admin_can_tighten_or_keep_team_model_caps(requested) -> None:
+    from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority
+
+    assert (
+        _check_team_model_budget_update_authority(
+            data=UpdateTeamRequest(team_id="t1", model_max_budget=requested),
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"),
+            existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS,
+        )
+        is None
+    )
+
+
+def test_team_model_cap_authority_skips_omitted_field_malformed_rows_and_proxy_admins() -> None:
+    from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority
+
+    team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin")
+    outcomes = (
+        _check_team_model_budget_update_authority(
+            data=UpdateTeamRequest(team_id="t1", max_budget=1.0),
+            user_api_key_dict=team_admin,
+            existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS,
+        ),
+        _check_team_model_budget_update_authority(
+            data=UpdateTeamRequest(team_id="t1", model_max_budget={}),
+            user_api_key_dict=team_admin,
+            existing_model_max_budget={"gpt-4o": "not-a-budget", "gpt-4o-mini": {"budget_duration": "1d"}},
+        ),
+        _check_team_model_budget_update_authority(
+            data=UpdateTeamRequest(team_id="t1", model_max_budget=None),
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+            existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS,
+        ),
+    )
+    assert outcomes == (None, None, None)
+
+
+@pytest.mark.asyncio
+async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_auth):
+    mock_db_client.jsonify_team_object = lambda db_data: db_data
+    mock_db_client.get_data = AsyncMock(return_value=None)
+    mock_db_client.update_data = AsyncMock(return_value=MagicMock())
+    mock_db_client.db = MagicMock()
+    mock_db_client.db.litellm_modeltable = MagicMock()
+    mock_db_client.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123"))
+
+    team_create_result = MagicMock(team_id="team-model-caps")
+    team_create_result.model_dump.return_value = {"team_id": "team-model-caps"}
+    mock_team_create = AsyncMock(return_value=team_create_result)
+    mock_db_client.db.litellm_teamtable = MagicMock()
+    mock_db_client.db.litellm_teamtable.create = mock_team_create
+    _wire_team_create_tx(mock_db_client)
+    mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
+    mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result)
+    mock_db_client.db.litellm_usertable = MagicMock()
+    mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock())
+
+    from fastapi import Request
+
+    from litellm.proxy._types import NewTeamRequest
+    from litellm.proxy.management_endpoints.team_endpoints import new_team
+
+    with patch("litellm.proxy.proxy_server.premium_user", True):  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        await new_team(
+            data=NewTeamRequest(
+                team_alias="model-caps",
+                model_max_budget={"gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}},
+            ),
+            http_request=MagicMock(spec=Request),
+            user_api_key_dict=mock_admin_auth,
+        )
+
+    team_data = mock_team_create.call_args.kwargs["data"]
+    assert team_data["model_max_budget"] == {
+        "gpt-4o": {"max_budget": 10.0, "budget_duration": "1d", "tpm_limit": None, "rpm_limit": None}
+    }
+
+
+@pytest.mark.asyncio
+async def test_new_team_rejects_unenforceable_model_max_budget(mock_db_client, mock_admin_auth):
+    from fastapi import Request
+
+    from litellm.proxy._types import NewTeamRequest, ProxyException
+    from litellm.proxy.management_endpoints.team_endpoints import new_team
+
+    mock_db_client.db.litellm_teamtable.create = AsyncMock()
+
+    with patch("litellm.proxy.proxy_server.premium_user", True), pytest.raises(ProxyException) as exc:  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        await new_team(
+            data=NewTeamRequest(team_alias="model-caps", model_max_budget={"gpt-4o": {"max_budget": 10.0}}),
+            http_request=MagicMock(spec=Request),
+            user_api_key_dict=mock_admin_auth,
+        )
+
+    assert exc.value.code == "400"
+    assert "budget_duration" in str(exc.value.message)
+    mock_db_client.db.litellm_teamtable.create.assert_not_awaited()
+
+
+def _existing_team_with_model_caps(caps):
+    existing = MagicMock()
+    existing.team_id = "standalone-team-123"
+    existing.organization_id = None
+    existing.max_budget = None
+    existing.model_id = None
+    existing.model_max_budget = caps
+    existing.model_dump.return_value = {
+        "team_id": "standalone-team-123",
+        "organization_id": None,
+        "model_max_budget": caps,
+        "members_with_roles": [{"user_id": "team-admin-model-caps", "role": "admin"}],
+    }
+    return existing
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cleared_with", [{}, None], ids=["empty_mapping", "null"])
+async def test_update_team_clearing_model_max_budget_writes_an_empty_mapping(
+    disable_audit_logging_for_mocked_team, cleared_with
+):
+    from fastapi import Request
+
+    from litellm.proxy.management_endpoints.team_endpoints import update_team
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        patch("litellm.proxy.proxy_server.premium_user", True),  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+    ):
+        mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(
+            return_value=_existing_team_with_model_caps(_EXISTING_TEAM_MODEL_CAPS)
+        )
+        mock_prisma.jsonify_team_object = lambda db_data: db_data
+        mock_cache.async_get_cache = AsyncMock(return_value=None)
+        mock_cache.async_set_cache = AsyncMock()
+        updated = _existing_team_with_model_caps({})
+        updated.litellm_model_table = None
+        mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated)
+
+        await update_team(
+            data=UpdateTeamRequest(team_id="standalone-team-123", model_max_budget=cleared_with),
+            http_request=MagicMock(spec=Request),
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
+        )
+
+    assert mock_prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["model_max_budget"] == {}
+
+
+@pytest.mark.asyncio
+async def test_update_team_model_max_budget_raise_blocked_for_team_admin():
+    from fastapi import Request
+
+    from litellm.proxy._types import ProxyException
+    from litellm.proxy.management_endpoints.team_endpoints import update_team
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        patch("litellm.proxy.proxy_server.premium_user", True),  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()),  # test-quality-ok: stubs the audit write so the test observes only the team update result
+    ):
+        mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(
+            return_value=_existing_team_with_model_caps(_EXISTING_TEAM_MODEL_CAPS)
+        )
+        mock_cache.async_get_cache = AsyncMock(return_value=None)
+        mock_prisma.db.litellm_teamtable.update = AsyncMock()
+
+        with pytest.raises(ProxyException) as exc:
+            await update_team(
+                data=UpdateTeamRequest(
+                    team_id="standalone-team-123",
+                    model_max_budget={
+                        **_EXISTING_TEAM_MODEL_CAPS,
+                        "gpt-4o": {"max_budget": 100.0, "budget_duration": "1d"},
+                    },
+                ),
+                http_request=MagicMock(spec=Request),
+                user_api_key_dict=UserAPIKeyAuth(
+                    user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin-model-caps", models=[]
+                ),
+            )
+
+    assert exc.value.code == "403"
+    assert "proxy admin" in str(exc.value.message).lower()
+    mock_prisma.db.litellm_teamtable.update.assert_not_awaited()

tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py

@@ -2541,7 +2541,7 @@ class TestRecordPartialUsageForFailure:
             function_id="test-partial-usage-failure",
         )
 
-    def _interrupted_chunks(self):
+    def _interrupted_chunks(self, *, model: str = "claude-sonnet-5"):
         return [
             self._sse(
                 "message_start",
@@ -2551,7 +2551,7 @@ class TestRecordPartialUsageForFailure:
                         "id": "msg_abc",
                         "type": "message",
                         "role": "assistant",
-                        "model": "claude-sonnet-5",
+                        "model": model,
                         "content": [],
                         "stop_reason": None,
                         "stop_sequence": None,
@@ -2588,7 +2588,7 @@ class TestRecordPartialUsageForFailure:
         AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
             litellm_logging_obj=logging_obj,
             request_body={"model": "claude-unpriced-test-model", "stream": True},
-            all_chunks=self._interrupted_chunks(),
+            all_chunks=self._interrupted_chunks(model="claude-unpriced-test-model"),
         )
 
         usage = logging_obj.model_call_details["combined_usage_object"]

tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py

@@ -47,6 +47,19 @@ def test_team_and_user_state_round_trips_through_metadata():
     )
 
 
+def test_team_model_max_budget_rides_on_the_token():
+    """The team's per-model caps must reach the token, or the auth check and the spend hook never see them."""
+    token = UserAPIKeyAuth(token="hashed", team_id="t1")
+    team_model_max_budget = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}
+    carry_team_and_user_budget_state(
+        valid_token=token,
+        team_object=LiteLLM_TeamTable(team_id="t1", model_max_budget=team_model_max_budget),
+        user_object=None,
+    )
+
+    assert token.team_model_max_budget == team_model_max_budget
+
+
 def test_missing_objects_leave_no_metadata_and_no_snapshot():
     token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1")
     carry_team_and_user_budget_state(valid_token=token, team_object=None, user_object=None)