chore: merge litellm_internal_staging into litellm_lit_4738_table_pagination

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

.github/actions/setup-uv-with-retries/action.yml

@@ -1,11 +1,7 @@
 name: "Set up uv with retries"
 description: >-
-  Install uv via astral-sh/setup-uv, retrying on transient failures. Even with
-  an exact pinned version, the action resolves the artifact URL by fetching
-  https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a
-  single request with no retry, timeout, or fallback, so one connection-level
-  network error ("fetch failed") fails the whole job before any test runs.
-  Retrying the full step covers the manifest fetch and the binary download.
+  Install uv via astral-sh/setup-uv, retrying the full setup step so manifest
+  resolution and binary downloads get fresh attempts after transient failures.
 
 inputs:
   version:
@@ -18,7 +14,7 @@ runs:
     - name: Set up uv (attempt 1)
       id: attempt-1
       continue-on-error: true
-      uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+      uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
       with:
         version: ${{ inputs.version }}
 
@@ -31,7 +27,7 @@ runs:
       id: attempt-2
       if: steps.attempt-1.outcome == 'failure'
       continue-on-error: true
-      uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+      uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
       with:
         version: ${{ inputs.version }}
 
@@ -42,6 +38,6 @@ runs:
 
     - name: Set up uv (attempt 3)
       if: steps.attempt-2.outcome == 'failure'
-      uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+      uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
       with:
         version: ${{ inputs.version }}

.github/ci-coverage-allowlist.yml

@@ -79,6 +79,11 @@ test_paths:
       - tests/load_tests/test_otel_load_test.py
       - tests/load_tests/test_vertex_embeddings_load_test.py
       - tests/load_tests/test_vertex_load_tests.py
+  - reason: >-
+      Env-gated saturation benchmark requires a live proxy and provider credentials, so it is run
+      locally rather than in pull-request jobs
+    paths:
+      - tests/load_tests/test_granian_admission_saturation.py
   - reason: >-
       A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
       localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a

.github/workflows/test-unit.yml

@@ -96,7 +96,6 @@ jobs:
           - shard: misc
             artifact-name: misc
             test-path: >-
-              tests/sdk_function_trace
               tests/test_litellm/batches
               tests/test_litellm/secret_managers
               tests/test_litellm/a2a_protocol

basedpyright-code-budget.json

@@ -3,7 +3,7 @@
     "limit": 14074
   },
   "reportArgumentType": {
-    "limit": 2214
+    "limit": 2206
   },
   "reportAssignmentType": {
     "limit": 319
@@ -57,7 +57,7 @@
     "limit": 5601
   },
   "reportMissingTypeArgument": {
-    "limit": 15287
+    "limit": 15285
   },
   "reportMissingTypeStubs": {
     "limit": 40
@@ -99,19 +99,19 @@
     "limit": 0
   },
   "reportUnknownArgumentType": {
-    "limit": 44362
+    "limit": 44360
   },
   "reportUnknownLambdaType": {
     "limit": 109
   },
   "reportUnknownMemberType": {
-    "limit": 38323
+    "limit": 38311
   },
   "reportUnknownParameterType": {
     "limit": 19624
   },
   "reportUnknownVariableType": {
-    "limit": 29861
+    "limit": 29847
   },
   "reportUnnecessaryCast": {
     "limit": 111

litellm-rust/Cargo.lock

@@ -1435,14 +1435,17 @@ dependencies = [
  "aws-sigv4",
  "aws-smithy-runtime-api",
  "aws-types",
+ "base64",
  "rand 0.8.7",
  "reqwest",
+ "rstest",
  "serde",
  "serde_json",
  "sha2 0.10.9",
  "thiserror 2.0.19",
  "tokio",
  "tracing",
+ "tracing-subscriber",
 ]
 
 [[package]]
@@ -1461,7 +1464,6 @@ dependencies = [
  "tokio",
  "tokio-tungstenite",
  "tracing",
- "tracing-subscriber",
 ]
 
 [[package]]
@@ -1511,6 +1513,16 @@ version = "0.3.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
 
+[[package]]
+name = "mime_guess"
+version = "2.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e"
+dependencies = [
+ "mime",
+ "unicase",
+]
+
 [[package]]
 name = "mio"
 version = "1.2.2"
@@ -1969,6 +1981,7 @@ dependencies = [
  "hyper-util",
  "js-sys",
  "log",
+ "mime_guess",
  "percent-encoding",
  "pin-project-lite",
  "quinn",
@@ -2736,6 +2749,12 @@ version = "1.20.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20"
 
+[[package]]
+name = "unicase"
+version = "2.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142"
+
 [[package]]
 name = "unicode-ident"
 version = "1.0.24"

litellm-rust/Cargo.toml

@@ -24,7 +24,7 @@ pyo3 = "0.29.2"
 pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
 pythonize = "0.29.0"
 rand = "0.8"
-reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
+reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
 rstest = "0.26.1"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0", features = ["float_roundtrip"] }

litellm-rust/crates/ai-gateway/Cargo.toml

@@ -32,6 +32,7 @@ subtle = { workspace = true, optional = true }
 # SHA-256 hash_token) so the plaintext credential never enters a log payload.
 sha2 = { workspace = true, optional = true }
 pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
+tower = { version = "0.5.3", features = ["util"], optional = true }
 
 [features]
 default = []
@@ -39,6 +40,7 @@ server = ["dep:axum", "dep:subtle", "dep:sha2"]
 # Build the gateway's config from the proxy YAML via an embedded Python
 # interpreter (links libpython; requires `litellm` importable at runtime).
 python-config = ["dep:pyo3"]
+trace-parity = ["server", "dep:tower", "litellm-core/observability"]
 
 [dev-dependencies]
 futures-channel = "0.3"

litellm-rust/crates/ai-gateway/src/lib.rs

@@ -28,6 +28,8 @@ pub mod auth;
 pub mod routes;
 #[cfg(feature = "server")]
 pub mod state;
+#[cfg(feature = "trace-parity")]
+pub mod trace_parity;
 
 mod constants;
 pub mod integrations;

litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs

@@ -12,6 +12,7 @@ use litellm_core::providers::azure_ai::ocr::transformation::{
     AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG,
 };
 use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
+use litellm_core::providers::reducto::ocr::transformation as reducto;
 use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
 use litellm_core::providers::vertex_ai::ocr::transformation::{
     VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
@@ -39,6 +40,7 @@ pub(super) fn ocr_provider_config(
 ) -> Option<&'static dyn OcrProviderConfig> {
     match provider {
         "mistral" => Some(&MISTRAL_OCR_CONFIG),
+        "reducto" => reducto::config_for_model(model),
         "azure_ai" if is_azure_document_intelligence_model(model) => {
             Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
         }
@@ -334,6 +336,7 @@ fn operation_status(response_json: &Value) -> Result<&str, Error> {
     }
 }
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub(super) async fn poll_document_intelligence(
     operation_url: &str,
     original_url: &str,
@@ -392,9 +395,11 @@ pub(super) async fn poll_document_intelligence(
 
 #[cfg(test)]
 mod tests {
-    use super::*;
+    use litellm_core::ocr::transformation::OcrResponseHandling;
     use serde_json::json;
 
+    use super::*;
+
     #[test]
     fn blocks_private_and_metadata_ips() {
         assert!(is_blocked_ip("127.0.0.1".parse().unwrap()));
@@ -438,4 +443,87 @@ mod tests {
 
         assert_eq!(transformed, document);
     }
+
+    #[test]
+    fn truncate_error_body_passes_short_strings_through() {
+        let body = "Unauthorized";
+        assert_eq!(truncate_error_body(body), "Unauthorized");
+    }
+
+    #[test]
+    fn truncate_error_body_caps_long_payloads() {
+        let body = "x".repeat(306);
+        let truncated = truncate_error_body(&body);
+
+        assert!(truncated.ends_with("... (truncated)"));
+        let prefix_chars = truncated
+            .strip_suffix("... (truncated)")
+            .expect("truncated marker present")
+            .chars()
+            .count();
+        assert_eq!(prefix_chars, 256);
+    }
+
+    #[test]
+    fn truncate_error_body_does_not_split_multibyte_chars() {
+        let body = "é".repeat(266);
+        let truncated = truncate_error_body(&body);
+        assert!(truncated.is_char_boundary(truncated.len()));
+    }
+
+    #[test]
+    fn ocr_dispatch_supports_migrated_providers() {
+        assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
+        assert!(
+            ocr_provider_config("azure_ai", "pixtral-12b-2409")
+                .expect("azure ai config resolves")
+                .requires_data_uri_document()
+        );
+        assert_eq!(
+            ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
+                .expect("document intelligence config resolves")
+                .response_handling(),
+            OcrResponseHandling::AzureDocumentIntelligencePoll
+        );
+        assert!(
+            ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
+                .expect("vertex deepseek config resolves")
+                .supported_ocr_params()
+                .contains(&"temperature")
+        );
+        assert!(ocr_provider_config("openai", "gpt-4o").is_none());
+    }
+
+    #[test]
+    fn string_headers_accepts_string_values() {
+        let headers = json!({
+            "x-trace-id": "trace-1"
+        })
+        .as_object()
+        .unwrap()
+        .clone();
+
+        assert_eq!(
+            string_headers(Some(headers)).expect("string headers accepted"),
+            vec![("x-trace-id".to_string(), "trace-1".to_string())]
+        );
+    }
+
+    #[test]
+    fn string_headers_rejects_non_string_values() {
+        let headers = json!({
+            "x-retry-count": 3
+        })
+        .as_object()
+        .unwrap()
+        .clone();
+
+        let err = string_headers(Some(headers)).expect_err("non-string header rejected");
+        assert_eq!(
+            err,
+            Error::InvalidRequest(
+                "OCR extra_headers.x-retry-count must be a string, got number".to_string()
+            )
+        );
+    }
 }

litellm-rust/crates/ai-gateway/src/ocr/handler.rs

@@ -50,7 +50,11 @@ pub(crate) async fn execute_ocr_provider_call(
         .await?;
         return Ok(request
             .config
-            .transform_ocr_response(&request.model, response_json)?
+            .transform_ocr_response_with_params(
+                &request.model,
+                response_json,
+                &request.optional_params,
+            )?
             .into_json());
     }
 
@@ -71,6 +75,10 @@ pub(crate) async fn execute_ocr_provider_call(
 
     Ok(request
         .config
-        .transform_ocr_response(&request.model, response_json)?
+        .transform_ocr_response_with_params(
+            &request.model,
+            response_json,
+            &request.optional_params,
+        )?
         .into_json())
 }

litellm-rust/crates/ai-gateway/src/ocr/hooks.rs

@@ -1,11 +1,15 @@
 use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
 use litellm_core::error::Error;
+use litellm_core::providers::reducto::ocr::transformation::{
+    build_upload_request, extract_document_source, extract_upload_file_id,
+};
 use serde_json::{Map, Value, json};
 use std::future::Future;
 use std::pin::Pin;
 
-use super::common_utils::{convert_document_url_to_data_uri, string_headers};
+use super::common_utils::{convert_document_url_to_data_uri, string_headers, truncate_error_body};
 use super::types::{PreparedOcrRequest, ProviderOcrRequest};
+use crate::client::http_client;
 use crate::integrations::custom_guardrail::{
     CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
 };
@@ -89,22 +93,39 @@ impl OcrLifecycleHooks {
         )?;
         let model = request.model.clone();
         let custom_llm_provider = request.custom_llm_provider.clone();
-        let document = if config.requires_data_uri_document() {
+        let is_reducto = custom_llm_provider == "reducto";
+        let document = if is_reducto {
+            let guarded_document = self
+                .run_during_call_guardrails(&model, &custom_llm_provider, &url, request.document)
+                .await?;
+            upload_reducto_document(
+                &guarded_document,
+                request.api_base.as_deref(),
+                request.timeout,
+                &upstream_headers,
+            )
+            .await?
+        } else if config.requires_data_uri_document() {
             convert_document_url_to_data_uri(request.document).await?
         } else {
             request.document
         };
+        let optional_params = request.optional_params;
         let body = config
-            .transform_ocr_request(&request.model, document, request.optional_params)?
+            .transform_ocr_request(&request.model, document, optional_params.clone())?
             .data;
-        let body = self
-            .run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
-            .await?;
+        let body = if is_reducto {
+            body
+        } else {
+            self.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
+                .await?
+        };
         Ok(ProviderOcrRequest {
             model,
             config,
             url,
             body,
+            optional_params,
             upstream_headers,
             timeout: request.timeout,
         })
@@ -165,6 +186,63 @@ impl OcrLifecycleHooks {
     }
 }
 
+async fn upload_reducto_document(
+    document: &Value,
+    api_base: Option<&str>,
+    timeout: Option<std::time::Duration>,
+    upstream_headers: &[(String, String)],
+) -> Result<Value, Error> {
+    let source = extract_document_source(document)?;
+    let Some(authorization) = upstream_headers
+        .iter()
+        .find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
+        .map(|(_, value)| value.as_str())
+    else {
+        return Err(Error::Auth(
+            "Reducto upload requires an Authorization header".to_string(),
+        ));
+    };
+    let Some(upload) = build_upload_request(source, authorization, api_base) else {
+        return Ok(document.clone());
+    };
+    let part = reqwest::multipart::Part::bytes(upload.bytes)
+        .file_name(upload.file_name)
+        .mime_str(&upload.mime_type)
+        .map_err(|error| Error::InvalidRequest(error.to_string()))?;
+    let form = reqwest::multipart::Form::new().part("file", part);
+    let mut request_builder = http_client().post(upload.url).multipart(form);
+    for (name, value) in upstream_headers {
+        if !name.eq_ignore_ascii_case("content-type")
+            && !name.eq_ignore_ascii_case("content-length")
+        {
+            request_builder = request_builder.header(name, value);
+        }
+    }
+    if let Some(timeout) = timeout {
+        request_builder = request_builder.timeout(timeout);
+    }
+    let response = request_builder
+        .send()
+        .await
+        .map_err(|error| Error::Network(error.to_string()))?;
+    let status = response.status();
+    let body = response
+        .text()
+        .await
+        .map_err(|error| Error::Network(error.to_string()))?;
+    if !status.is_success() {
+        return Err(Error::Http {
+            status: status.as_u16(),
+            body: truncate_error_body(&body),
+        });
+    }
+    let response_json: Value = serde_json::from_str(&body).map_err(|error| {
+        Error::InvalidResponse(format!("invalid Reducto upload response JSON: {error}"))
+    })?;
+    let file_id = extract_upload_file_id(&response_json)?;
+    Ok(json!({"type": "document_url", "document_url": file_id}))
+}
+
 impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLifecycleHooks {
     type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
     type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;

litellm-rust/crates/ai-gateway/src/ocr/mod.rs

@@ -24,4 +24,151 @@ pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
 }
 
 #[cfg(test)]
-mod tests;
+mod tests {
+    use serde_json::{Map, json};
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+    use tokio::net::{TcpListener, TcpStream};
+
+    use super::{OcrRequest, ocr};
+    use crate::integrations::types::RequestMetadata;
+
+    async fn read_http_request(socket: &mut TcpStream) -> String {
+        let mut request = Vec::new();
+        let mut buffer = [0_u8; 1024];
+        let header_end = loop {
+            let n = socket.read(&mut buffer).await.expect("reads request");
+            if n == 0 {
+                break request.len();
+            }
+            request.extend_from_slice(&buffer[..n]);
+            if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
+                break position + 4;
+            }
+        };
+        let headers = String::from_utf8_lossy(&request[..header_end]);
+        let content_length = headers
+            .lines()
+            .find_map(|line| {
+                let (name, value) = line.split_once(':')?;
+                name.eq_ignore_ascii_case("content-length")
+                    .then(|| value.trim().parse::<usize>().ok())
+                    .flatten()
+            })
+            .unwrap_or(0);
+        while request.len().saturating_sub(header_end) < content_length {
+            let n = socket.read(&mut buffer).await.expect("reads body");
+            if n == 0 {
+                break;
+            }
+            request.extend_from_slice(&buffer[..n]);
+        }
+        String::from_utf8(request).expect("request is utf8")
+    }
+
+    fn base_ocr_request(model: &str) -> OcrRequest<'_> {
+        OcrRequest {
+            model,
+            document: json!({
+                "type": "document_url",
+                "document_url": "https://example.com/doc.pdf"
+            }),
+            api_key: Some("sk-test"),
+            api_base: None,
+            custom_llm_provider: None,
+            extra_headers: None,
+            optional_params: Map::new(),
+            timeout: None,
+            callbacks: Vec::new(),
+            guardrails: Vec::new(),
+            request_metadata: RequestMetadata::default(),
+            litellm_call_id: None,
+        }
+    }
+
+    #[tokio::test]
+    async fn reducto_file_upload_then_parse_maps_response() {
+        let listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("test listener binds");
+        let address = listener.local_addr().expect("listener has local address");
+        let server = tokio::spawn(async move {
+            let (mut upload_socket, _) = listener.accept().await.expect("accepts upload request");
+            let upload_request = read_http_request(&mut upload_socket).await;
+            let upload_body = r#"{"file_id":"reducto://uploaded.pdf"}"#;
+            let upload_response = format!(
+                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+                upload_body.len(),
+                upload_body
+            );
+            upload_socket
+                .write_all(upload_response.as_bytes())
+                .await
+                .expect("writes upload response");
+
+            let (mut parse_socket, _) = listener.accept().await.expect("accepts parse request");
+            let parse_request = read_http_request(&mut parse_socket).await;
+            let parse_body = r#"{"job_id":"job_123","usage":{"num_pages":3,"credits":3},"result":{"chunks":[{"content":"Page 1 block A","blocks":[{"content":"Page 1 block A","bbox":{"page":1},"kind":"text"}]},{"content":"Page 2 block A","blocks":[{"content":"Page 2 block A","bbox":{"page":2},"kind":"table"}]},{"content":"Page 1 block B","blocks":[{"content":"Page 1 block B","bbox":{"page":1},"kind":"text"}]},{"content":"Page 3 block A","blocks":[{"content":"Page 3 block A","bbox":{"page":3},"kind":"figure"}]}]}}"#;
+            let parse_response = format!(
+                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+                parse_body.len(),
+                parse_body
+            );
+            parse_socket
+                .write_all(parse_response.as_bytes())
+                .await
+                .expect("writes parse response");
+            (upload_request, parse_request)
+        });
+        let api_base = format!("http://{address}");
+        let mut request = base_ocr_request("reducto/parse-v3");
+        request.api_base = Some(&api_base);
+        request.api_key = None;
+        request.extra_headers = Some(Map::from_iter([
+            ("Authorization".to_string(), json!("Bearer test-key")),
+            ("x-trace-id".to_string(), json!("trace-1")),
+        ]));
+        request.document = json!({
+            "type": "document_url",
+            "document_url": "data:application/pdf;base64,JVBERi0xLjQ="
+        });
+        request.optional_params = Map::from_iter([
+            (
+                "formatting".to_string(),
+                json!({"table_output_format": "html"}),
+            ),
+            ("retrieval".to_string(), json!({"chunk_mode": "section"})),
+            ("settings".to_string(), json!({"ocr_system": "standard"})),
+        ]);
+
+        let response = ocr(request).await.expect("Reducto OCR succeeds");
+
+        assert_eq!(response["pages"].as_array().map(Vec::len), Some(3));
+        assert_eq!(
+            response["pages"][0]["markdown"],
+            "Page 1 block A\n\nPage 1 block B"
+        );
+        assert_eq!(response["pages"][1]["markdown"], "Page 2 block A");
+        assert_eq!(response["pages"][2]["markdown"], "Page 3 block A");
+        assert_eq!(response["usage_info"]["pages_processed"], 3);
+        assert_eq!(response["usage_info"]["credits"], 3);
+        assert_eq!(response["provider_native_response"]["job_id"], "job_123");
+        let (upload_request, parse_request) = server.await.expect("server task completes");
+        assert!(
+            upload_request
+                .to_ascii_lowercase()
+                .contains("authorization: bearer test-key")
+        );
+        assert!(upload_request.contains("application/pdf"));
+        assert!(upload_request.contains("%PDF-1.4"));
+        assert!(upload_request.contains("x-trace-id: trace-1"));
+        assert!(
+            parse_request
+                .to_ascii_lowercase()
+                .contains("authorization: bearer test-key")
+        );
+        assert!(parse_request.contains(r#""input":"reducto://uploaded.pdf""#));
+        assert!(parse_request.contains(r#""table_output_format":"html""#));
+        assert!(parse_request.contains(r#""chunk_mode":"section""#));
+        assert!(parse_request.contains(r#""ocr_system":"standard""#));
+    }
+}

litellm-rust/crates/ai-gateway/src/ocr/prepare.rs

@@ -2,6 +2,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
+use serde_json::{Map, Value};
 
 use super::common_utils::ocr_provider_config;
 use super::hooks::OcrLifecycleHooks;
@@ -28,17 +29,33 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
     let model = provider_info.model.to_string();
     let custom_llm_provider = provider_info.custom_llm_provider.to_string();
     let config = ocr_provider_config(&custom_llm_provider, &model)
-        .ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()));
+        .ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()))
+        .and_then(|config| {
+            validate_request_format(config, &request.optional_params, &custom_llm_provider)?;
+            Ok(config)
+        });
     let optional_params = match &config {
         Ok(config) => {
             let supported = config.supported_ocr_params();
-            config.map_ocr_params(
+            let mut mapped = config.map_ocr_params(
                 &request
                     .optional_params
-                    .into_iter()
+                    .iter()
                     .filter(|(name, _)| supported.contains(&name.as_str()))
+                    .map(|(name, value)| (name.clone(), value.clone()))
                     .collect(),
-            )
+            );
+            for name in [
+                "vertex_project",
+                "vertex_ai_project",
+                "vertex_location",
+                "vertex_ai_location",
+            ] {
+                if let Some(value) = request.optional_params.get(name) {
+                    mapped.insert(name.to_string(), value.clone());
+                }
+            }
+            mapped
         }
         Err(_) => request.optional_params,
     };
@@ -64,6 +81,26 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
     }
 }
 
+fn validate_request_format(
+    config: &'static dyn litellm_core::ocr::transformation::OcrProviderConfig,
+    optional_params: &Map<String, Value>,
+    provider: &str,
+) -> Result<(), litellm_core::Error> {
+    let Some(format) = optional_params.get("req_format") else {
+        return Ok(());
+    };
+    match format.as_str() {
+        Some("litellm") => Ok(()),
+        Some("native") if config.supported_ocr_params().contains(&"req_format") => Ok(()),
+        Some("native") => Err(litellm_core::Error::InvalidRequest(format!(
+            "`req_format=native` is not supported for provider {provider}"
+        ))),
+        _ => Err(litellm_core::Error::InvalidRequest(format!(
+            "Invalid `req_format`: {format}. Expected `litellm` or `native`"
+        ))),
+    }
+}
+
 fn new_ocr_call_id() -> String {
     static COUNTER: AtomicU64 = AtomicU64::new(1);
     let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
@@ -73,3 +110,54 @@ fn new_ocr_call_id() -> String {
         .unwrap_or(0);
     format!("ocr-{timestamp}-{sequence}")
 }
+
+#[cfg(test)]
+mod tests {
+    use litellm_core::error::Error;
+    use serde_json::{Map, json};
+
+    use super::{OcrRequest, prepare_ocr_call};
+    use crate::integrations::types::RequestMetadata;
+
+    fn base_ocr_request(model: &str) -> OcrRequest<'_> {
+        OcrRequest {
+            model,
+            document: json!({
+                "type": "document_url",
+                "document_url": "https://example.com/doc.pdf"
+            }),
+            api_key: Some("sk-test"),
+            api_base: None,
+            custom_llm_provider: None,
+            extra_headers: None,
+            optional_params: Map::new(),
+            timeout: None,
+            callbacks: Vec::new(),
+            guardrails: Vec::new(),
+            request_metadata: RequestMetadata::default(),
+            litellm_call_id: None,
+        }
+    }
+
+    fn request_with_format(format: &str) -> OcrRequest<'_> {
+        let mut request = base_ocr_request("mistral/mistral-ocr-latest");
+        request.optional_params = Map::from_iter([("req_format".to_string(), json!(format))]);
+        request
+    }
+
+    #[test]
+    fn native_format_rejected_for_provider_without_support_as_bad_request() {
+        let prepared = prepare_ocr_call(request_with_format("native"));
+        assert!(
+            matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("not supported for provider"))
+        );
+    }
+
+    #[test]
+    fn unknown_format_rejected_for_provider_without_support_as_bad_request() {
+        let prepared = prepare_ocr_call(request_with_format("raw"));
+        assert!(
+            matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("Invalid `req_format`"))
+        );
+    }
+}

litellm-rust/crates/ai-gateway/src/ocr/types.rs

@@ -53,6 +53,7 @@ pub(crate) struct ProviderOcrRequest {
     pub(crate) config: &'static dyn OcrProviderConfig,
     pub(crate) url: String,
     pub(crate) body: Value,
+    pub(crate) optional_params: Map<String, Value>,
     pub(crate) upstream_headers: Vec<(String, String)>,
     pub(crate) timeout: Option<Duration>,
 }

litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs

@@ -21,6 +21,12 @@ pub fn router() -> Router<AppState> {
     Router::new().route(MESSAGES_ROUTE_PATH, post(handle))
 }
 
+#[tracing::instrument(
+    name = "messages_gateway_route",
+    target = "litellm::function_trace",
+    level = "trace",
+    skip_all
+)]
 async fn handle(
     _auth: RequireMasterKey,
     State(state): State<AppState>,

litellm-rust/crates/ai-gateway/src/routes/messages/service.rs

@@ -12,6 +12,12 @@ pub(crate) enum MessagesResponse {
     Stream(reqwest::Response),
 }
 
+#[tracing::instrument(
+    name = "messages_gateway_service",
+    target = "litellm::function_trace",
+    level = "trace",
+    skip_all
+)]
 pub async fn run(
     router: &Arc<Router>,
     body: Value,

litellm-rust/crates/ai-gateway/src/trace_parity.rs

@@ -0,0 +1,65 @@
+//! Harness-only in-process adapters. Never mounted as production routes.
+
+use std::sync::Arc;
+
+use axum::body::{Body, to_bytes};
+use axum::http::header::{AUTHORIZATION, CONTENT_TYPE};
+use axum::http::{Request, StatusCode};
+use litellm_core::Error;
+use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
+use serde::Serialize;
+use serde_json::Value;
+use tower::ServiceExt;
+
+use crate::io::realtime_pool::RealtimePool;
+use crate::routes;
+use crate::state::AppState;
+
+#[derive(Debug, Serialize)]
+pub struct GatewayResponse {
+    pub status: u16,
+    pub body: Value,
+}
+
+pub async fn messages_request(
+    model_alias: String,
+    provider_model: String,
+    api_base: String,
+    body: Value,
+) -> Result<GatewayResponse, Error> {
+    let state = AppState {
+        router: Arc::new(ModelRouter::new(vec![Deployment {
+            model_name: model_alias,
+            litellm_params: LiteLLMParams {
+                model: provider_model,
+                api_key: Some("trace-provider-key".to_string()),
+                api_base: Some(api_base),
+            },
+        }])),
+        master_key: Some(Arc::from("trace-master-key")),
+        loggers: Arc::new(Vec::new()),
+        realtime_pool: RealtimePool::disabled(),
+    };
+    let request = Request::builder()
+        .method("POST")
+        .uri("/v1/messages")
+        .header(AUTHORIZATION, "Bearer trace-master-key")
+        .header(CONTENT_TYPE, "application/json")
+        .body(Body::from(body.to_string()))
+        .map_err(|error| Error::InvalidRequest(error.to_string()))?;
+    let response = routes::app(state)
+        .oneshot(request)
+        .await
+        .map_err(|error| match error {})?;
+    let status: StatusCode = response.status();
+    let bytes = to_bytes(response.into_body(), usize::MAX)
+        .await
+        .map_err(|error| Error::InvalidResponse(error.to_string()))?;
+    let body = serde_json::from_slice(&bytes).map_err(|error| {
+        Error::InvalidResponse(format!("gateway returned invalid JSON: {error}"))
+    })?;
+    Ok(GatewayResponse {
+        status: status.as_u16(),
+        body,
+    })
+}

litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs

@@ -1,23 +1,19 @@
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
-use litellm_core::error::Error;
-use litellm_core::http_utils::has_header;
-use litellm_core::ocr::transformation::OcrResponseHandling;
-use serde_json::{Map, Value, json};
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
-use tokio::net::{TcpListener, TcpStream};
-
-use super::common_utils::{ocr_provider_config, string_headers, truncate_error_body};
-use super::{OcrRequest, ocr};
-use crate::integrations::custom_guardrail::{
+use litellm_ai_gateway::integrations::custom_guardrail::{
     CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
     GuardrailFuture, GuardrailRequest,
 };
-use crate::integrations::custom_logger::{
+use litellm_ai_gateway::integrations::custom_logger::{
     CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
 };
-use crate::integrations::types::RequestMetadata;
+use litellm_ai_gateway::integrations::types::RequestMetadata;
+use litellm_ai_gateway::ocr::{OcrRequest, ocr};
+use litellm_core::error::Error;
+use serde_json::{Map, Value, json};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::{TcpListener, TcpStream};
 
 async fn read_http_headers(socket: &mut TcpStream) -> String {
     let mut request = Vec::new();
@@ -136,6 +132,7 @@ struct RecordingOcrGuardrail {
     hooks: Vec<GuardrailEventHook>,
     events: Mutex<Vec<&'static str>>,
     block_pre_call: bool,
+    block_during_call: bool,
 }
 
 impl RecordingOcrGuardrail {
@@ -144,6 +141,7 @@ impl RecordingOcrGuardrail {
             hooks,
             events: Mutex::new(Vec::new()),
             block_pre_call: false,
+            block_during_call: false,
         }
     }
 
@@ -152,6 +150,16 @@ impl RecordingOcrGuardrail {
             hooks: vec![GuardrailEventHook::PreCall],
             events: Mutex::new(Vec::new()),
             block_pre_call: true,
+            block_during_call: false,
+        }
+    }
+
+    fn blocking_during_call() -> Self {
+        Self {
+            hooks: vec![GuardrailEventHook::DuringCall],
+            events: Mutex::new(Vec::new()),
+            block_pre_call: false,
+            block_during_call: true,
         }
     }
 
@@ -193,91 +201,95 @@ impl CustomGuardrail for RecordingOcrGuardrail {
     ) -> GuardrailFuture<'a> {
         Box::pin(async move {
             self.events.lock().unwrap().push("async_moderation_hook");
+            if self.block_during_call {
+                return Ok(GuardrailDecision::Block(GuardrailError::blocked(
+                    "blocked before provider",
+                )));
+            }
             request.data["body"]["guarded_during"] = json!(true);
             Ok(GuardrailDecision::Mask(request))
         })
     }
 }
 
-#[test]
-fn truncate_error_body_passes_short_strings_through() {
-    let body = "Unauthorized";
-    assert_eq!(truncate_error_body(body), "Unauthorized");
+fn base_ocr_request(model: &str) -> OcrRequest<'_> {
+    OcrRequest {
+        model,
+        document: json!({
+            "type": "document_url",
+            "document_url": "https://example.com/doc.pdf"
+        }),
+        api_key: Some("sk-test"),
+        api_base: None,
+        custom_llm_provider: None,
+        extra_headers: None,
+        optional_params: Map::new(),
+        timeout: None,
+        callbacks: Vec::new(),
+        guardrails: Vec::new(),
+        request_metadata: RequestMetadata::default(),
+        litellm_call_id: None,
+    }
 }
 
-#[test]
-fn truncate_error_body_caps_long_payloads() {
-    let body = "x".repeat(306);
-    let truncated = truncate_error_body(&body);
+#[tokio::test]
+async fn reducto_during_call_guardrail_blocks_before_upload() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("test listener binds");
+    let address = listener.local_addr().expect("listener has local address");
+    let api_base = format!("http://{address}");
+    let guardrail = Arc::new(RecordingOcrGuardrail::blocking_during_call());
+    let mut request = base_ocr_request("reducto/parse-v3");
+    request.api_base = Some(&api_base);
+    request.document = json!({
+        "type": "document_url",
+        "document_url": "data:application/pdf;base64,JVBERi0xLjQ="
+    });
+    request.guardrails = vec![guardrail.clone()];
 
-    assert!(truncated.ends_with("... (truncated)"));
-    let prefix_chars = truncated
-        .strip_suffix("... (truncated)")
-        .expect("truncated marker present")
-        .chars()
-        .count();
-    assert_eq!(prefix_chars, 256);
+    let error = ocr(request).await.expect_err("guardrail blocks upload");
+
+    assert!(matches!(error, Error::InvalidRequest(_)));
+    assert_eq!(guardrail.events(), vec!["async_moderation_hook"]);
+    let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
+    assert!(accepted.is_err(), "upload socket should not be touched");
 }
 
-#[test]
-fn truncate_error_body_does_not_split_multibyte_chars() {
-    let body = "é".repeat(266);
-    let truncated = truncate_error_body(&body);
-    assert!(truncated.is_char_boundary(truncated.len()));
-}
+#[tokio::test]
+async fn reducto_upload_error_body_is_truncated() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("test listener binds");
+    let address = listener.local_addr().expect("listener has local address");
+    let server = tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.expect("accepts upload request");
+        let _request = read_http_request(&mut socket).await;
+        let body = "x".repeat(300);
+        let response = format!(
+            "HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+            body.len(),
+            body
+        );
+        socket
+            .write_all(response.as_bytes())
+            .await
+            .expect("writes upload response");
+    });
+    let api_base = format!("http://{address}");
+    let mut request = base_ocr_request("reducto/parse-v3");
+    request.api_base = Some(&api_base);
+    request.document = json!({
+        "type": "document_url",
+        "document_url": "data:application/pdf;base64,JVBERi0xLjQ="
+    });
+
+    let error = ocr(request).await.expect_err("upload should fail");
 
-#[test]
-fn ocr_dispatch_supports_migrated_providers() {
-    assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
     assert!(
-        ocr_provider_config("azure_ai", "pixtral-12b-2409")
-            .expect("azure ai config resolves")
-            .requires_data_uri_document()
+        matches!(error, Error::Http { status: 500, body } if body.chars().count() < 300 && body.ends_with("... (truncated)"))
     );
-    assert_eq!(
-        ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
-            .expect("document intelligence config resolves")
-            .response_handling(),
-        OcrResponseHandling::AzureDocumentIntelligencePoll
-    );
-    assert!(
-        ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
-            .expect("vertex deepseek config resolves")
-            .supported_ocr_params()
-            .contains(&"temperature")
-    );
-    assert!(ocr_provider_config("openai", "gpt-4o").is_none());
-}
-
-#[test]
-fn string_headers_accepts_string_values() {
-    let headers = json!({
-        "x-trace-id": "trace-1"
-    })
-    .as_object()
-    .unwrap()
-    .clone();
-
-    assert_eq!(
-        string_headers(Some(headers)).expect("string headers accepted"),
-        vec![("x-trace-id".to_string(), "trace-1".to_string())]
-    );
-}
-
-#[test]
-fn auth_header_detection_is_case_insensitive() {
-    let headers = vec![
-        ("x-trace-id".to_string(), "trace-1".to_string()),
-        ("authorization".to_string(), "Bearer sk-test".to_string()),
-    ];
-
-    assert!(has_header(&headers, "authorization"));
-
-    let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())];
-    assert!(has_header(&headers, "authorization"));
-
-    let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())];
-    assert!(!has_header(&headers, "authorization"));
+    server.await.expect("server task completes");
 }
 
 #[tokio::test]
@@ -595,21 +607,3 @@ async fn document_intelligence_poll_uses_resolved_subscription_key() {
         "{poll_request}"
     );
 }
-
-#[test]
-fn string_headers_rejects_non_string_values() {
-    let headers = json!({
-        "x-retry-count": 3
-    })
-    .as_object()
-    .unwrap()
-    .clone();
-
-    let err = string_headers(Some(headers)).expect_err("non-string header rejected");
-    assert_eq!(
-        err,
-        Error::InvalidRequest(
-            "OCR extra_headers.x-retry-count must be a string, got number".to_string()
-        )
-    );
-}

litellm-rust/crates/core/Cargo.toml

@@ -6,12 +6,14 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
+base64.workspace = true
 rand.workspace = true
 reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
+tracing-subscriber = { workspace = true, optional = true }
 sha2.workspace = true
 aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
 aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
@@ -30,6 +32,9 @@ bedrock-auth = [
     "dep:aws-types",
     "dep:aws-smithy-runtime-api",
 ]
+observability = ["dep:tracing-subscriber"]
 
 [dev-dependencies]
+rstest.workspace = true
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
+tracing-subscriber.workspace = true

litellm-rust/crates/core/src/constants.rs

@@ -41,3 +41,5 @@ pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
 /// `litellm/litellm_core_utils/prompt_templates/factory.py`.
 pub const EMPTY_TEXT_PLACEHOLDER: &str =
     "[System: Empty message content sanitised to satisfy protocol]";
+
+pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";

litellm-rust/crates/core/src/http_utils.rs

@@ -101,6 +101,23 @@ mod tests {
         assert!(!has_header(&headers, "authorization"));
     }
 
+    #[test]
+    fn auth_header_detection_is_case_insensitive() {
+        let headers = vec![
+            ("x-trace-id".to_string(), "trace-1".to_string()),
+            ("authorization".to_string(), "Bearer sk-test".to_string()),
+        ];
+
+        assert!(has_header(&headers, "authorization"));
+
+        let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())];
+
+        assert!(has_header(&headers, "authorization"));
+
+        let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())];
+        assert!(!has_header(&headers, "authorization"));
+    }
+
     #[test]
     fn bearer_detection_requires_a_non_empty_token() {
         assert!(has_bearer_auth(&[(

litellm-rust/crates/core/src/lib.rs

@@ -6,6 +6,8 @@ pub mod constants;
 pub mod error;
 pub mod http_utils;
 pub mod messages;
+#[cfg(any(feature = "observability", test))]
+pub mod observability;
 pub mod ocr;
 pub mod providers;
 pub mod realtime;

litellm-rust/crates/core/src/observability/function_trace.rs

@@ -0,0 +1,215 @@
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+
+use serde::Serialize;
+use tracing::span::{Attributes, Id};
+use tracing::{Dispatch, Subscriber};
+use tracing_subscriber::layer::Context;
+use tracing_subscriber::prelude::*;
+use tracing_subscriber::registry::LookupSpan;
+use tracing_subscriber::{Layer, Registry};
+
+use super::function_trace_filter;
+
+#[derive(Clone, Debug, PartialEq, Serialize)]
+pub struct FunctionTraceEvent {
+    pub id: usize,
+    pub parent_id: Option<usize>,
+    pub function: &'static str,
+    pub module_path: Option<&'static str>,
+    pub file: Option<&'static str>,
+    pub line: Option<u32>,
+}
+
+#[derive(Clone, Default)]
+pub struct FunctionTrace {
+    events: Arc<Mutex<Vec<FunctionTraceEvent>>>,
+    span_events: Arc<Mutex<HashMap<Id, usize>>>,
+}
+
+impl FunctionTrace {
+    pub fn dispatcher(&self) -> Dispatch {
+        Dispatch::new(
+            Registry::default().with(
+                FunctionTraceLayer {
+                    trace: self.clone(),
+                }
+                .with_filter(function_trace_filter()),
+            ),
+        )
+    }
+
+    pub fn events(&self) -> Vec<FunctionTraceEvent> {
+        self.events
+            .lock()
+            .unwrap_or_else(|error| error.into_inner())
+            .clone()
+    }
+}
+
+struct FunctionTraceLayer {
+    trace: FunctionTrace,
+}
+
+impl<S> Layer<S> for FunctionTraceLayer
+where
+    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
+{
+    fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
+        let parent_id = context.span(id).and_then(|span| {
+            let span_events = self
+                .trace
+                .span_events
+                .lock()
+                .unwrap_or_else(|error| error.into_inner());
+            span.scope()
+                .skip(1)
+                .find_map(|ancestor| span_events.get(&ancestor.id()).copied())
+        });
+        let mut events = self
+            .trace
+            .events
+            .lock()
+            .unwrap_or_else(|error| error.into_inner());
+        let event_id = events.len();
+        events.push(FunctionTraceEvent {
+            id: event_id,
+            parent_id,
+            function: attributes.metadata().name(),
+            module_path: attributes.metadata().module_path(),
+            file: attributes.metadata().file(),
+            line: attributes.metadata().line(),
+        });
+        self.trace
+            .span_events
+            .lock()
+            .unwrap_or_else(|error| error.into_inner())
+            .insert(id.clone(), event_id);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::constants::FUNCTION_TRACE_TARGET;
+
+    use super::*;
+
+    fn event(
+        id: usize,
+        parent_id: Option<usize>,
+        function: &'static str,
+    ) -> (usize, Option<usize>, &'static str) {
+        (id, parent_id, function)
+    }
+
+    fn structural_events(
+        events: &[FunctionTraceEvent],
+    ) -> Vec<(usize, Option<usize>, &'static str)> {
+        events
+            .iter()
+            .map(|event| (event.id, event.parent_id, event.function))
+            .collect()
+    }
+
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+    async fn outer() {
+        tokio::task::yield_now().await;
+        inner().await;
+    }
+
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+    async fn inner() {
+        tokio::task::yield_now().await;
+    }
+
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+    async fn concurrent_parent() {
+        tokio::join!(inner(), inner());
+    }
+
+    #[tokio::test]
+    async fn concurrent_futures_keep_separate_traces_across_yields() {
+        use tracing::instrument::WithSubscriber;
+
+        let first = FunctionTrace::default();
+        let second = FunctionTrace::default();
+        let outside = FunctionTrace::default();
+
+        async {
+            tokio::join!(
+                outer().with_subscriber(first.dispatcher()),
+                inner().with_subscriber(second.dispatcher()),
+            );
+            inner().await;
+        }
+        .with_subscriber(outside.dispatcher())
+        .await;
+
+        assert_eq!(
+            structural_events(&first.events()),
+            vec![event(0, None, "outer"), event(1, Some(0), "inner")],
+        );
+        assert_eq!(
+            structural_events(&second.events()),
+            vec![event(0, None, "inner")],
+        );
+        assert_eq!(
+            structural_events(&outside.events()),
+            vec![event(0, None, "inner")],
+        );
+    }
+
+    #[tokio::test]
+    async fn concurrent_siblings_keep_the_same_parent() {
+        use tracing::instrument::WithSubscriber;
+
+        let trace = FunctionTrace::default();
+        concurrent_parent()
+            .with_subscriber(trace.dispatcher())
+            .await;
+
+        assert_eq!(
+            structural_events(&trace.events()),
+            vec![
+                event(0, None, "concurrent_parent"),
+                event(1, Some(0), "inner"),
+                event(2, Some(0), "inner"),
+            ]
+        );
+    }
+
+    #[test]
+    fn records_matching_spans_in_creation_order() {
+        let trace = FunctionTrace::default();
+        let dispatch = trace.dispatcher();
+
+        tracing::dispatcher::with_default(&dispatch, || {
+            let _ignored = tracing::trace_span!(target: "other", "ignored");
+            let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
+            let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level");
+            let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
+        });
+
+        assert_eq!(
+            structural_events(&trace.events()),
+            vec![event(0, None, "same_name"), event(1, None, "same_name")]
+        );
+    }
+
+    #[test]
+    fn records_matching_span_nesting_depth() {
+        let trace = FunctionTrace::default();
+        let dispatch = trace.dispatcher();
+
+        tracing::dispatcher::with_default(&dispatch, || {
+            let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer");
+            let _outer_guard = outer.enter();
+            let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner");
+        });
+
+        assert_eq!(
+            structural_events(&trace.events()),
+            vec![event(0, None, "outer"), event(1, Some(0), "inner")]
+        );
+    }
+}

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

@@ -0,0 +1,59 @@
+use tracing::span::Id;
+use tracing::{Level, Metadata, Subscriber};
+use tracing_subscriber::filter::{FilterFn, LevelFilter, filter_fn};
+use tracing_subscriber::layer::Context;
+use tracing_subscriber::registry::LookupSpan;
+
+use crate::constants::FUNCTION_TRACE_TARGET;
+
+pub mod function_trace;
+
+pub use function_trace::{FunctionTrace, FunctionTraceEvent};
+
+pub fn function_trace_filter() -> FilterFn<impl Fn(&Metadata<'_>) -> bool> {
+    filter_fn(|metadata| {
+        metadata.is_span()
+            && metadata.target() == FUNCTION_TRACE_TARGET
+            && *metadata.level() == Level::TRACE
+    })
+    .with_max_level_hint(LevelFilter::TRACE)
+}
+
+pub fn span_depth<S>(context: &Context<'_, S>, id: &Id) -> usize
+where
+    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
+{
+    context
+        .span(id)
+        .map(|span| span.scope().skip(1).count())
+        .unwrap_or_default()
+}
+
+#[cfg(test)]
+mod tests {
+    use tracing::instrument::WithSubscriber;
+
+    use super::*;
+
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+    async fn instrumented_with_literal_target() {}
+
+    #[tokio::test]
+    async fn literal_instrument_target_matches_filter_constant() {
+        assert_eq!(FUNCTION_TRACE_TARGET, "litellm::function_trace");
+
+        let trace = FunctionTrace::default();
+        instrumented_with_literal_target()
+            .with_subscriber(trace.dispatcher())
+            .await;
+
+        let events = trace.events();
+        assert_eq!(events.len(), 1);
+        assert_eq!(events[0].id, 0);
+        assert_eq!(events[0].parent_id, None);
+        assert_eq!(events[0].function, "instrumented_with_literal_target");
+        assert_eq!(events[0].module_path, Some(module_path!()));
+        assert_eq!(events[0].file, Some(file!()));
+        assert!(events[0].line.is_some());
+    }
+}

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

@@ -51,6 +51,15 @@ pub trait OcrProviderConfig: Sync {
         response_json: Value,
     ) -> Result<OcrResponseData, Error>;
 
+    fn transform_ocr_response_with_params(
+        &self,
+        model: &str,
+        response_json: Value,
+        _optional_params: &Map<String, Value>,
+    ) -> Result<OcrResponseData, Error> {
+        self.transform_ocr_response(model, response_json)
+    }
+
     fn complete_url(
         &self,
         api_base: Option<&str>,

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

@@ -1,5 +1,5 @@
 use serde::{Deserialize, Serialize};
-use serde_json::Value;
+use serde_json::{Map, Value};
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub struct OcrRequestData {
@@ -14,16 +14,25 @@ pub struct OcrResponseData {
     pub document_annotation: Option<Value>,
     pub usage_info: Option<Value>,
     pub object: String,
+    pub extra_fields: Map<String, Value>,
+    pub provider_native_response: Option<Value>,
 }
 
 impl OcrResponseData {
     pub fn into_json(self) -> Value {
-        serde_json::json!({
+        let mut response = serde_json::json!({
             "pages": self.pages,
             "model": self.model,
             "document_annotation": self.document_annotation,
             "usage_info": self.usage_info,
             "object": self.object,
-        })
+        });
+        if let Value::Object(object) = &mut response {
+            object.extend(self.extra_fields);
+            if let Some(native_response) = self.provider_native_response {
+                object.insert("provider_native_response".to_string(), native_response);
+            }
+        }
+        response
     }
 }

litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs

@@ -134,6 +134,8 @@ impl OcrProviderConfig for MistralOcrConfig {
             document_annotation,
             usage_info,
             object: "ocr".to_string(),
+            extra_fields: Map::new(),
+            provider_native_response: None,
         })
     }
 

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

@@ -4,4 +4,5 @@ pub mod azure_ai;
 pub mod bedrock;
 pub mod mistral;
 pub mod openai;
+pub mod reducto;
 pub mod vertex_ai;

litellm-rust/crates/core/src/providers/reducto/mod.rs

@@ -0,0 +1 @@
+pub mod ocr;

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

@@ -0,0 +1,4 @@
+pub mod transformation;
+
+#[cfg(test)]
+mod tests;

litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs

@@ -0,0 +1,202 @@
+use rstest::{fixture, rstest};
+use serde_json::{Value, json};
+
+use super::transformation::*;
+use crate::ocr::transformation::OcrProviderConfig;
+
+#[fixture]
+fn parse_response() -> Value {
+    json!({
+        "job_id": "job_123",
+        "usage": {"num_pages": 3, "credits": 3},
+        "result": {
+            "chunks": [
+                {
+                    "content": "Page 1 block A",
+                    "blocks": [{
+                        "content": "Page 1 block A",
+                        "bbox": {"page": 1},
+                        "kind": "text",
+                    }],
+                },
+                {
+                    "content": "Page 2 block A",
+                    "blocks": [{
+                        "content": "Page 2 block A",
+                        "bbox": {"page": 2},
+                        "kind": "table",
+                    }],
+                },
+                {
+                    "content": "Page 1 block B",
+                    "blocks": [{
+                        "content": "Page 1 block B",
+                        "bbox": {"page": 1},
+                        "kind": "text",
+                    }],
+                },
+                {
+                    "content": "Page 3 block A",
+                    "blocks": [{
+                        "content": "Page 3 block A",
+                        "bbox": {"page": 3},
+                        "kind": "figure",
+                    }],
+                },
+            ],
+        },
+    })
+}
+
+#[rstest]
+fn test_parse_v3_file_upload_and_response_mapping(parse_response: Value) {
+    let source = classify_document_source("data:application/pdf;base64,JVBERi0xLjQ=")
+        .expect("PDF data URI should be valid");
+    let upload = build_upload_request(
+        source,
+        "Bearer test-key",
+        Some("https://platform.reducto.ai"),
+    )
+    .expect("data URI should require upload");
+    assert_eq!(upload.url, "https://platform.reducto.ai/upload");
+    assert_eq!(upload.authorization, "Bearer test-key");
+    assert_eq!(upload.file_name, "document");
+    assert_eq!(upload.mime_type, "application/pdf");
+    assert_eq!(upload.bytes, b"%PDF-1.4");
+
+    let optional_params = json!({
+        "formatting": {"table_output_format": "html"},
+        "retrieval": {"chunk_mode": "section"},
+        "settings": {"ocr_system": "standard"},
+    })
+    .as_object()
+    .expect("params should be an object")
+    .clone();
+    let request = build_parse_v3_request("reducto://uploaded.pdf", optional_params);
+    assert_eq!(
+        request.data,
+        json!({
+            "input": "reducto://uploaded.pdf",
+            "formatting": {"table_output_format": "html"},
+            "retrieval": {"chunk_mode": "section"},
+            "settings": {"ocr_system": "standard"},
+        })
+    );
+
+    let transformed = transform_reducto_response("parse-v3", parse_response.clone())
+        .expect("response should transform");
+    assert_eq!(
+        transformed.usage_info,
+        Some(json!({"pages_processed": 3, "credits": 3}))
+    );
+    assert_eq!(transformed.pages.len(), 3);
+    assert_eq!(
+        transformed.pages[0],
+        json!({
+            "index": 0,
+            "markdown": "Page 1 block A\n\nPage 1 block B",
+            "blocks": [
+                {"content": "Page 1 block A", "bbox": {"page": 1}, "kind": "text"},
+                {"content": "Page 1 block B", "bbox": {"page": 1}, "kind": "text"},
+            ],
+        })
+    );
+    assert_eq!(transformed.pages[1]["markdown"], "Page 2 block A");
+    assert_eq!(transformed.pages[2]["markdown"], "Page 3 block A");
+    assert_eq!(transformed.provider_native_response, Some(parse_response));
+}
+
+#[rstest]
+fn test_parse_v3_reducto_id_passthrough_skips_upload(parse_response: Value) {
+    let document = json!({
+        "type": "document_url",
+        "document_url": "reducto://already-uploaded.pdf",
+    });
+    let source = extract_document_source(&document).expect("Reducto ID should be valid");
+    assert!(build_upload_request(source.clone(), "Bearer test-key", None).is_none());
+    assert_eq!(
+        source,
+        ReductoDocumentSource::FileId("reducto://already-uploaded.pdf".to_string())
+    );
+
+    let request = REDUCTO_PARSE_V3_CONFIG
+        .transform_ocr_request(
+            "parse-v3",
+            document,
+            json!({"retrieval": {"chunk_mode": "section"}})
+                .as_object()
+                .expect("params should be object")
+                .clone(),
+        )
+        .expect("direct ID should transform");
+    assert_eq!(request.data["input"], "reducto://already-uploaded.pdf");
+    assert_eq!(request.data["retrieval"]["chunk_mode"], "section");
+
+    let response = REDUCTO_PARSE_V3_CONFIG
+        .transform_ocr_response("parse-v3", parse_response)
+        .expect("response should transform");
+    assert!(
+        response.pages[0]["markdown"]
+            .as_str()
+            .expect("markdown should be string")
+            .starts_with("Page 1 block A")
+    );
+}
+
+#[rstest]
+fn test_parse_legacy_wraps_enhance_under_options() {
+    let request = build_parse_legacy_request(
+        "reducto://legacy.pdf",
+        json!({"enhance": {"agentic": [{"type": "table"}]}})
+            .as_object()
+            .expect("params should be object"),
+    );
+    assert_eq!(
+        request.data,
+        json!({
+            "document_url": "reducto://legacy.pdf",
+            "options": {"enhance": {"agentic": [{"type": "table"}]}},
+        })
+    );
+}
+
+#[rstest]
+fn test_parse_v3_image_data_uri_upload_uses_image_mime() {
+    let source = classify_document_source("data:image/png;base64,iVBORw0KGgo=")
+        .expect("PNG data URI should be valid");
+    let upload = build_upload_request(
+        source,
+        "Bearer programmatic-key",
+        Some("https://custom.reducto.test/"),
+    )
+    .expect("data URI should require upload");
+    assert_eq!(upload.url, "https://custom.reducto.test/upload");
+    assert_eq!(upload.authorization, "Bearer programmatic-key");
+    assert_eq!(upload.mime_type, "image/png");
+    assert_eq!(upload.bytes, b"\x89PNG\r\n\x1a\n");
+}
+
+#[rstest]
+#[case::http("http://example.com/document.pdf")]
+#[case::https("https://example.com/document.pdf")]
+fn test_parse_v3_rejects_plain_http_urls(#[case] source: &str) {
+    let error = classify_document_source(source).expect_err("plain URL should be rejected");
+    assert!(error.to_string().contains("upload the file first"));
+}
+
+#[rstest]
+fn test_parse_v3_uses_programmatic_api_key_over_env() {
+    let key = resolve_api_key(Some("passed-key"), &|_| Some("env-reducto-key".to_string()))
+        .expect("explicit key should resolve");
+    assert_eq!(key, "passed-key");
+
+    let headers = REDUCTO_PARSE_V3_CONFIG
+        .validate_environment(Vec::new(), Some("passed-key"), &|_| {
+            Some("env-reducto-key".to_string())
+        })
+        .expect("headers should validate");
+    assert_eq!(
+        headers,
+        vec![("Authorization".to_string(), "Bearer passed-key".to_string())]
+    );
+}

litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs

@@ -0,0 +1,407 @@
+use std::collections::BTreeMap;
+
+use base64::Engine;
+use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
+use serde_json::{Map, Value, json};
+
+use crate::error::{Error, json_type_name};
+use crate::ocr::transformation::OcrProviderConfig;
+use crate::ocr::types::{OcrRequestData, OcrResponseData};
+
+pub const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
+pub const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
+pub const REDUCTO_ID_PREFIX: &str = "reducto://";
+
+const PARSE_V3_SUPPORTED_OCR_PARAMS: &[&str] = &["formatting", "retrieval", "settings"];
+const PARSE_LEGACY_SUPPORTED_OCR_PARAMS: &[&str] = &["enhance"];
+const MISSING_KEY_MESSAGE: &str = "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()";
+const DATA_URI_UPLOAD_REQUIRED: &str =
+    "Reducto data URI upload must complete before OCR request transformation";
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum ReductoDocumentSource {
+    FileId(String),
+    Upload { bytes: Vec<u8>, mime_type: String },
+}
+
+#[derive(Clone, PartialEq, Eq)]
+pub struct ReductoUploadRequest {
+    pub url: String,
+    pub authorization: String,
+    pub file_name: &'static str,
+    pub bytes: Vec<u8>,
+    pub mime_type: String,
+}
+
+pub struct ReductoParseV3Config;
+pub struct ReductoParseLegacyConfig;
+
+pub const REDUCTO_PARSE_V3_CONFIG: ReductoParseV3Config = ReductoParseV3Config;
+pub const REDUCTO_PARSE_LEGACY_CONFIG: ReductoParseLegacyConfig = ReductoParseLegacyConfig;
+
+pub fn config_for_model(model: &str) -> Option<&'static dyn OcrProviderConfig> {
+    match model {
+        "parse-v3" => Some(&REDUCTO_PARSE_V3_CONFIG),
+        "parse-legacy" => Some(&REDUCTO_PARSE_LEGACY_CONFIG),
+        _ => None,
+    }
+}
+
+pub fn normalize_api_base(api_base: Option<&str>) -> String {
+    api_base
+        .map(str::trim)
+        .filter(|base| !base.is_empty())
+        .unwrap_or(REDUCTO_API_BASE)
+        .trim_end_matches('/')
+        .to_string()
+}
+
+pub fn parse_url(api_base: Option<&str>) -> String {
+    format!("{}/parse", normalize_api_base(api_base))
+}
+
+pub fn upload_url(api_base: Option<&str>) -> String {
+    format!("{}/upload", normalize_api_base(api_base))
+}
+
+pub fn resolve_api_key(
+    api_key: Option<&str>,
+    env_lookup: &dyn Fn(&str) -> Option<String>,
+) -> Result<String, Error> {
+    api_key
+        .map(str::trim)
+        .filter(|key| !key.is_empty())
+        .map(str::to_string)
+        .or_else(|| {
+            env_lookup(REDUCTO_API_KEY_ENV)
+                .map(|key| key.trim().to_string())
+                .filter(|key| !key.is_empty())
+        })
+        .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
+}
+
+pub fn extract_document_source(document: &Value) -> Result<ReductoDocumentSource, Error> {
+    let document = document.as_object().ok_or_else(|| Error::InvalidType {
+        expected: "object",
+        actual: json_type_name(document),
+    })?;
+    let source = document
+        .get("document_url")
+        .and_then(Value::as_str)
+        .filter(|source| !source.is_empty())
+        .or_else(|| document.get("image_url").and_then(Value::as_str))
+        .ok_or_else(|| {
+            Error::InvalidRequest(
+                "Reducto expected OCR preprocessing to produce document_url or image_url"
+                    .to_string(),
+            )
+        })?;
+    classify_document_source(source)
+}
+
+pub fn classify_document_source(source: &str) -> Result<ReductoDocumentSource, Error> {
+    if source.starts_with(REDUCTO_ID_PREFIX) {
+        return Ok(ReductoDocumentSource::FileId(source.to_string()));
+    }
+    if source.starts_with("http://") || source.starts_with("https://") {
+        return Err(Error::InvalidRequest(
+            "Reducto requires type='file' (auto-uploaded) or a reducto:// id. Plain http(s) URLs are not supported; upload the file first."
+                .to_string(),
+        ));
+    }
+    if !source.starts_with("data:") {
+        return Err(Error::InvalidRequest(
+            "Reducto requires a reducto:// id or a base64 data URI after OCR preprocessing."
+                .to_string(),
+        ));
+    }
+
+    let (header, encoded) = source
+        .split_once(',')
+        .ok_or_else(|| Error::InvalidRequest("Invalid Reducto data URI provided.".to_string()))?;
+    if !header.split(';').any(|part| part == "base64") {
+        return Err(Error::InvalidRequest(
+            "Reducto only supports base64-encoded data URIs.".to_string(),
+        ));
+    }
+
+    let mime_type = header
+        .strip_prefix("data:")
+        .and_then(|header| header.split(';').next())
+        .filter(|mime| !mime.is_empty())
+        .unwrap_or("application/octet-stream")
+        .to_string();
+    let bytes = BASE64_STANDARD.decode(encoded).map_err(|_| {
+        Error::InvalidRequest("Invalid Reducto base64 payload provided.".to_string())
+    })?;
+
+    Ok(ReductoDocumentSource::Upload { bytes, mime_type })
+}
+
+pub fn build_upload_request(
+    source: ReductoDocumentSource,
+    authorization: &str,
+    api_base: Option<&str>,
+) -> Option<ReductoUploadRequest> {
+    let ReductoDocumentSource::Upload { bytes, mime_type } = source else {
+        return None;
+    };
+
+    Some(ReductoUploadRequest {
+        url: upload_url(api_base),
+        authorization: authorization.to_string(),
+        file_name: "document",
+        bytes,
+        mime_type,
+    })
+}
+
+pub fn extract_upload_file_id(response_json: &Value) -> Result<&str, Error> {
+    response_json
+        .as_object()
+        .and_then(|response| response.get("file_id"))
+        .and_then(Value::as_str)
+        .filter(|file_id| !file_id.is_empty())
+        .ok_or_else(|| {
+            Error::InvalidResponse(format!(
+                "Reducto /upload returned 200 without a file_id; got payload={response_json}"
+            ))
+        })
+}
+
+pub fn build_parse_v3_request(
+    file_id: &str,
+    optional_params: Map<String, Value>,
+) -> OcrRequestData {
+    let data = std::iter::once(("input".to_string(), Value::String(file_id.to_string())))
+        .chain(optional_params)
+        .collect();
+    OcrRequestData {
+        data: Value::Object(data),
+        files: None,
+    }
+}
+
+pub fn build_parse_legacy_request(
+    file_id: &str,
+    optional_params: &Map<String, Value>,
+) -> OcrRequestData {
+    let options = optional_params
+        .get("enhance")
+        .filter(|enhance| !enhance.is_null())
+        .map(|enhance| json!({"options": {"enhance": enhance}}));
+    let data = match options {
+        Some(Value::Object(options)) => std::iter::once((
+            "document_url".to_string(),
+            Value::String(file_id.to_string()),
+        ))
+        .chain(options)
+        .collect(),
+        _ => Map::from_iter([(
+            "document_url".to_string(),
+            Value::String(file_id.to_string()),
+        )]),
+    };
+    OcrRequestData {
+        data: Value::Object(data),
+        files: None,
+    }
+}
+
+fn source_file_id(document: &Value) -> Result<String, Error> {
+    match extract_document_source(document)? {
+        ReductoDocumentSource::FileId(file_id) => Ok(file_id),
+        ReductoDocumentSource::Upload { .. } => Err(Error::Unsupported(DATA_URI_UPLOAD_REQUIRED)),
+    }
+}
+
+fn page_number(block: &Map<String, Value>) -> Option<i64> {
+    let page = block.get("bbox")?.as_object()?.get("page")?;
+    page.as_i64()
+        .or_else(|| page.as_u64().and_then(|page| i64::try_from(page).ok()))
+        .or_else(|| page.as_str().and_then(|page| page.parse().ok()))
+}
+
+fn chunks(result: &Map<String, Value>) -> &[Value] {
+    result
+        .get("chunks")
+        .and_then(Value::as_array)
+        .map(Vec::as_slice)
+        .unwrap_or_default()
+}
+
+fn build_pages(result: &Map<String, Value>) -> Vec<Value> {
+    let blocks_by_page = chunks(result)
+        .iter()
+        .filter_map(Value::as_object)
+        .filter_map(|chunk| chunk.get("blocks").and_then(Value::as_array))
+        .flatten()
+        .filter_map(|block| block.as_object().map(|object| (block, object)))
+        .filter_map(|(block, object)| page_number(object).map(|page| (page, block.clone())))
+        .fold(
+            BTreeMap::<i64, Vec<Value>>::new(),
+            |mut pages, (page, block)| {
+                pages.entry(page).or_default().push(block);
+                pages
+            },
+        );
+
+    if blocks_by_page.is_empty() {
+        let markdown = chunks(result)
+            .iter()
+            .filter_map(Value::as_object)
+            .filter_map(|chunk| chunk.get("content").and_then(Value::as_str))
+            .filter(|content| !content.is_empty())
+            .collect::<Vec<_>>()
+            .join("\n\n");
+        return if markdown.is_empty() {
+            Vec::new()
+        } else {
+            vec![json!({"index": 0, "markdown": markdown})]
+        };
+    }
+
+    blocks_by_page
+        .into_iter()
+        .map(|(page, blocks)| {
+            let markdown = blocks
+                .iter()
+                .filter_map(Value::as_object)
+                .filter_map(|block| block.get("content").and_then(Value::as_str))
+                .filter(|content| !content.is_empty())
+                .collect::<Vec<_>>()
+                .join("\n\n");
+            json!({
+                "index": page.saturating_sub(1).max(0),
+                "markdown": markdown,
+                "blocks": blocks,
+            })
+        })
+        .collect()
+}
+
+pub fn transform_reducto_response(
+    model: &str,
+    response_json: Value,
+) -> Result<OcrResponseData, Error> {
+    let response = response_json
+        .as_object()
+        .ok_or_else(|| Error::InvalidType {
+            expected: "object",
+            actual: json_type_name(&response_json),
+        })?;
+    let empty_result = Map::new();
+    let result = match response.get("result") {
+        Some(Value::Object(result)) => result,
+        Some(Value::Null) => &empty_result,
+        Some(_) => {
+            return Err(Error::InvalidResponse(
+                "Reducto result must be an object".to_string(),
+            ));
+        }
+        None => response,
+    };
+    let usage = response
+        .get("usage")
+        .and_then(Value::as_object)
+        .cloned()
+        .unwrap_or_default();
+    let usage_info = Some(json!({
+        "pages_processed": usage.get("num_pages").cloned().unwrap_or(Value::Null),
+        "credits": usage.get("credits").cloned().unwrap_or(Value::Null),
+    }));
+
+    Ok(OcrResponseData {
+        pages: build_pages(result),
+        model: model.to_string(),
+        document_annotation: None,
+        usage_info,
+        object: "ocr".to_string(),
+        extra_fields: Map::new(),
+        provider_native_response: Some(response_json),
+    })
+}
+
+impl OcrProviderConfig for ReductoParseV3Config {
+    fn supported_ocr_params(&self) -> &'static [&'static str] {
+        PARSE_V3_SUPPORTED_OCR_PARAMS
+    }
+
+    fn transform_ocr_request(
+        &self,
+        _model: &str,
+        document: Value,
+        optional_params: Map<String, Value>,
+    ) -> Result<OcrRequestData, Error> {
+        let file_id = source_file_id(&document)?;
+        Ok(build_parse_v3_request(&file_id, optional_params))
+    }
+
+    fn transform_ocr_response(
+        &self,
+        model: &str,
+        response_json: Value,
+    ) -> Result<OcrResponseData, Error> {
+        transform_reducto_response(model, response_json)
+    }
+
+    fn complete_url(
+        &self,
+        api_base: Option<&str>,
+        _model: &str,
+        _optional_params: &Map<String, Value>,
+        _env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> Result<String, Error> {
+        Ok(parse_url(api_base))
+    }
+
+    fn resolve_api_key(
+        &self,
+        api_key: Option<&str>,
+        env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> Result<String, Error> {
+        resolve_api_key(api_key, env_lookup)
+    }
+}
+
+impl OcrProviderConfig for ReductoParseLegacyConfig {
+    fn supported_ocr_params(&self) -> &'static [&'static str] {
+        PARSE_LEGACY_SUPPORTED_OCR_PARAMS
+    }
+
+    fn transform_ocr_request(
+        &self,
+        _model: &str,
+        document: Value,
+        optional_params: Map<String, Value>,
+    ) -> Result<OcrRequestData, Error> {
+        let file_id = source_file_id(&document)?;
+        Ok(build_parse_legacy_request(&file_id, &optional_params))
+    }
+
+    fn transform_ocr_response(
+        &self,
+        model: &str,
+        response_json: Value,
+    ) -> Result<OcrResponseData, Error> {
+        transform_reducto_response(model, response_json)
+    }
+
+    fn complete_url(
+        &self,
+        api_base: Option<&str>,
+        _model: &str,
+        _optional_params: &Map<String, Value>,
+        _env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> Result<String, Error> {
+        Ok(parse_url(api_base))
+    }
+
+    fn resolve_api_key(
+        &self,
+        api_key: Option<&str>,
+        env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> Result<String, Error> {
+        resolve_api_key(api_key, env_lookup)
+    }
+}

litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs

@@ -212,6 +212,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
         MISTRAL_OCR_CONFIG.supported_ocr_params()
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_ocr_request(
         &self,
         model: &str,
@@ -229,6 +230,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
         MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn complete_url(
         &self,
         api_base: Option<&str>,
@@ -253,10 +255,21 @@ impl OcrProviderConfig for VertexAiOcrConfig {
 }
 
 impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn supported_ocr_params(&self) -> &'static [&'static str] {
         DEEPSEEK_SUPPORTED_OCR_PARAMS
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+    fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
+        non_default_params
+            .iter()
+            .filter(|(name, _)| DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&name.as_str()))
+            .map(|(name, value)| (name.clone(), value.clone()))
+            .collect()
+    }
+
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_ocr_request(
         &self,
         model: &str,
@@ -283,6 +296,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
         })
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_ocr_response(
         &self,
         model: &str,
@@ -335,9 +349,12 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
             document_annotation: object.get("document_annotation").cloned(),
             usage_info,
             object: "ocr".to_string(),
+            extra_fields: Map::new(),
+            provider_native_response: None,
         })
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn complete_url(
         &self,
         api_base: Option<&str>,
@@ -360,6 +377,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use rstest::rstest;
 
     #[test]
     fn vertex_mistral_url_uses_project_location_and_model() {
@@ -411,6 +429,22 @@ mod tests {
         );
     }
 
+    #[rstest]
+    #[case::bare_model("deepseek-ocr-maas")]
+    #[case::namespaced_model("deepseek-ai/deepseek-ocr-maas")]
+    fn vertex_deepseek_request_uses_single_provider_namespace(#[case] model: &str) {
+        let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG
+            .transform_ocr_request(
+                model,
+                json!({"type": "image_url", "image_url": "data:image/png;base64,AA=="}),
+                Map::new(),
+            )
+            .expect("request transforms")
+            .data;
+
+        assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
+    }
+
     #[test]
     fn vertex_deepseek_response_wraps_markdown_content() {
         let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG

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

@@ -14,11 +14,15 @@ default = ["abi3"]
 abi3 = ["pyo3/abi3-py310"]
 extension-module = ["pyo3/extension-module"]
 panic-test = []
+trace-parity = [
+    "dep:tracing",
+    "litellm-core/observability",
+    "litellm-ai-gateway/trace-parity",
+]
 
 [dependencies]
 futures-util.workspace = true
-tracing.workspace = true
-tracing-subscriber.workspace = true
+tracing = { workspace = true, optional = true }
 litellm-core = { workspace = true, features = ["bedrock-auth"] }
 litellm-ai-gateway = { workspace = true, default-features = false }
 litellm-python-interop.workspace = true
@@ -31,6 +35,7 @@ tokio.workspace = true
 [dev-dependencies]
 criterion = "0.8.2"
 tokio-tungstenite.workspace = true
+tracing.workspace = true
 
 [[bench]]
 name = "serialization"

litellm-rust/crates/python-bridge/src/constants.rs

@@ -1 +0,0 @@
-pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";

litellm-rust/crates/python-bridge/src/function_trace.rs

@@ -1,216 +1,22 @@
 use std::future::Future;
-use std::sync::{Arc, Mutex};
 
+use litellm_core::observability::{FunctionTrace, FunctionTraceEvent};
 use serde::Serialize;
 use tracing::instrument::WithSubscriber;
-use tracing::span::{Attributes, Id};
-use tracing::{Dispatch, Level, Subscriber};
-use tracing_subscriber::filter::{LevelFilter, filter_fn};
-use tracing_subscriber::layer::Context;
-use tracing_subscriber::prelude::*;
-use tracing_subscriber::registry::LookupSpan;
-use tracing_subscriber::{Layer, Registry};
-
-use crate::constants::FUNCTION_TRACE_TARGET;
 
 #[derive(Serialize)]
-#[serde(untagged)]
-pub(crate) enum TraceResponse<T> {
-    Plain(T),
-    Traced {
-        response: T,
-        trace: Vec<FunctionTraceEvent>,
-    },
+pub(crate) struct TracedResponse<T> {
+    response: T,
+    trace: Vec<FunctionTraceEvent>,
 }
 
-pub(crate) async fn trace_call<T, E>(
+pub(crate) async fn capture<T, E>(
     future: impl Future<Output = Result<T, E>>,
-    enabled: bool,
-) -> Result<TraceResponse<T>, E> {
-    if !enabled {
-        return future.await.map(TraceResponse::Plain);
-    }
+) -> Result<TracedResponse<T>, E> {
     let trace = FunctionTrace::default();
     let response = future.with_subscriber(trace.dispatcher()).await?;
-    Ok(TraceResponse::Traced {
+    Ok(TracedResponse {
         response,
         trace: trace.events(),
     })
 }
-
-#[derive(Clone, Debug, PartialEq, Serialize)]
-pub struct FunctionTraceEvent {
-    pub function: &'static str,
-    pub depth: usize,
-}
-
-#[derive(Clone, Default)]
-pub struct FunctionTrace {
-    events: Arc<Mutex<Vec<FunctionTraceEvent>>>,
-}
-
-impl FunctionTrace {
-    pub fn dispatcher(&self) -> Dispatch {
-        let filter = filter_fn(|metadata| {
-            metadata.is_span()
-                && metadata.target() == FUNCTION_TRACE_TARGET
-                && *metadata.level() == Level::TRACE
-        })
-        .with_max_level_hint(LevelFilter::TRACE);
-        Dispatch::new(
-            Registry::default().with(
-                FunctionTraceLayer {
-                    trace: self.clone(),
-                }
-                .with_filter(filter),
-            ),
-        )
-    }
-
-    pub fn events(&self) -> Vec<FunctionTraceEvent> {
-        self.events
-            .lock()
-            .unwrap_or_else(|error| error.into_inner())
-            .clone()
-    }
-}
-
-struct FunctionTraceLayer {
-    trace: FunctionTrace,
-}
-
-impl<S> Layer<S> for FunctionTraceLayer
-where
-    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
-{
-    fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
-        let depth = context
-            .span(id)
-            .map(|span| span.scope().skip(1).count())
-            .unwrap_or_default();
-        self.trace
-            .events
-            .lock()
-            .unwrap_or_else(|error| error.into_inner())
-            .push(FunctionTraceEvent {
-                function: attributes.metadata().name(),
-                depth,
-            });
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
-    async fn outer() {
-        tokio::task::yield_now().await;
-        inner().await;
-    }
-
-    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
-    async fn inner() {
-        tokio::task::yield_now().await;
-    }
-
-    #[tokio::test]
-    async fn concurrent_futures_keep_separate_traces_across_yields() {
-        use tracing::instrument::WithSubscriber;
-
-        let first = FunctionTrace::default();
-        let second = FunctionTrace::default();
-        let outside = FunctionTrace::default();
-
-        async {
-            tokio::join!(
-                outer().with_subscriber(first.dispatcher()),
-                inner().with_subscriber(second.dispatcher()),
-            );
-            inner().await;
-        }
-        .with_subscriber(outside.dispatcher())
-        .await;
-
-        assert_eq!(
-            first.events(),
-            vec![
-                FunctionTraceEvent {
-                    function: "outer",
-                    depth: 0
-                },
-                FunctionTraceEvent {
-                    function: "inner",
-                    depth: 1
-                },
-            ],
-        );
-        assert_eq!(
-            second.events(),
-            vec![FunctionTraceEvent {
-                function: "inner",
-                depth: 0
-            }],
-        );
-        assert_eq!(
-            outside.events(),
-            vec![FunctionTraceEvent {
-                function: "inner",
-                depth: 0
-            }],
-        );
-    }
-
-    #[test]
-    fn records_matching_spans_in_creation_order() {
-        let trace = FunctionTrace::default();
-        let dispatch = trace.dispatcher();
-
-        tracing::dispatcher::with_default(&dispatch, || {
-            let _ignored = tracing::trace_span!(target: "other", "ignored");
-            let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
-            let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level");
-            let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
-        });
-
-        assert_eq!(
-            trace.events(),
-            vec![
-                FunctionTraceEvent {
-                    function: "same_name",
-                    depth: 0,
-                },
-                FunctionTraceEvent {
-                    function: "same_name",
-                    depth: 0,
-                },
-            ]
-        );
-    }
-
-    #[test]
-    fn records_matching_span_nesting_depth() {
-        let trace = FunctionTrace::default();
-        let dispatch = trace.dispatcher();
-
-        tracing::dispatcher::with_default(&dispatch, || {
-            let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer");
-            let _outer_guard = outer.enter();
-            let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner");
-        });
-
-        assert_eq!(
-            trace.events(),
-            vec![
-                FunctionTraceEvent {
-                    function: "outer",
-                    depth: 0,
-                },
-                FunctionTraceEvent {
-                    function: "inner",
-                    depth: 1,
-                },
-            ]
-        );
-    }
-}

litellm-rust/crates/python-bridge/src/lib.rs

@@ -1,8 +1,8 @@
-mod constants;
 mod diagnostics;
 mod errors;
 mod execution;
-pub mod function_trace;
+#[cfg(feature = "trace-parity")]
+mod function_trace;
 mod marshal;
 mod routes;
 
@@ -115,9 +115,43 @@ mod tests {
                 .extract::<Vec<String>>()
                 .expect("module names should be strings")
                 .into_iter()
-                .filter(|name| !name.starts_with("__"))
+                .filter(|name| !name.starts_with('_'))
                 .collect();
             assert_eq!(public_names, expected);
+
+            #[cfg(not(feature = "trace-parity"))]
+            assert!(!module.hasattr("_trace").expect("module lookup should work"));
+
+            #[cfg(feature = "trace-parity")]
+            {
+                let trace = module
+                    .getattr("_trace")
+                    .expect("trace build should expose its diagnostic namespace");
+                let trace_names: Vec<String> = trace
+                    .cast::<PyModule>()
+                    .expect("trace namespace should be a module")
+                    .dict()
+                    .keys()
+                    .extract::<Vec<String>>()
+                    .expect("trace names should be strings")
+                    .into_iter()
+                    .filter(|name| !name.starts_with("__"))
+                    .collect();
+                assert_eq!(
+                    trace_names,
+                    [
+                        "ocr",
+                        "aocr",
+                        "transcription",
+                        "atranscription",
+                        "messages",
+                        "amessages",
+                        "chat_completions",
+                        "achat_completions",
+                        "gateway_messages",
+                    ]
+                );
+            }
         });
     }
 

litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs

@@ -54,16 +54,16 @@ bridge_route! {
     required = {
         model: String,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        audio: Value,
+        audio: serde_json::Value,
     },
     optional = {
         api_key: Option<String>,
         api_base: Option<String>,
         custom_llm_provider: Option<String>,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        extra_headers: Option<Value>,
+        extra_headers: Option<serde_json::Value>,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        optional_params: Option<Value>,
+        optional_params: Option<serde_json::Value>,
         timeout_seconds: Option<f64>,
     },
     prepare = prepare_transcription,

litellm-rust/crates/python-bridge/src/routes/chat_completions.rs

@@ -73,16 +73,16 @@ bridge_route! {
     required = {
         model: String,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        messages: Value,
+        messages: serde_json::Value,
     },
     optional = {
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        optional_params: Option<Value>,
+        optional_params: Option<serde_json::Value>,
         api_key: Option<String>,
         api_base: Option<String>,
         custom_llm_provider: Option<String>,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        extra_headers: Option<Value>,
+        extra_headers: Option<serde_json::Value>,
         timeout_seconds: Option<f64>,
     },
     prepare = prepare_chat_completions,

litellm-rust/crates/python-bridge/src/routes/definition.rs

@@ -20,43 +20,33 @@ macro_rules! bridge_route {
         }
 
         #[pyfunction]
-        #[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))]
+        #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
         #[allow(clippy::too_many_arguments)]
         fn $sync_name(
             py: pyo3::Python<'_>,
             $($(#[$required_attr])* $required_name: $required_type,)*
             $($(#[$optional_attr])* $optional_name: $optional_type,)*
-            trace: bool,
         ) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
             let future = $prepare($inputs {
                 $($required_name,)*
                 $($optional_name),*
             })?;
-            $crate::execution::run_sync(
-                py,
-                $crate::function_trace::trace_call(future, trace),
-                $map_error,
-            )
+            $crate::execution::run_sync(py, future, $map_error)
         }
 
         #[pyfunction]
-        #[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))]
+        #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
         #[allow(clippy::too_many_arguments)]
         fn $async_name(
             py: pyo3::Python<'_>,
             $($(#[$required_attr])* $required_name: $required_type,)*
             $($(#[$optional_attr])* $optional_name: $optional_type,)*
-            trace: bool,
         ) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
             let future = $prepare($inputs {
                 $($required_name,)*
                 $($optional_name),*
             })?;
-            $crate::execution::run_async(
-                py,
-                $crate::function_trace::trace_call(future, trace),
-                $map_error,
-            )
+            $crate::execution::run_async(py, future, $map_error)
         }
 
         pub(super) fn register(
@@ -67,6 +57,71 @@ macro_rules! bridge_route {
             $crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($async_name, module)?)?;
             Ok(())
         }
+
+        #[cfg(feature = "trace-parity")]
+        mod trace {
+            use pyo3::prelude::*;
+            use super::{$inputs, $map_error, $prepare};
+
+            #[pyfunction]
+            #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
+            #[allow(clippy::too_many_arguments)]
+            fn $sync_name(
+                py: pyo3::Python<'_>,
+                $($(#[$required_attr])* $required_name: $required_type,)*
+                $($(#[$optional_attr])* $optional_name: $optional_type,)*
+            ) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
+                let future = $prepare($inputs {
+                    $($required_name,)*
+                    $($optional_name),*
+                })?;
+                $crate::execution::run_sync(
+                    py,
+                    $crate::function_trace::capture(future),
+                    $map_error,
+                )
+            }
+
+            #[pyfunction]
+            #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
+            #[allow(clippy::too_many_arguments)]
+            fn $async_name(
+                py: pyo3::Python<'_>,
+                $($(#[$required_attr])* $required_name: $required_type,)*
+                $($(#[$optional_attr])* $optional_name: $optional_type,)*
+            ) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
+                let future = $prepare($inputs {
+                    $($required_name,)*
+                    $($optional_name),*
+                })?;
+                $crate::execution::run_async(
+                    py,
+                    $crate::function_trace::capture(future),
+                    $map_error,
+                )
+            }
+
+            pub(super) fn register(
+                module: &pyo3::Bound<'_, pyo3::types::PyModule>,
+            ) -> pyo3::PyResult<()> {
+                $crate::routes::definition::add_function(
+                    module,
+                    pyo3::wrap_pyfunction!($sync_name, module)?,
+                )?;
+                $crate::routes::definition::add_function(
+                    module,
+                    pyo3::wrap_pyfunction!($async_name, module)?,
+                )?;
+                Ok(())
+            }
+        }
+
+        #[cfg(feature = "trace-parity")]
+        pub(super) fn register_trace(
+            module: &pyo3::Bound<'_, pyo3::types::PyModule>,
+        ) -> pyo3::PyResult<()> {
+            trace::register(module)
+        }
     };
 }
 
@@ -130,20 +185,26 @@ mod tests {
         ) -> PyResult<impl Future<Output = Result<String, Error>> + Send + 'static> {
             FUTURE_DROPPED.store(false, Ordering::SeqCst);
             let drop_guard = (inputs.value == "pending").then_some(DropGuard);
-            Ok(async move {
-                let _drop_guard = drop_guard;
-                tokio::task::yield_now().await;
-                match inputs.value.as_str() {
-                    "error" => Err(Error::InvalidRequest("synthetic error".to_string())),
-                    "map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())),
-                    "panic" => panic!("synthetic panic"),
-                    "pending" => {
-                        pending::<()>().await;
-                        unreachable!()
-                    }
-                    _ => Ok(inputs.value),
+            Ok(execute_echo(inputs, drop_guard))
+        }
+
+        #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+        async fn execute_echo(
+            inputs: EchoInputs,
+            drop_guard: Option<DropGuard>,
+        ) -> Result<String, Error> {
+            let _drop_guard = drop_guard;
+            tokio::task::yield_now().await;
+            match inputs.value.as_str() {
+                "error" => Err(Error::InvalidRequest("synthetic error".to_string())),
+                "map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())),
+                "panic" => panic!("synthetic panic"),
+                "pending" => {
+                    pending::<()>().await;
+                    unreachable!()
                 }
-            })
+                _ => Ok(inputs.value),
+            }
         }
 
         fn map_error(error: Error) -> PyErr {
@@ -164,22 +225,22 @@ mod tests {
                 (
                     "ocr",
                     "aocr",
-                    "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)",
+                    "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
                 ),
                 (
                     "transcription",
                     "atranscription",
-                    "(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)",
+                    "(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
                 ),
                 (
                     "messages",
                     "amessages",
-                    "(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)",
+                    "(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)",
                 ),
                 (
                     "chat_completions",
                     "achat_completions",
-                    "(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)",
+                    "(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)",
                 ),
             ];
 
@@ -411,6 +472,32 @@ asyncio.run(exercise())
         });
     }
 
+    #[cfg(feature = "trace-parity")]
+    #[test]
+    fn diagnostic_route_returns_the_response_and_filtered_trace() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "synthetic").expect("module should be created");
+            synthetic::register_trace(&module).expect("trace routes should register");
+            let locals = PyDict::new(py);
+            locals
+                .set_item("routes", &module)
+                .expect("module should enter Python locals");
+            let code = CString::new(
+                r#"
+result = routes.echo("traced")
+assert result == {
+    "response": "traced",
+    "trace": [{"function": "execute_echo", "depth": 0}],
+}
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("diagnostic route should return its response and trace");
+        });
+    }
+
     #[test]
     fn route_registration_rejects_duplicate_python_names() {
         Python::initialize();

litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs

@@ -0,0 +1,29 @@
+use pyo3::prelude::*;
+use serde_json::Value;
+
+use crate::errors::core_error_to_pyerr;
+
+#[pyfunction]
+fn gateway_messages<'py>(
+    py: Python<'py>,
+    model_alias: String,
+    provider_model: String,
+    api_base: String,
+    #[pyo3(from_py_with = litellm_python_interop::from_py)] body: Value,
+) -> PyResult<Bound<'py, PyAny>> {
+    let future = litellm_ai_gateway::trace_parity::messages_request(
+        model_alias,
+        provider_model,
+        api_base,
+        body,
+    );
+    crate::execution::run_async(
+        py,
+        crate::function_trace::capture(future),
+        core_error_to_pyerr,
+    )
+}
+
+pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    super::definition::add_function(module, wrap_pyfunction!(gateway_messages, module)?)
+}

litellm-rust/crates/python-bridge/src/routes/messages.rs

@@ -50,14 +50,14 @@ bridge_route! {
     required = {
         model: String,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        body: Value,
+        body: serde_json::Value,
     },
     optional = {
         api_key: Option<String>,
         api_base: Option<String>,
         custom_llm_provider: Option<String>,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        extra_headers: Option<Value>,
+        extra_headers: Option<serde_json::Value>,
         timeout_seconds: Option<f64>,
     },
     prepare = prepare_messages,

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

@@ -3,6 +3,9 @@ use pyo3::prelude::*;
 #[macro_use]
 mod definition;
 
+#[cfg(feature = "trace-parity")]
+mod gateway_messages;
+
 mod audio_transcription;
 mod chat_completions;
 mod messages;
@@ -12,5 +15,16 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
     ocr::register(module)?;
     audio_transcription::register(module)?;
     messages::register(module)?;
-    chat_completions::register(module)
+    chat_completions::register(module)?;
+    #[cfg(feature = "trace-parity")]
+    {
+        let trace = PyModule::new(module.py(), "_trace")?;
+        ocr::register_trace(&trace)?;
+        audio_transcription::register_trace(&trace)?;
+        messages::register_trace(&trace)?;
+        chat_completions::register_trace(&trace)?;
+        gateway_messages::register_trace(&trace)?;
+        module.add_submodule(&trace)?;
+    }
+    Ok(())
 }

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

@@ -56,16 +56,16 @@ bridge_route! {
     required = {
         model: String,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        document: Value,
+        document: serde_json::Value,
     },
     optional = {
         api_key: Option<String>,
         api_base: Option<String>,
         custom_llm_provider: Option<String>,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        extra_headers: Option<Value>,
+        extra_headers: Option<serde_json::Value>,
         #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        optional_params: Option<Value>,
+        optional_params: Option<serde_json::Value>,
         timeout_seconds: Option<f64>,
     },
     prepare = prepare_ocr,

litellm-rust/crates/python-bridge/src/routes/runtime.rs

@@ -1,423 +0,0 @@
-use std::future::Future;
-use std::panic::AssertUnwindSafe;
-use std::time::Duration;
-
-use futures_util::FutureExt;
-use litellm_core::error::Error;
-use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
-use pyo3::exceptions::PyRuntimeError;
-use pyo3::prelude::*;
-use serde::Serialize;
-use tokio::runtime::{Handle, Runtime};
-use tokio::time::{self, MissedTickBehavior};
-
-pub(super) fn run_sync<T, F>(
-    py: Python<'_>,
-    future: F,
-    map_error: fn(Error) -> PyErr,
-) -> PyResult<Py<PyAny>>
-where
-    T: Serialize + Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
-{
-    run_sync_on(
-        py,
-        pyo3_async_runtimes::tokio::get_runtime(),
-        future,
-        map_error,
-    )
-}
-
-fn run_sync_on<T, F>(
-    py: Python<'_>,
-    runtime: &Runtime,
-    future: F,
-    map_error: fn(Error) -> PyErr,
-) -> PyResult<Py<PyAny>>
-where
-    T: Serialize + Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
-{
-    if Handle::try_current().is_ok() {
-        return Err(PyRuntimeError::new_err(
-            "synchronous native routes cannot run from a Tokio context; use the async route",
-        ));
-    }
-
-    let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
-    let result = map_core_result(result, map_error)?;
-    Pythonized(result).into_pyobject(py).map(Bound::unbind)
-}
-
-pub(super) fn run_async<T, F>(
-    py: Python<'_>,
-    future: F,
-    map_error: fn(Error) -> PyErr,
-) -> PyResult<Bound<'_, PyAny>>
-where
-    T: Serialize + Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
-{
-    pyo3_async_runtimes::tokio::future_into_py(py, async move {
-        let result = catch_route_panic(future).await?;
-        let result = map_core_result(result, map_error)?;
-        Ok(Pythonized(result))
-    })
-}
-
-fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
-    match result {
-        Ok(value) => Ok(value),
-        Err(error) => Err(
-            std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error)))
-                .map_err(panic_to_pyerr)?,
-        ),
-    }
-}
-
-async fn catch_route_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
-where
-    F: Future<Output = Result<T, Error>>,
-{
-    AssertUnwindSafe(future)
-        .catch_unwind()
-        .await
-        .map_err(panic_to_pyerr)
-}
-
-async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
-where
-    F: Future<Output = Result<T, Error>>,
-{
-    let future = catch_route_panic(future);
-    tokio::pin!(future);
-
-    let signal_interval = Duration::from_millis(50);
-    let mut signal_checks =
-        time::interval_at(time::Instant::now() + signal_interval, signal_interval);
-    signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay);
-    loop {
-        tokio::select! {
-            result = &mut future => return result,
-            _ = signal_checks.tick() => Python::attach(|py| py.check_signals())?,
-        }
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use std::ffi::CString;
-    use std::future::poll_fn;
-    use std::sync::atomic::{AtomicUsize, Ordering};
-    use std::sync::{Arc, mpsc};
-    use std::task::Poll;
-    use std::thread;
-    use std::time::Instant;
-
-    use pyo3::panic::PanicException;
-    use pyo3::types::{PyDict, PyModule};
-    use serde::Serializer;
-    use tokio::runtime::Builder;
-
-    use super::*;
-
-    fn runtime_error(error: Error) -> PyErr {
-        PyRuntimeError::new_err(error.to_string())
-    }
-
-    fn panicking_error_mapper(_error: Error) -> PyErr {
-        panic!("error mapper panicked")
-    }
-
-    struct PanickingOutput;
-
-    static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0);
-
-    impl Serialize for PanickingOutput {
-        fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
-        where
-            S: Serializer,
-        {
-            panic!("serializer panicked")
-        }
-    }
-
-    #[pyfunction]
-    fn async_serialization_panic(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
-        run_async(py, async { Ok(PanickingOutput) }, runtime_error)
-    }
-
-    #[pyfunction]
-    fn async_runtime_probe(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
-        run_async(
-            py,
-            async {
-                ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst);
-                Ok(true)
-            },
-            runtime_error,
-        )
-    }
-
-    #[pyfunction]
-    fn runtime_worker_count() -> usize {
-        pyo3_async_runtimes::tokio::get_runtime()
-            .metrics()
-            .num_workers()
-    }
-
-    #[pyfunction]
-    fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
-        let completion_deadline = Instant::now() + Duration::from_secs(2);
-        while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
-            if Instant::now() >= completion_deadline {
-                return false;
-            }
-            thread::sleep(Duration::from_millis(1));
-        }
-
-        let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
-        pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
-            let _ = heartbeat_tx.send(());
-        });
-        heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
-    }
-
-    fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
-        result
-            .expect("route should complete")
-            .bind(py)
-            .extract()
-            .expect("result should convert")
-    }
-
-    #[test]
-    fn sync_runner_polls_future_on_the_caller_thread() {
-        Python::initialize();
-        Python::attach(|py| {
-            let caller_thread = std::thread::current().id();
-            let result = run_sync(
-                py,
-                async move { Ok(std::thread::current().id() == caller_thread) },
-                runtime_error,
-            );
-
-            assert!(extract_bool(py, result));
-        });
-    }
-
-    #[test]
-    fn sync_runner_releases_gil_while_waiting() {
-        Python::initialize();
-        Python::attach(|py| {
-            let result = run_sync(
-                py,
-                async {
-                    let gil_acquired = tokio::time::timeout(
-                        Duration::from_secs(2),
-                        tokio::task::spawn_blocking(|| Python::attach(|_| true)),
-                    )
-                    .await;
-                    Ok(matches!(gil_acquired, Ok(Ok(true))))
-                },
-                runtime_error,
-            );
-
-            assert!(extract_bool(py, result));
-        });
-    }
-
-    #[test]
-    fn sync_runner_rejects_calls_from_a_tokio_context() {
-        Python::initialize();
-        let runtime = Builder::new_current_thread()
-            .enable_all()
-            .build()
-            .expect("runtime should build");
-
-        let error = runtime.block_on(async {
-            Python::attach(|py| {
-                run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
-                    .expect_err("sync route should reject a nested Tokio runtime")
-            })
-        });
-
-        assert_eq!(
-            error.to_string(),
-            "RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route"
-        );
-    }
-
-    #[test]
-    fn sync_runner_can_drive_a_current_thread_runtime() {
-        Python::initialize();
-        let runtime = Builder::new_current_thread()
-            .enable_all()
-            .build()
-            .expect("runtime should build");
-        Python::attach(|py| {
-            let result = run_sync_on(
-                py,
-                &runtime,
-                async {
-                    tokio::task::yield_now().await;
-                    Ok(true)
-                },
-                runtime_error,
-            );
-            assert!(extract_bool(py, result));
-        });
-    }
-
-    #[test]
-    fn sync_runner_maps_a_panicked_future() {
-        Python::initialize();
-        Python::attach(|py| {
-            let error = run_sync::<bool, _>(
-                py,
-                poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
-                runtime_error,
-            )
-            .expect_err("panicked route should become a Python exception");
-
-            assert!(error.is_instance_of::<PanicException>(py));
-            assert_eq!(error.to_string(), "PanicException: route future panicked");
-        });
-    }
-
-    #[test]
-    fn sync_runner_maps_a_panicked_error_mapper() {
-        Python::initialize();
-        Python::attach(|py| {
-            let error = run_sync::<bool, _>(
-                py,
-                async { Err(Error::InvalidRequest("invalid".to_string())) },
-                panicking_error_mapper,
-            )
-            .expect_err("panicked mapper should become a Python exception");
-
-            assert!(error.is_instance_of::<PanicException>(py));
-            assert_eq!(error.to_string(), "PanicException: error mapper panicked");
-        });
-    }
-
-    #[test]
-    fn sync_runner_surfaces_serializer_panics() {
-        Python::initialize();
-        Python::attach(|py| {
-            let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error)
-                .expect_err("serializer panic should become a Python exception");
-
-            assert!(error.is_instance_of::<PanicException>(py));
-            assert_eq!(error.to_string(), "PanicException: serializer panicked");
-        });
-    }
-
-    #[test]
-    fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
-        Python::initialize();
-        let barrier = Arc::new(tokio::sync::Barrier::new(2));
-        let callers: Vec<_> = (0..2)
-            .map(|_| {
-                let barrier = Arc::clone(&barrier);
-                thread::spawn(move || {
-                    Python::attach(|py| {
-                        extract_bool(
-                            py,
-                            run_sync(
-                                py,
-                                async move {
-                                    Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait())
-                                        .await
-                                        .is_ok())
-                                },
-                                runtime_error,
-                            ),
-                        )
-                    })
-                })
-            })
-            .collect();
-        let results: Vec<_> = callers
-            .into_iter()
-            .map(|caller| caller.join().expect("caller should not panic"))
-            .collect();
-
-        assert_eq!(results, vec![true, true]);
-    }
-
-    #[test]
-    fn async_runner_surfaces_serializer_panics() {
-        Python::initialize();
-        Python::attach(|py| {
-            let module = PyModule::new(py, "runtime").expect("module should be created");
-            module
-                .add_function(
-                    wrap_pyfunction!(async_serialization_panic, &module)
-                        .expect("function should wrap"),
-                )
-                .expect("function should register");
-            let locals = PyDict::new(py);
-            locals
-                .set_item("runtime", &module)
-                .expect("module should enter Python locals");
-            let code = CString::new(
-                r#"
-import asyncio
-
-async def exercise():
-    try:
-        await runtime.async_serialization_panic()
-    except BaseException as error:
-        assert type(error).__name__ == "PanicException"
-        assert str(error) == "serializer panicked"
-    else:
-        raise AssertionError("serializer panic was not raised")
-
-asyncio.run(exercise())
-"#,
-            )
-            .expect("Python source should not contain null bytes");
-            py.run(&code, Some(&locals), Some(&locals))
-                .expect("serializer panic should reach the Python awaiter");
-        });
-    }
-
-    #[test]
-    fn async_result_delivery_does_not_stall_tokio_workers() {
-        Python::initialize();
-        ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst);
-        Python::attach(|py| {
-            let module = PyModule::new(py, "runtime").expect("module should be created");
-            for function in [
-                wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"),
-                wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"),
-                wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"),
-            ] {
-                module
-                    .add_function(function)
-                    .expect("function should register");
-            }
-            let locals = PyDict::new(py);
-            locals
-                .set_item("runtime", &module)
-                .expect("module should enter Python locals");
-            let code = CString::new(
-                r#"
-import asyncio
-
-async def exercise():
-    worker_count = runtime.runtime_worker_count()
-    awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)]
-    assert runtime.runtime_is_responsive(worker_count)
-    assert await asyncio.gather(*awaitables) == [True] * worker_count
-
-asyncio.run(exercise())
-"#,
-            )
-            .expect("Python source should not contain null bytes");
-            py.run(&code, Some(&locals), Some(&locals))
-                .expect("result delivery should leave Tokio workers responsive");
-        });
-    }
-}

litellm/anthropic_interface/exceptions/exceptions.py

@@ -1,8 +1,9 @@
 """Anthropic error format type definitions."""
 
+from collections.abc import Mapping
 from typing import Literal
 
-from typing_extensions import Required, TypedDict
+from typing_extensions import NotRequired, ReadOnly, Required, TypedDict
 
 # Known Anthropic error types
 # Source: https://docs.anthropic.com/en/api/errors
@@ -23,6 +24,7 @@ class AnthropicErrorDetail(TypedDict):
 
     type: AnthropicErrorType
     message: str
+    provider_specific_fields: NotRequired[ReadOnly[Mapping[str, object]]]
 
 
 class AnthropicErrorResponse(TypedDict, total=False):

litellm/caching/caching.py

@@ -100,6 +100,7 @@ class Cache:
         qdrant_semantic_cache_vector_size: int | None = None,
         semantic_cache_embedding_max_input_tokens: int | None = None,
         semantic_cache_embedding_timeout: float | None = None,
+        semantic_cache_scope: str = SemanticCacheScope.KEY.value,
         # GCP IAM authentication parameters
         gcp_service_account: str | None = None,
         gcp_ssl_ca_certs: str | None = None,
@@ -127,6 +128,7 @@ class Cache:
             similarity_threshold (float, optional): The similarity threshold for semantic-caching, Required if type is "redis-semantic" or "qdrant-semantic".
             semantic_cache_embedding_max_input_tokens (int, optional): Truncate prompts to this many tokens before embedding them for semantic caching. Defaults to the embedding deployment's configured max_input_tokens.
             semantic_cache_embedding_timeout (float, optional): Seconds a semantic-cache lookup may spend embedding the prompt before it gives up and lets the request continue to the LLM. Defaults to SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS.
+            semantic_cache_scope (str, optional): "key" isolates semantic-cache buckets per key/team/org. "end_user" additionally isolates per end user (falls back to the key scope when the request carries no end-user id). Defaults to "key".
 
             # Disk Cache Args
             disk_cache_dir (str, optional): The directory for the disk cache. Defaults to None.
@@ -274,6 +276,7 @@ class Cache:
         self.redis_flush_size = redis_flush_size
         self.ttl = ttl
         self.mode: CacheMode = mode or CacheMode.default_on
+        self.semantic_cache_scope: str = SemanticCacheScope(semantic_cache_scope).value
 
         if self.type == LiteLLMCacheType.LOCAL and default_in_memory_ttl is not None:
             self.ttl = default_in_memory_ttl
@@ -301,6 +304,7 @@ class Cache:
         "user_api_key_team_id",
         "user_api_key_org_id",
     )
+    _SEMANTIC_CACHE_END_USER_SCOPE_FIELD: Final = "user_api_key_end_user_id"
 
     def _is_semantic_cache(self) -> bool:
         return self.type in (
@@ -309,19 +313,21 @@ class Cache:
             LiteLLMCacheType.VALKEY_SEMANTIC,
         )
 
-    def _get_semantic_cache_tenant_scope(self, kwargs: dict) -> str:
-        metadata: Final[dict] = kwargs.get("metadata") or {}
-        litellm_params: Final[dict] = kwargs.get("litellm_params") or {}
-        metadata_in_litellm_params: Final[dict] = litellm_params.get("metadata") or {}
+    def _semantic_cache_scope_fields(self) -> tuple[str, ...]:
+        if self.semantic_cache_scope == SemanticCacheScope.END_USER:
+            return (*self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS, self._SEMANTIC_CACHE_END_USER_SCOPE_FIELD)
+        return self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS
 
-        scope = ""
-        for field in self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS:
-            value = metadata.get(field)
-            if value is None:
-                value = metadata_in_litellm_params.get(field)
-            if value is not None:
-                scope += f"{field}: {value}"
-        return scope
+    def _get_semantic_cache_tenant_scope(self, kwargs: dict) -> str:
+        litellm_params: Final[dict] = kwargs.get("litellm_params") or {}
+        metadata_sources: Final[tuple[dict, ...]] = tuple(
+            source.get(key) or {} for source in (kwargs, litellm_params) for key in ("metadata", "litellm_metadata")
+        )
+        scope_values: Final = (
+            (field, next((source[field] for source in metadata_sources if source.get(field) is not None), None))
+            for field in self._semantic_cache_scope_fields()
+        )
+        return "".join(f"{field}: {value}" for field, value in scope_values if value is not None)
 
     def get_cache_key(self, **kwargs) -> str:
         """

litellm/caching/redis_cluster_node_isolation.py

@@ -19,13 +19,13 @@ connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-cov
 retry-exhaustion) is unchanged from upstream, since those already carry real evidence the
 topology changed.
 
-redis-py 8.x fixed this upstream with gentler machinery than this override's
-``node.disconnect()`` (which also kills connections other coroutines are mid-operation
-on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per
-killed connection): it marks in-use connections for reconnect only after their current
-operation completes, disconnects only the idle pooled ones, and defers reinitialization
-to the outer retry loop. When the installed ``ClusterNode`` has that per-connection
-recovery API, the factory returns the base ``RedisCluster`` unmodified.
+redis-py 8.x recovers connections per-connection, so the copied override is not used. Upstream
+still flips the shared ``_initialize`` flag on any node's timeout, funneling every concurrent
+caller through the reinit lock and, if ``CLUSTER SLOTS`` lands on the slow node, into a full
+teardown. For those versions the factory returns a thin wrapper around upstream's
+``_execute_command`` that clears the flag again after an isolated timeout (a ConnectionError,
+a third consecutive timeout on the same node, or a concurrent request from any other command
+or ``aclose()`` still reinits).
 """
 
 import asyncio
@@ -44,6 +44,8 @@ class _ClusterNodeAttrs(Protocol):
     mode; typing ``target_node`` as this Protocol at the one boundary keeps the override's
     own logic fully typed without a banned ``typing.cast``."""
 
+    name: str
+
     async def execute_command(
         self,
         *args: object,
@@ -78,18 +80,20 @@ class _ClusterAttrs(Protocol):
 #: this override can't see (Python won't error -- it'll just run our now-stale copy), so
 #: construction logs a loud warning rather than silently trusting an unverified copy.
 _VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"})
+_CONSECUTIVE_TIMEOUTS_BEFORE_REINIT: Final = 3
 
 
-def get_litellm_async_redis_cluster_class(
+def get_litellm_async_redis_cluster_class(  # noqa: C901  # supports redis-py version-specific cluster implementations
     cluster_node_class: type | None = None,
+    base_cluster_class: type | None = None,
 ) -> type["_AsyncRedisClusterType"]:
-    """Returns the base ``RedisCluster`` when the installed redis-py already recovers a
-    node-level connection error per-connection (8.x+), else builds the ``RedisCluster``
-    subclass with the per-node isolation fix for older versions whose upstream branch
-    tears down the whole cluster client.
+    """Returns a timeout-tolerant ``RedisCluster`` subclass when installed redis-py already
+    recovers node-level connections per-connection (8.x+), else builds the ``RedisCluster``
+    subclass with the per-node isolation fix for older versions whose upstream branch tears
+    down the whole cluster client.
 
-    ``cluster_node_class`` exists for dependency injection in tests; production callers
-    leave it unset and the installed ``ClusterNode`` is used.
+    ``cluster_node_class`` and ``base_cluster_class`` exist for dependency injection in tests;
+    production callers leave them unset and the installed redis-py classes are used.
 
     Imported lazily because this module is reachable from a base ``import litellm`` while
     redis is not a base dependency. Cheap to call repeatedly: the underlying redis
@@ -118,13 +122,68 @@ def get_litellm_async_redis_cluster_class(
     from redis.exceptions import TimeoutError as _RedisTimeoutError
 
     node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode
+    base_class: Final = base_cluster_class if base_cluster_class is not None else _BaseAsyncRedisCluster
     if hasattr(node_class, "update_active_connections_for_reconnect"):
         verbose_logger.debug(
-            "redis-py %s recovers a node-level connection error per-connection upstream; "
-            "using the base RedisCluster without litellm's node-isolation override.",
+            "redis-py %s recovers node connections per-connection upstream; using "
+            "LiteLLM's timeout-tolerant RedisCluster wrapper.",
             redis.__version__,
         )
-        return _BaseAsyncRedisCluster
+
+        class LiteLLMAsyncRedisClusterTimeoutTolerant(
+            base_class  # pyright: ignore[reportGeneralTypeIssues, reportUntypedBaseClass]  # the injected base class is selected at runtime
+        ):
+            def __init__(
+                self,
+                *args: object,
+                **kwargs: object,  # kwargs-ok: passes redis-py's constructor kwargs through untouched
+            ) -> None:
+                self._litellm_initialize = False
+                self._litellm_reinit_requests = 0
+                self._litellm_tolerated_timeouts = 0
+                super().__init__(*args, **kwargs)
+                self._litellm_consecutive_timeouts: dict[  # mutable-ok: per-node counter updated on the command hot path
+                    str, int
+                ] = {}
+
+            @property
+            def _initialize(self) -> bool:
+                return self._litellm_initialize
+
+            @_initialize.setter
+            def _initialize(self, value: bool) -> None:
+                if value:
+                    self._litellm_reinit_requests += 1
+                self._litellm_initialize = value
+
+            async def _execute_command(
+                self,
+                target_node: _ClusterNodeAttrs,
+                *args: object,
+                **kwargs: object,  # kwargs-ok: matches redis-py's own command dispatch signature
+            ) -> object:
+                outstanding_before: Final = self._litellm_reinit_requests - self._litellm_tolerated_timeouts
+                pending_before: Final = self._litellm_initialize
+                try:
+                    result: Final = await super()._execute_command(target_node, *args, **kwargs)
+                except _RedisTimeoutError:
+                    timeouts: Final = self._litellm_consecutive_timeouts.get(target_node.name, 0) + 1
+                    if timeouts >= _CONSECUTIVE_TIMEOUTS_BEFORE_REINIT:
+                        self._litellm_consecutive_timeouts.pop(target_node.name, None)
+                        raise
+                    self._litellm_consecutive_timeouts[target_node.name] = timeouts
+                    self._litellm_tolerated_timeouts += 1
+                    if (
+                        not pending_before
+                        and self._litellm_reinit_requests - self._litellm_tolerated_timeouts == outstanding_before
+                    ):
+                        self._initialize = False
+                    raise
+                if self._litellm_consecutive_timeouts:
+                    self._litellm_consecutive_timeouts.pop(target_node.name, None)
+                return result
+
+        return LiteLLMAsyncRedisClusterTimeoutTolerant
 
     if redis.__version__ not in _VERIFIED_REDIS_VERSIONS:
         verbose_logger.warning(

litellm/constants.py

@@ -151,6 +151,7 @@ DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_
 MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60"))
 MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200"))
 MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600"))
+MCP_SSO_ASSERTION_CACHE_TTL_SECONDS: Final = int(os.getenv("MCP_SSO_ASSERTION_CACHE_TTL_SECONDS", "60"))
 
 # Default npm cache directory for STDIO MCP servers.
 # npm/npx needs a writable cache dir; in containers the default (~/.npm)

litellm/integrations/azure_storage/azure_storage.py

@@ -36,7 +36,7 @@ AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default"
 def _cached_credential_chain_token_provider() -> Callable[[], str]:
     return get_azure_ad_token_provider(
         azure_scope=AZURE_STORAGE_TOKEN_SCOPE,
-        azure_credential=AzureCredentialType.DefaultAzureCredential,
+        azure_credential=AzureCredentialType.DeploymentIdentityCredential,
     )
 
 

litellm/litellm_core_utils/litellm_logging.py

@@ -10,7 +10,7 @@ import subprocess
 import sys
 import time
 import traceback
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Callable, Iterator, Mapping, Sequence
 from datetime import datetime as dt_object
 from functools import lru_cache
 from types import MappingProxyType, TracebackType
@@ -5878,14 +5878,28 @@ def _get_status_fields(
     #########################################################
     # Map - guardrail_information.guardrail_status to guardrail_status
     #########################################################
-    guardrail_status: GuardrailStatus = "not_run"
-    if guardrail_information and isinstance(guardrail_information, list):
-        for information in guardrail_information:
-            if isinstance(information, dict):
-                raw_status = information.get("guardrail_status", "not_run")
-                if raw_status != "not_run":
-                    guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run")
-                    break
+    # Severity order, least severe first. The status aggregates across ALL
+    # guardrail entries rather than taking the first non-"not_run" one: a
+    # pre_call guardrail that passed (e.g. a mask) records its entry before a
+    # later guardrail's block, and first-wins would report a blocked request
+    # as "success".
+    GUARDRAIL_STATUS_SEVERITY: Final[tuple[GuardrailStatus, ...]] = (
+        "not_run",
+        "success",
+        "guardrail_failed_to_respond",
+        "guardrail_intervened",
+    )
+    entries: Final[Sequence[object]] = guardrail_information if isinstance(guardrail_information, list) else ()
+    raw_statuses: Final[Iterator[object]] = (
+        entry.get("guardrail_status", "not_run") for entry in entries if isinstance(entry, dict)
+    )
+    # A guardrail is free to write any value here, and an unhashable one would
+    # raise TypeError on the mapping lookup and drop the whole payload.
+    guardrail_status: Final[GuardrailStatus] = max(
+        (GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") for raw_status in raw_statuses if isinstance(raw_status, str)),
+        key=GUARDRAIL_STATUS_SEVERITY.index,
+        default="not_run",
+    )
 
     return StandardLoggingPayloadStatusFields(llm_api_status=llm_api_status, guardrail_status=guardrail_status)
 

litellm/llms/anthropic/experimental_pass_through/messages/utils.py

@@ -6,6 +6,7 @@ from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams
 from litellm.types.llms.anthropic_messages.anthropic_response import (
     AnthropicMessagesResponse,
 )
+from litellm.types.llms.openai import ChatCompletionSystemMessage
 
 if TYPE_CHECKING:
     from litellm.exceptions import ContentPolicyViolationError
@@ -36,6 +37,16 @@ def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "
     )
 
 
+def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMessage | None:
+    """
+    Return the Anthropic Messages top-level ``system`` (a string or a list of text
+    blocks) as an OpenAI-style system message, or None when the request has none.
+    """
+    if not isinstance(system, (str, list)) or not system:
+        return None
+    return ChatCompletionSystemMessage(role="system", content=system)
+
+
 @lru_cache(maxsize=1)
 def _anthropic_messages_optional_param_keys() -> frozenset[str]:
     """

litellm/llms/snowflake/chat/transformation.py

@@ -8,23 +8,31 @@ Routes to native Cortex REST API endpoints based on model:
 Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api
 """
 
+import copy
 import json
+import re
 from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict
 
 import httpx
 from typing_extensions import ReadOnly
 
-from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
+from litellm.litellm_core_utils.prompt_templates.factory import (
+    anthropic_process_openai_file_message,
+    convert_to_anthropic_tool_result,
+    create_anthropic_image_param,
+    select_anthropic_content_block_type_for_file,
+)
+from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser
+from litellm.llms.anthropic.chat.transformation import AnthropicConfig
+from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload
+from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolMessage
 from litellm.types.utils import (
-    ChatCompletionMessageToolCall,
-    ChatCompletionUsageBlock,
     Choices,
-    Function,
     GenericStreamingChunk,
     Message,
     ModelResponse,
-    Usage,
+    ModelResponseStream,
 )
 
 from ...base_llm.base_model_iterator import BaseModelResponseIterator
@@ -93,6 +101,103 @@ def _is_claude_model(model: str) -> bool:
     return any(name.startswith(p) for p in _CLAUDE_MODEL_PREFIXES)
 
 
+def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object:
+    """One OpenAI ``image_url`` block in the native shape Cortex accepts.
+
+    Cortex documents base64 sources only, so remote URLs are inlined the way every
+    other base64-only Anthropic dialect (Bedrock invoke, Vertex) inlines them, and
+    pdf/text data URIs become document blocks rather than malformed image blocks.
+    """
+    image_url: Final = block.get("image_url")
+    url: Final = image_url if isinstance(image_url, str) else _image_url_field(image_url, "url")
+    if not url:
+        return block
+
+    converted: Final = (
+        anthropic_process_openai_file_message({"type": "file", "file": {"file_data": url}})
+        if select_anthropic_content_block_type_for_file(_data_uri_media_type(url)) == "document"
+        else create_anthropic_image_param(
+            image_url if isinstance(image_url, dict) else url,  # mutable-ok: caller's JSON block
+            format=_image_url_field(image_url, "format"),
+            is_bedrock_invoke=True,
+        )
+    )
+    cache_control: Final = block.get("cache_control")
+    if cache_control is None:
+        return converted
+    return {**converted, "cache_control": cache_control}  # mutable-ok: JSON wire block
+
+
+def _image_url_field(image_url: object, key: str) -> str | None:
+    value: Final = image_url.get(key) if isinstance(image_url, dict) else None
+    return value if isinstance(value, str) else None
+
+
+def _data_uri_media_type(url: str) -> str:
+    match: Final = re.match(r"data:([^;,]+)", url)
+    return match.group(1) if match else ""
+
+
+def _convert_image_url_blocks_to_anthropic(content: object) -> object:
+    if not isinstance(content, list):
+        return content
+    return [  # mutable-ok: JSON wire blocks
+        _convert_image_url_to_anthropic(block)
+        if isinstance(block, Mapping) and block.get("type") == "image_url"
+        else block
+        for block in content
+    ]
+
+
+def _convert_tool_result_to_anthropic(
+    content: object, tool_call_id: str, cache_control: object
+) -> Mapping[str, object]:
+    """The Anthropic ``tool_result`` block for one OpenAI tool message.
+
+    Delegating to the shared converter keeps image, document and per-block cache
+    breakpoints identical to every other Anthropic dialect; only the plain-string
+    and non-list shapes it does not model are handled here.
+    """
+    if not isinstance(content, list):
+        plain: Final[dict[str, object]] = {  # mutable-ok: JSON wire block
+            "type": "tool_result",
+            "tool_use_id": tool_call_id,
+            "content": content if isinstance(content, str) else json.dumps(content),
+        }
+        return {**plain, "cache_control": cache_control} if cache_control is not None else plain
+    converted: Final = convert_to_anthropic_tool_result(
+        ChatCompletionToolMessage(role="tool", tool_call_id=tool_call_id, content=content),
+        force_base64=True,
+    )
+    if cache_control is None:
+        return converted
+    return {**converted, "cache_control": cache_control}  # mutable-ok: JSON wire block
+
+
+def _signed_thinking_blocks(msg: object) -> list[dict[str, object]]:  # mutable-ok: JSON wire blocks
+    """The assistant turn's thinking blocks that can legally be echoed back.
+
+    Only signed blocks round-trip: Cortex rejects a thinking block whose signature is
+    missing, which is what an unsigned block from a non-thinking turn would produce.
+    """
+    blocks: Final = msg.get("thinking_blocks") if isinstance(msg, dict) else getattr(msg, "thinking_blocks", None)
+    if not isinstance(blocks, list):
+        return []  # mutable-ok: JSON wire blocks
+    return [  # mutable-ok: JSON wire blocks
+        dict(block)
+        for block in blocks
+        if isinstance(block, Mapping) and (block.get("signature") or block.get("type") == "redacted_thinking")
+    ]
+
+
+def _clean_input_schema(schema: object) -> object:  # mutable-ok: JSON schema copy
+    return (
+        {key: value for key, value in schema.items() if key != "$schema"}
+        if isinstance(schema, Mapping)
+        else schema  # mutable-ok: JSON schema copy
+    )  # mutable-ok: JSON schema copy
+
+
 class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
     """
     Snowflake Cortex REST API — unified provider.
@@ -178,7 +283,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
                 if "description" in func:
                     anthropic_tool["description"] = func["description"]
                 if "parameters" in func:
-                    anthropic_tool["input_schema"] = func["parameters"]
+                    anthropic_tool["input_schema"] = _clean_input_schema(func["parameters"])
                 else:
                     anthropic_tool["input_schema"] = {
                         "type": "object",
@@ -186,10 +291,16 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
                     }
                 anthropic_tools.append(anthropic_tool)
             else:
-                anthropic_tools.append(tool)
+                anthropic_tools.append(
+                    {**tool, "input_schema": _clean_input_schema(tool["input_schema"])}  # mutable-ok: JSON wire tool
+                    if "input_schema" in tool
+                    else tool
+                )
         return anthropic_tools
 
-    def _extract_system_and_messages(self, messages: list[AllMessageValues]) -> tuple[str | None, list[dict]]:
+    def _extract_system_and_messages(  # mutable-ok: JSON wire messages
+        self, messages: list[AllMessageValues]
+    ) -> tuple[list[dict] | None, list[dict]]:
         """
         Split messages into system prompt and conversation turns for Anthropic format.
 
@@ -197,26 +308,39 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
         - assistant messages with tool_calls → tool_use content blocks
         - tool role messages → user role with tool_result content blocks
         """
-        system_parts: Final[list[str]] = []
-        conversation: Final[list[dict]] = []
+        system_parts: Final[list[dict]] = []  # mutable-ok: JSON wire messages
+        conversation: Final[list[dict]] = []  # mutable-ok: JSON wire messages
 
         for msg in messages:
             if isinstance(msg, dict):
                 role = msg.get("role", "")
                 content: Any = msg.get("content", "")
+                msg_cache_control: object = msg.get("cache_control")
             else:
                 role = getattr(msg, "role", "")
                 content = getattr(msg, "content", "")
+                msg_cache_control = getattr(msg, "cache_control", None)
 
             if role == "system":
                 if isinstance(content, str) and content:
-                    system_parts.append(content)
+                    system_parts.append({"type": "text", "text": content})  # mutable-ok: JSON wire system block
                 elif isinstance(content, list):
-                    system_parts.append("\n".join(b.get("text", "") for b in content if b.get("type") == "text"))
+                    system_parts.extend(
+                        {  # mutable-ok: JSON wire system block
+                            "type": "text",
+                            "text": block.get("text", ""),
+                            **(
+                                {"cache_control": block["cache_control"]} if "cache_control" in block else {}
+                            ),  # mutable-ok: JSON wire block
+                        }
+                        for block in content
+                        if isinstance(block, Mapping) and block.get("type") == "text"
+                    )
             elif role == "assistant":
                 tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None)
+                thinking_blocks = _signed_thinking_blocks(msg)
                 if tool_calls:
-                    content_blocks: list[dict[str, object]] = []
+                    content_blocks: list[dict[str, object]] = list(thinking_blocks)  # mutable-ok: JSON wire blocks
                     if content:
                         content_blocks.append({"type": "text", "text": content})
                     for tc in tool_calls:
@@ -239,18 +363,26 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
                             }
                         )
                     conversation.append({"role": "assistant", "content": content_blocks})
+                elif thinking_blocks:
+                    thinking_content = (
+                        [
+                            *thinking_blocks,
+                            *copy.deepcopy(content),
+                        ]
+                        if isinstance(content, list)
+                        else [*thinking_blocks, *([{"type": "text", "text": content}] if content else [])]
+                    )  # rebind-ok: loop-local normalized content
+                    conversation.append({"role": "assistant", "content": thinking_content})
                 else:
                     conversation.append({"role": "assistant", "content": content})
             elif role == "tool":
-                tool_call_id = (
+                tool_call_id_value = (
                     msg.get("tool_call_id", "") if isinstance(msg, dict) else getattr(msg, "tool_call_id", "")
                 )
-                tool_content = content if isinstance(content, str) else json.dumps(content)
-                tool_result_block = {
-                    "type": "tool_result",
-                    "tool_use_id": tool_call_id,
-                    "content": tool_content,
-                }
+                tool_call_id = (
+                    tool_call_id_value if isinstance(tool_call_id_value, str) else ""
+                )  # rebind-ok: normalized loop value
+                tool_result_block = _convert_tool_result_to_anthropic(content, tool_call_id, msg_cache_control)
                 if (
                     conversation
                     and conversation[-1]["role"] == "user"
@@ -260,11 +392,18 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
                 ):
                     conversation[-1]["content"].append(tool_result_block)
                 else:
-                    conversation.append({"role": "user", "content": [tool_result_block]})
+                    conversation.append(
+                        {"role": "user", "content": [tool_result_block]}  # mutable-ok: JSON wire message
+                    )  # mutable-ok: JSON wire message
             else:
-                conversation.append({"role": role, "content": content})
+                conversation.append(  # mutable-ok: JSON wire message
+                    {  # mutable-ok: JSON wire message
+                        "role": role,
+                        "content": _convert_image_url_blocks_to_anthropic(content),
+                    }  # mutable-ok: JSON wire message
+                )
 
-        system: Final[str | None] = "\n\n".join(system_parts) if system_parts else None
+        system: Final[list[dict] | None] = system_parts if system_parts else None  # mutable-ok: JSON wire messages
         return system, conversation
 
     def transform_request(
@@ -339,7 +478,9 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
         extra_body: dict,
     ) -> dict:
         """Anthropic Messages format for /messages endpoint."""
-        system, conversation = self._extract_system_and_messages(messages)
+        passthrough_system: Final = optional_params.pop("system", None)
+        extracted_system, conversation = self._extract_system_and_messages(messages)
+        system: Final = passthrough_system if passthrough_system is not None else extracted_system
 
         if "tools" in optional_params:
             optional_params["tools"] = self._transform_tools_to_anthropic(optional_params["tools"])
@@ -353,16 +494,19 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
 
         model_name: Final = model.removeprefix("snowflake/")
 
-        body: Final[dict[str, object]] = {
-            "model": model_name,
-            "messages": conversation,
-            "stream": stream,
-            **optional_params,
-            **extra_body,
-        }
-
+        body: Final[dict[str, object]] = normalize_cache_control_in_anthropic_payload(  # mutable-ok: JSON wire body
+            {  # mutable-ok: JSON wire body
+                "model": model_name,
+                "messages": conversation,
+                "stream": stream,
+                **optional_params,
+                **extra_body,  # mutable-ok: JSON wire body
+            }
+        )
         if system is not None:
-            body["system"] = system
+            body["system"] = normalize_cache_control_in_anthropic_payload(  # mutable-ok: JSON wire payload
+                {"system": system}  # mutable-ok: JSON wire payload
+            )["system"]
 
         if "max_tokens" not in body:
             body["max_tokens"] = 4096  # reasonable default; Anthropic API max varies by model
@@ -435,23 +579,10 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
             additional_args={"complete_input_dict": request_data},
         )
 
-        text_content = ""
-        tool_calls: Final = []
-
-        for block in response_json.get("content", []):
-            if block.get("type") == "text":
-                text_content += block.get("text", "")
-            elif block.get("type") == "tool_use":
-                tool_calls.append(
-                    ChatCompletionMessageToolCall(
-                        id=block.get("id", ""),
-                        type="function",
-                        function=Function(
-                            name=block.get("name", ""),
-                            arguments=json.dumps(block.get("input", {})),
-                        ),
-                    )
-                )
+        anthropic_config: Final = AnthropicConfig()
+        text_content, _, thinking_blocks, reasoning_content, tool_calls, _, _, _ = (
+            anthropic_config.extract_response_content(completion_response=dict(response_json))
+        )
 
         _stop_reason_map: Final = {
             "end_turn": "stop",
@@ -461,9 +592,13 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
         }
         finish_reason: Final = _stop_reason_map.get(response_json.get("stop_reason", "end_turn"), "stop")
 
-        message: Final = Message(content=text_content or None, role="assistant")
-        if tool_calls:
-            message.tool_calls = tool_calls
+        message: Final = Message(
+            content=text_content or None,
+            role="assistant",
+            tool_calls=tool_calls or None,
+            thinking_blocks=thinking_blocks,
+            reasoning_content=reasoning_content,
+        )
 
         choice: Final = Choices(
             finish_reason=finish_reason,
@@ -471,11 +606,13 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
             message=message,
         )
 
-        usage_data: Final = response_json.get("usage", {})
-        usage: Final = Usage(
-            prompt_tokens=usage_data.get("input_tokens", 0),
-            completion_tokens=usage_data.get("output_tokens", 0),
-            total_tokens=usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0),
+        # Cortex reports prompt-cache creation/read counts alongside input_tokens; the
+        # shared calculator folds them into prompt_tokens_details so cached input is
+        # visible and billed at its own rate.
+        usage: Final = anthropic_config.calculate_usage(
+            usage_object=response_json.get("usage", {}),
+            reasoning_content=reasoning_content,
+            completion_response=dict(response_json),
         )
 
         model_response.choices = [choice]
@@ -516,15 +653,19 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator):
         json_mode: bool | None = False,
     ):
         super().__init__(streaming_response=streaming_response, sync_stream=sync_stream)
-        self._tool_index = 0
-        self._tool_id = ""
-        self._tool_name = ""
-        self._input_tokens = 0
+        # Cortex streams the Anthropic SSE dialect on /messages, so its events are parsed
+        # by Anthropic's own parser: thinking deltas, signatures and prompt-cache usage
+        # all arrive the way they do on every other Anthropic-dialect provider.
+        self._anthropic_parser: Final = AnthropicStreamParser(
+            streaming_response=streaming_response,
+            sync_stream=sync_stream,
+            json_mode=json_mode,
+        )
 
-    def chunk_parser(self, chunk: dict) -> GenericStreamingChunk:
+    def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream:
         if "choices" in chunk:
             return self._parse_openai_chunk(chunk)
-        return self._parse_anthropic_chunk(chunk)
+        return self._anthropic_parser.chunk_parser(chunk)
 
     def _parse_openai_chunk(self, chunk: dict) -> GenericStreamingChunk:
         choices: Final = chunk.get("choices", [])
@@ -566,117 +707,3 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator):
             index=choice.get("index", 0),
             tool_use=tool_use,
         )
-
-    def _parse_anthropic_chunk(self, chunk: dict) -> GenericStreamingChunk:
-        event_type: Final = chunk.get("type", "")
-
-        if event_type == "message_start":
-            message: Final = chunk.get("message", {})
-            usage_data = message.get("usage", {})
-            self._input_tokens = usage_data.get("input_tokens", 0)
-            return GenericStreamingChunk(
-                text="",
-                is_finished=False,
-                finish_reason="",
-                usage=None,
-                index=0,
-                tool_use=None,
-            )
-
-        elif event_type == "content_block_delta":
-            delta = chunk.get("delta", {})
-            delta_type: Final = delta.get("type", "")
-
-            if delta_type == "text_delta":
-                return GenericStreamingChunk(
-                    text=delta.get("text", ""),
-                    is_finished=False,
-                    finish_reason="",
-                    usage=None,
-                    index=chunk.get("index", 0),
-                    tool_use=None,
-                )
-            elif delta_type == "input_json_delta":
-                return GenericStreamingChunk(
-                    text="",
-                    is_finished=False,
-                    finish_reason="",
-                    usage=None,
-                    index=chunk.get("index", 0),
-                    tool_use=ChatCompletionToolCallChunk(
-                        id=self._tool_id,
-                        type="function",
-                        function={
-                            "name": self._tool_name,
-                            "arguments": delta.get("partial_json", ""),
-                        },
-                        index=self._tool_index,
-                    ),
-                )
-
-        elif event_type == "content_block_start":
-            content_block: Final = chunk.get("content_block", {})
-            if content_block.get("type") == "tool_use":
-                self._tool_id = content_block.get("id", "")
-                self._tool_name = content_block.get("name", "")
-                self._tool_index = chunk.get("index", 0)
-                return GenericStreamingChunk(
-                    text="",
-                    is_finished=False,
-                    finish_reason="",
-                    usage=None,
-                    index=chunk.get("index", 0),
-                    tool_use=ChatCompletionToolCallChunk(
-                        id=self._tool_id,
-                        type="function",
-                        function={"name": self._tool_name, "arguments": ""},
-                        index=self._tool_index,
-                    ),
-                )
-
-        elif event_type == "message_delta":
-            delta = chunk.get("delta", {})
-            stop_reason: Final = delta.get("stop_reason", "")
-            usage_data = chunk.get("usage", {})
-            _stop_map: Final = {
-                "end_turn": "stop",
-                "max_tokens": "length",
-                "tool_use": "tool_calls",
-                "stop_sequence": "stop",
-            }
-            usage = None
-            if usage_data or self._input_tokens:
-                output_t: Final = usage_data.get("output_tokens", 0)
-                input_t: Final = self._input_tokens or usage_data.get("input_tokens", 0)
-                usage = ChatCompletionUsageBlock(
-                    prompt_tokens=input_t,
-                    completion_tokens=output_t,
-                    total_tokens=input_t + output_t,
-                )
-            return GenericStreamingChunk(
-                text="",
-                is_finished=True,
-                finish_reason=_stop_map.get(stop_reason, "stop"),
-                usage=usage,
-                index=0,
-                tool_use=None,
-            )
-
-        elif event_type == "message_stop":
-            return GenericStreamingChunk(
-                text="",
-                is_finished=True,
-                finish_reason="stop",
-                usage=None,
-                index=0,
-                tool_use=None,
-            )
-
-        return GenericStreamingChunk(
-            text="",
-            is_finished=False,
-            finish_reason="",
-            usage=None,
-            index=0,
-            tool_use=None,
-        )

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -3638,30 +3638,48 @@ class MCPServerManager:
         user_api_key_auth: UserAPIKeyAuth | None,
         raw_headers: Mapping[str, str] | None = None,
     ) -> None:
-        """Run the OBO exchange for a caller-supplied subject at the transport edge.
+        """Mint an exchange-backed server's upstream credential at the transport edge.
 
         Single-server routes call this before the MCP session opens, where an HTTP status and
         ``WWW-Authenticate`` still reach the client. A rejected subject raises the RFC 9728
         challenge and any other ``CredError`` maps onto its public HTTP status, so an exchange
         failure surfaces as a failure instead of the session continuing into an empty tool list.
         A successful exchange is cached by the exchanger, so the session's list/call reuses it.
+
+        Each mode pre-flights only where it would resolve the subject the session goes on to use,
+        which is what keeps the pre-flight from reaching a verdict the session would contradict.
+        ``oauth2_token_exchange`` mints from the caller's inbound bearer, so without one there is
+        nothing to exchange and the missing-subject case stays the preemptive challenge's job.
+        ``oauth2_id_jag`` is the mirror image: tool listing resolves it from the identity assertion
+        captured for this user at SSO login and never from the inbound bearer, so the pre-flight is
+        faithful exactly when no identity bearer was sent (a LiteLLM key in ``Authorization`` is not one),
+        and a caller that did send one is passed through
+        untouched rather than judged against a subject the listing will not use. That store-sourced
+        case is the one whose missing-assertion 412 and store-outage 503 the session cannot report.
+        Only OBO has a discovery challenge to raise; ID-JAG's failures are plain statuses whose body
+        already names what the user has to do, so they map through ``raise_public`` as at egress.
         """
-        if server.auth_type != MCPAuth.oauth2_token_exchange:
-            return
-        if not self._extract_bearer_token(oauth2_headers, None):
-            return
-        resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)
-        spec: Final = to_server_spec(resolved_server)
-        if spec is None or not isinstance(spec.config, TokenExchangeConfig):
-            return
         subject_token: Final = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
-        if subject_token is None:
+        match server.auth_type:
+            case MCPAuth.oauth2_token_exchange:
+                if not self._extract_bearer_token(oauth2_headers, None):
+                    return
+            case MCPAuth.oauth2_id_jag:
+                if subject_token is not None:
+                    return
+            case _:
+                return
+        resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)
+        spec: Final = _to_server_spec_fail_closed(resolved_server)
+        if spec is None or not isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)):
+            return
+        if subject_token is None and isinstance(spec.config, TokenExchangeConfig):
             raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path())
         match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
             case Ok(_):
                 return
             case Error(err):
-                if err.tag == "unauthorized":
+                if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig):
                     raise_token_exchange_challenge(
                         resolved_server,
                         root_path=get_server_root_path(),

litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py

@@ -11,7 +11,8 @@ being registered, so a gateway with no EMA upstream never stores bearer material
 The row is one encrypted payload per user, latest login wins. ``expires_at`` mirrors the
 id_token ``exp`` claim and is judged by the reader, never enforced by deletion here: an
 expired assertion with a refresh token is still renewable, and the DB row is the source of
-truth, the same contract as the per-user OAuth credential store.
+truth, the same contract as the per-user OAuth credential store. Reads use a per-process cache with
+TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against stale in-flight reads.
 """
 
 from __future__ import annotations
@@ -24,6 +25,8 @@ import jwt
 from pydantic import BaseModel, ConfigDict, SecretStr, TypeAdapter, ValidationError
 
 from litellm._logging import verbose_proxy_logger
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_SSO_ASSERTION_CACHE_TTL_SECONDS
 
 if TYPE_CHECKING:
     from litellm.proxy.utils import PrismaClient
@@ -45,6 +48,46 @@ class SSOIdentityAssertion(BaseModel):
     expires_at: datetime | None = None
 
 
+class SSOAssertionCache:
+    """Process-local read cache. ``invalidate`` bumps a process-wide epoch so a fetch that started
+    before a login cannot repopulate the old assertion after it."""
+
+    def __init__(self, ttl_seconds: int = MCP_SSO_ASSERTION_CACHE_TTL_SECONDS) -> None:
+        self._entries = InMemoryCache(
+            max_size_in_memory=MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE,
+            default_ttl=ttl_seconds,
+        )
+        self._epoch: int = 0
+
+    def epoch(self) -> int:
+        return self._epoch
+
+    def get(self, user_id: str) -> SSOIdentityAssertion | None:
+        cached: Final = self._entries.get_cache(  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
+            user_id
+        )
+        return cached if isinstance(cached, SSOIdentityAssertion) else None
+
+    def set_if_unchanged(self, user_id: str, assertion: SSOIdentityAssertion, seen_epoch: int) -> None:
+        if self._epoch != seen_epoch:
+            return
+        self._entries.set_cache(  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
+            user_id, assertion
+        )
+
+    def invalidate(self, user_id: str) -> None:
+        self._epoch += 1
+        self._entries.delete_cache(  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
+            user_id
+        )
+
+    def flush(self) -> None:
+        self._entries.flush_cache()  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
+
+
+_ASSERTION_CACHE: Final = SSOAssertionCache()
+
+
 class _IdTokenClaims(BaseModel):
     exp: float | None = None
     iss: str | None = None
@@ -107,7 +150,9 @@ async def ema_assertion_retention_enabled() -> bool:
     return row is not None
 
 
-async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAssertion) -> None:
+async def persist_sso_identity_assertion(
+    user_id: str, assertion: SSOIdentityAssertion, cache: SSOAssertionCache = _ASSERTION_CACHE
+) -> None:
     from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper  # noqa: PLC0415  # runtime global
     from litellm.proxy.proxy_server import prisma_client  # noqa: PLC0415  # runtime global
 
@@ -127,11 +172,10 @@ async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAss
             "update": {"assertion_b64": encoded},
         },
     )
+    cache.invalidate(user_id)
 
 
-async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | None:
-    """The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key
-    rotation), or unparseable. Expiry is not judged here; the reader owns that policy."""
+async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None:
     from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper  # noqa: PLC0415  # runtime global
     from litellm.proxy.proxy_server import prisma_client  # noqa: PLC0415  # runtime global
 
@@ -160,6 +204,21 @@ async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | N
     )
 
 
+async def fetch_sso_identity_assertion(
+    user_id: str, cache: SSOAssertionCache = _ASSERTION_CACHE
+) -> SSOIdentityAssertion | None:
+    """The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key
+    rotation), or unparseable. Expiry is not judged here; the reader owns that policy."""
+    cached: Final = cache.get(user_id)
+    if cached is not None:
+        return cached
+    seen_epoch: Final = cache.epoch()
+    assertion: Final = await _read_assertion_from_db(user_id)
+    if assertion is not None:
+        cache.set_if_unchanged(user_id, assertion, seen_epoch)
+    return assertion
+
+
 class AssertionStoreUnavailable(Exception):
     """Raised by ``fetch`` when the backing store is unreachable (e.g. the DB is down).
 
@@ -189,9 +248,12 @@ class DbSSOAssertionStore:
     from credential resolution and from the upstream-401 retry.
     """
 
+    def __init__(self, cache: SSOAssertionCache = _ASSERTION_CACHE) -> None:
+        self._cache = cache
+
     async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
         try:
-            return await fetch_sso_identity_assertion(user_id)
+            return await fetch_sso_identity_assertion(user_id, cache=self._cache)
         except Exception as exc:  # noqa: BLE001  # any driver/storage failure is an outage, not an absence
             raise AssertionStoreUnavailable(str(exc)) from exc
 

litellm/proxy/_experimental/mcp_server/server.py

@@ -3851,15 +3851,15 @@ if MCP_AVAILABLE:
 
                 raise_token_exchange_challenge(server, root_path=get_server_root_path())
 
-            # token_exchange (OBO) with a subject present: run the exchange here at the transport
-            # edge, so a rejected subject raises the RFC 9728 challenge (and a gateway fault its
-            # public status) instead of the session opening and list_tools masking the failure as
-            # an empty tool list. Gated to single-server routes; the multi-server aggregate keeps
-            # absorbing per-server auth failures so one bad server cannot 401 the whole connect.
+            # Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run
+            # the exchange here at the transport edge, so a rejected subject raises the RFC 9728
+            # challenge and any other failure its public status, instead of the session opening and
+            # list_tools masking it as an empty tool list. The manager owns which modes pre-flight
+            # and what each mints from. Gated to single-server routes the key may reach; the
+            # multi-server aggregate keeps absorbing per-server auth failures so one bad server
+            # cannot 401 the whole connect.
             if (
                 server
-                and server.auth_type == MCPAuth.oauth2_token_exchange
-                and oauth2_headers
                 and len(mcp_servers or []) == 1
                 and server.server_id
                 in frozenset(

litellm/proxy/_types.py

@@ -737,6 +737,9 @@ class LiteLLMRoutes(enum.Enum):
             "/.well-known/litellm-ui-config",
             "/public/model_hub",
             "/public/v1/model_hub",
+            "/public/v1/model_hub/providers",
+            "/public/v1/model_hub/modes",
+            "/public/v1/model_hub/features",
             "/public/model_hub/info",
             "/public/agent_hub",
             "/public/mcp_hub",
@@ -2404,6 +2407,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
     """
 
     completion_model: str | None = Field(None, description="proxy level default model for all chat completion calls")
+    max_in_flight_requests_per_worker: int | None = Field(
+        None, gt=0, description="maximum concurrent requests handled by each worker"
+    )
+    max_queued_requests_per_worker: int | None = Field(
+        None, ge=0, description="maximum requests waiting for a worker slot"
+    )
+    admission_queue_timeout_seconds: float = Field(
+        1.0, gt=0, description="maximum time a request waits for a worker slot"
+    )
     plugins: list[PluginConfig] | None = Field(
         None, description="external services registered as embeddable UI plugins"
     )

litellm/proxy/anthropic_endpoints/endpoints.py

@@ -9,7 +9,7 @@ from fastapi.responses import JSONResponse
 
 import litellm
 from litellm._logging import verbose_proxy_logger
-from litellm.anthropic_interface.exceptions import AnthropicExceptionMapping
+from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping
 from litellm.integrations.custom_guardrail import ModifyResponseException
 from litellm.llms.anthropic.experimental_pass_through.context_management import (
     AnthropicContextManagementError,
@@ -30,6 +30,27 @@ from litellm.types.utils import TokenCountResponse
 router: Final = APIRouter()
 
 
+def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSONResponse:
+    from litellm.proxy.proxy_server import (
+        _close_dangling_otel_server_span,  # pyright: ignore[reportPrivateUsage]  # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does
+    )
+
+    status_code: Final = int(exc.code) if exc.code is not None and exc.code.isdigit() else 500
+    _close_dangling_otel_server_span(request, status_code, exc=exc)
+    envelope: Final = AnthropicExceptionMapping.transform_to_anthropic_error(
+        status_code=status_code,
+        raw_message=exc.message,
+        request_id=request.headers.get("x-request-id"),
+    )
+    if not exc.provider_specific_fields:
+        return JSONResponse(status_code=status_code, content=envelope, headers=exc.headers)
+    content: Final[AnthropicErrorResponse] = {
+        **envelope,
+        "error": {**envelope["error"], "provider_specific_fields": exc.provider_specific_fields},
+    }
+    return JSONResponse(status_code=status_code, content=content, headers=exc.headers)
+
+
 def _strip_total_tokens_from_anthropic_response(response: Any) -> None:
     """Remove the OpenAI-flavored `usage.total_tokens` field that LiteLLM
     injects into Anthropic /v1/messages responses.
@@ -195,7 +216,7 @@ async def anthropic_response(
         verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e)
 
         if isinstance(e, ProxyException):
-            raise
+            return _anthropic_error_json_response(e, request)
 
         # Extract model_id from request metadata (same as success path)
         litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
@@ -216,15 +237,18 @@ async def anthropic_response(
         )
 
         if isinstance(e, HTTPException):
-            raise proxy_exception_from_http_exception(e, headers)
+            return _anthropic_error_json_response(proxy_exception_from_http_exception(e, headers), request)
 
         error_msg: Final = f"{e}"
-        raise ProxyException(
-            message=getattr(e, "message", error_msg),
-            type=getattr(e, "type", "None"),
-            param=getattr(e, "param", "None"),
-            code=getattr(e, "status_code", 500),
-            headers=headers,
+        return _anthropic_error_json_response(
+            ProxyException(
+                message=getattr(e, "message", error_msg),
+                type=getattr(e, "type", "None"),
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", 500),
+                headers=headers,
+            ),
+            request,
         )
 
 

litellm/proxy/auth/model_checks.py

@@ -387,33 +387,22 @@ def _get_wildcard_models(
     all_wildcard_models: Final = []
     for model in unique_models:
         if _check_wildcard_routing(model=model):
-            if return_wildcard_routes:  # will add the wildcard route to the list eg: anthropic/*.
+            if return_wildcard_routes:
                 all_wildcard_models.append(model)
 
-            ## get litellm params from model
-            if llm_router is not None:
-                model_list = llm_router.get_model_list(model_name=model, team_id=team_id)
-                if model_list:
-                    for router_model in model_list:
-                        wildcard_models = get_known_models_from_wildcard(
+            models_to_remove.add(model)
+
+            model_list = llm_router.get_model_list(model_name=model, team_id=team_id) if llm_router else None
+            if model_list:
+                for router_model in model_list:
+                    all_wildcard_models.extend(
+                        get_known_models_from_wildcard(
                             wildcard_model=model,
                             litellm_params=LiteLLM_Params(**router_model["litellm_params"]),
                         )
-                        all_wildcard_models.extend(wildcard_models)
-                else:
-                    # Router has no deployment for this wildcard (e.g., BYOK team models)
-                    # Fall back to expanding from known provider models
-                    wildcard_models = get_known_models_from_wildcard(wildcard_model=model, litellm_params=None)
-                    if wildcard_models:
-                        models_to_remove.add(model)
-                        all_wildcard_models.extend(wildcard_models)
+                    )
             else:
-                # get all known provider models
-                wildcard_models = get_known_models_from_wildcard(wildcard_model=model, litellm_params=None)
-
-                if wildcard_models:
-                    models_to_remove.add(model)
-                    all_wildcard_models.extend(wildcard_models)
+                all_wildcard_models.extend(get_known_models_from_wildcard(wildcard_model=model, litellm_params=None))
 
     for model in models_to_remove:
         unique_models.remove(model)

litellm/proxy/health_endpoints/_health_endpoints.py

@@ -50,6 +50,9 @@ from litellm.proxy.health_check import (
     perform_health_check,
     run_with_timeout,
 )
+from litellm.proxy.middleware.admission_control_middleware import (
+    get_admission_control_stats,
+)
 from litellm.proxy.middleware.in_flight_requests_middleware import (
     get_in_flight_requests,
 )
@@ -63,6 +66,13 @@ from litellm.secret_managers.main import get_secret_bool
 #### Health ENDPOINTS ####
 
 
+class _HealthBacklogResponse(TypedDict):
+    in_flight_requests: ReadOnly[int]
+    admitted_requests: ReadOnly[int]
+    queued_requests: ReadOnly[int]
+    rejected_requests: ReadOnly[int]
+
+
 def _reject_os_environ_references(params: dict) -> None:
     """
     Validate that the provided params do not contain any ``os.environ/``
@@ -1759,7 +1769,14 @@ async def health_backlog():
     for the event loop to get to them, adding latency before LiteLLM even starts
     its own timer.
     """
-    return {"in_flight_requests": get_in_flight_requests()}
+    stats: Final = get_admission_control_stats()
+    response: Final[_HealthBacklogResponse] = {
+        "in_flight_requests": get_in_flight_requests(),
+        "admitted_requests": stats.admitted,
+        "queued_requests": stats.queued,
+        "rejected_requests": stats.rejected_total,
+    }
+    return response
 
 
 @router.get(

litellm/proxy/list_api/in_memory.py

@@ -141,3 +141,15 @@ class InMemoryListExecutor(Generic[TRow]):
     async def find_many(self, plan: QueryPlan) -> Sequence[TRow]:
         page: Final = _ordered(self._matching(plan.where), plan.order)[plan.skip : plan.skip + plan.take]
         return await self.enrich_page(tuple(row for _, row in page))
+
+    async def distinct(self, field: str, where: tuple[Predicate, ...]) -> Sequence[str]:
+        """A repeated field contributes each of its elements, so a facet over `providers`
+        lists providers rather than the tuples rows happen to carry."""
+        cells: Final = (cells.get(field) for cells, _ in self._matching(where))
+        values: Final = (
+            value
+            for cell in cells
+            for value in (cell if isinstance(cell, tuple) else (cell,))
+            if isinstance(value, str) and value
+        )
+        return tuple(sorted(frozenset(values)))

litellm/proxy/list_api/list_framework.py

@@ -28,12 +28,15 @@ from litellm.proxy.list_api.common import (
     PROBLEM_TYPE_BASE,
     ManagementProblem,
     build_list_links,
+    build_page_links,
     escape_like,
     unknown_query_param_problem,
 )
 from litellm.types.proxy.management_endpoints.management_v1 import (
+    FacetListResponse,
     ListMeta,
     ListResponse,
+    PageMeta,
     ProblemDetail,
 )
 
@@ -186,6 +189,13 @@ class ListExecutor(Protocol[TRow_co]):
     async def find_many(self, plan: QueryPlan) -> Sequence[TRow_co]: ...
 
 
+class FacetExecutor(Protocol):
+    """The half of a facet that knows the rows. Separate from `ListExecutor` so a SQL
+    executor is not forced to implement `distinct` to keep serving entity lists."""
+
+    async def distinct(self, field: str, where: tuple[Predicate, ...]) -> Sequence[str]: ...
+
+
 def order_by_sql(order: tuple[SortKey, ...]) -> str:
     """`ORDER BY` body for a plan, NULLS LAST in both directions.
 
@@ -515,6 +525,78 @@ def build_query_plan(
     )
 
 
+def _facet_allowed_params(spec: ListSpec[TRow, TOut]) -> tuple[str, ...]:
+    """A facet's values are always ascending, so `sort` is not one of its parameters."""
+    return tuple(name for name in _allowed_params(spec) if name != SORT_PARAM)
+
+
+def _facet_where(
+    spec: ListSpec[TRow, TOut],
+    params: Mapping[str, str],
+    caller: UserAPIKeyAuth,
+) -> tuple[Predicate, ...] | ProblemDetail:
+    scope_predicates: Final = _scope_predicates(spec.scope(caller))
+    if isinstance(scope_predicates, ProblemDetail):
+        return scope_predicates
+    filters: Final = _parse_filters(spec, params)
+    if isinstance(filters, ProblemDetail):
+        return filters
+    search: Final = _search_predicate(spec, params)
+    return scope_predicates + filters + ((search,) if search is not None else ())
+
+
+async def handle_facet(
+    spec: ListSpec[TRow, TOut],
+    executor: FacetExecutor,
+    request: Request,
+    caller: UserAPIKeyAuth,
+    field: str,
+) -> FacetListResponse:
+    """The distinct values one column takes over a filtered query on a resource.
+
+    Carries the parent's parameters so a filter dropdown offers exactly the values the
+    table can show, and `has_more` rather than a total, which would cost a COUNT(*) over
+    the whole match set on every keystroke.
+    """
+    params: Final = request.query_params
+    unknown: Final = tuple(sorted(name for name in params if name == SORT_PARAM or not _is_known_param(spec, name)))
+    if unknown:
+        raise ManagementProblem(unknown_query_param_problem(unknown=unknown, allowed=_facet_allowed_params(spec)))
+
+    duplicates: Final = _duplicate_params(request)
+    if duplicates:
+        raise ManagementProblem(
+            _problem(
+                "duplicate-query-parameter",
+                "Duplicate query parameter",
+                400,
+                f"Repeated query parameter(s): {', '.join(duplicates)}. Each may appear once; "
+                f"use a comma-separated list for multiple filter values.",
+            )
+        )
+
+    page: Final = _parse_page(params)
+    if isinstance(page, ProblemDetail):
+        raise ManagementProblem(page)
+    page_size: Final = _parse_page_size(spec, params)
+    if isinstance(page_size, ProblemDetail):
+        raise ManagementProblem(page_size)
+
+    where: Final = _facet_where(spec, params, caller)
+    if isinstance(where, ProblemDetail):
+        raise ManagementProblem(where)
+
+    values: Final = await executor.distinct(field, where)
+    skip: Final = (page - 1) * page_size
+    window: Final = values[skip : skip + page_size + 1]
+    has_more: Final = len(window) > page_size
+    return FacetListResponse(
+        data=tuple(window[:page_size]),
+        meta=PageMeta(page=page, page_size=page_size, has_more=has_more),
+        links=build_page_links(request=request, page=page, has_more=has_more),
+    )
+
+
 def _duplicate_params(request: Request) -> tuple[str, ...]:
     names: Final = tuple(name for name, _ in request.query_params.multi_items())
     return tuple(sorted(frozenset(name for name in names if names.count(name) > 1)))

litellm/proxy/management_endpoints/policy_endpoints/endpoints.py

@@ -12,7 +12,7 @@ All /policy management endpoints
 import copy
 import json
 import os
-from collections.abc import AsyncIterator
+from collections.abc import AsyncGenerator, AsyncIterator
 from typing import TYPE_CHECKING, Final, Literal, cast
 
 from fastapi import APIRouter, Depends, HTTPException, Request
@@ -20,6 +20,7 @@ from fastapi.responses import Response, StreamingResponse
 from pydantic import BaseModel, Field
 from typing_extensions import TypedDict
 
+import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import (
     COMPETITOR_LLM_TEMPERATURE,
@@ -32,6 +33,10 @@ from litellm.llms.openai.chat.guardrail_translation.handler import (
 )
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.common_utils.sse_keepalive import (
+    SSE_COMMENT_PING,
+    wrap_sse_stream_with_keepalive_pings,
+)
 from litellm.proxy.guardrails.guardrail_hooks.custom_code import (
     RESPONSE_REJECTION_GUARDRAIL_CODE,
     CustomCodeGuardrail,
@@ -811,7 +816,7 @@ async def _stream_competitor_events(
     llm_enrichment: dict,
     brand_name: str,
     model: str,
-) -> AsyncIterator[str]:
+) -> AsyncGenerator[str, None]:
     """Stream competitor names as SSE events, then emit a final 'done' event."""
     competitors: Final[list[str]] = list(data.competitors or [])
 
@@ -883,7 +888,11 @@ async def enrich_policy_template_stream(
     model: Final = data.model or DEFAULT_COMPETITOR_DISCOVERY_MODEL
 
     return StreamingResponse(
-        _stream_competitor_events(data, template, llm_enrichment, brand_name, model),
+        wrap_sse_stream_with_keepalive_pings(
+            _stream_competitor_events(data, template, llm_enrichment, brand_name, model),
+            ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
+            ping_chunk=SSE_COMMENT_PING,
+        ),
         media_type="text/event-stream",
         headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
     )

litellm/proxy/management_endpoints/scim/scim_v2.py

@@ -1123,13 +1123,15 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou
             user_id=user_id,
             user_email=user_id,  # We don't have email from group membership
             user_alias=None,
-            teams=[],  # Teams will be added separately
             metadata={"created_via": created_via},
             auto_create_key=False,
             user_role=default_role,
         )
 
-        created_user: Final = await new_user(data=new_user_request)
+        created_user: Final = await new_user(
+            data=new_user_request,
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
         verbose_proxy_logger.info("Created user %s via %s", user_id, created_via)
         return created_user
 
@@ -1699,7 +1701,7 @@ async def create_user(
             user_id=user_id,
             user_email=user_data["user_email"],
             user_alias=user_data["user_alias"],
-            teams=user_data["teams"],
+            teams=user_data["teams"] or None,
             metadata=metadata,
             auto_create_key=False,
             user_role=resolved_role if admin_group is not None else default_role,
@@ -1717,6 +1719,7 @@ async def create_user(
 
         created_user: Final = await new_user(
             data=new_user_request,
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
         )
 
         scim_user: Final = await ScimTransformations.transform_litellm_user_to_scim_user(user=created_user)
@@ -1771,22 +1774,25 @@ async def update_user(
             roles=user_data["roles"],
         )
 
+        # SCIM User.groups is readOnly (RFC 7643 4.1.2): IdPs sync membership via /Groups and send
+        # no groups or `groups: []` on profile PUTs, so empty means unspecified, not "remove from every team"
+        target_teams: Final = user_data["teams"] or existing_user.teams
         await _handle_team_membership_changes(
             user_id=user_id,
-            existing_teams=existing_user.teams or [],
-            new_teams=user_data["teams"],
+            existing_teams=existing_user.teams,
+            new_teams=target_teams,
         )
 
         update_data: Final = {
             "user_email": user_data["user_email"],
             "user_alias": user_data["user_alias"],
             "sso_user_id": user_data["sso_user_id"],
-            "teams": user_data["teams"],
+            "teams": target_teams,
             "metadata": safe_dumps(metadata),
         }
 
         admin_group: Final = await _get_scim_admin_group()
-        if admin_group is not None:
+        if admin_group is not None and user_data["teams"]:
             update_data["user_role"] = _resolve_scim_user_role(
                 user.groups or [], admin_group, _default_scim_user_role()
             )

litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py

@@ -4,7 +4,7 @@ usage/spend data by querying the aggregated daily activity endpoints.
 """
 
 import json
-from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence
 from datetime import date
 from typing import Any, Final, Literal, Protocol, cast, overload
 
@@ -543,7 +543,7 @@ async def stream_usage_ai_chat(
     model: str | None = None,
     user_id: str | None = None,
     is_admin: bool = False,
-) -> AsyncIterator[str]:
+) -> AsyncGenerator[str, None]:
     """Stream SSE events: status → tool_call → chunk → done."""
     resolved_model: Final = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL
     truncated: Final = messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages

litellm/proxy/management_endpoints/usage_endpoints/endpoints.py

@@ -10,8 +10,13 @@ from fastapi import APIRouter, Depends, Request
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
 
+import litellm
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.common_utils.sse_keepalive import (
+    SSE_COMMENT_PING,
+    wrap_sse_stream_with_keepalive_pings,
+)
 
 router: Final = APIRouter()
 
@@ -56,11 +61,15 @@ async def usage_ai_chat(
     messages: Final = [{"role": m.role, "content": m.content} for m in data.messages]
 
     return StreamingResponse(
-        stream_usage_ai_chat(
-            messages=messages,
-            model=data.model,
-            user_id=user_id,
-            is_admin=is_admin,
+        wrap_sse_stream_with_keepalive_pings(
+            stream_usage_ai_chat(
+                messages=messages,
+                model=data.model,
+                user_id=user_id,
+                is_admin=is_admin,
+            ),
+            ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
+            ping_chunk=SSE_COMMENT_PING,
         ),
         media_type="text/event-stream",
         headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},

litellm/proxy/middleware/admission_control_middleware.py

@@ -0,0 +1,315 @@
+import asyncio
+import os
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Annotated, Final, Protocol, TypeAlias, runtime_checkable
+
+from pydantic import Field, TypeAdapter, ValidationError
+from starlette.responses import JSONResponse
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from litellm._logging import verbose_proxy_logger
+
+_EXEMPT_PATHS: Final[frozenset[str]] = frozenset(
+    {
+        "/health/liveliness",
+        "/health/liveness",
+        "/health/readiness",
+        "/health/readiness/details",
+        "/health/backlog",
+        "/health/drain",
+        "/metrics",
+        "/metrics/",
+    }
+)
+
+
+@dataclass(frozen=True, slots=True)
+class AdmissionControlSettings:
+    max_in_flight_requests: int
+    max_queued_requests: int
+    queue_timeout_seconds: float
+
+
+AdmissionControlSettingsGetter: TypeAlias = Callable[[], AdmissionControlSettings | None]  # mutable-ok: Callable params
+
+
+@dataclass(frozen=True, slots=True)
+class AdmissionControlStats:
+    admitted: int
+    queued: int
+    rejected_total: int
+
+
+@runtime_checkable
+class _Gauge(Protocol):
+    def inc(self, amount: float = 1) -> None: ...
+
+    def dec(self, amount: float = 1) -> None: ...
+
+
+@runtime_checkable
+class _CounterChild(Protocol):
+    def inc(self, amount: float = 1) -> None: ...
+
+
+@runtime_checkable
+class _Counter(Protocol):
+    def labels(self, reason: str) -> _CounterChild: ...
+
+
+@dataclass(frozen=True, slots=True)
+class AdmissionControlMetrics:
+    admitted_gauge: _Gauge
+    queued_gauge: _Gauge
+    rejected_counter: _Counter
+
+
+AdmissionControlMetricsFactory: TypeAlias = Callable[[], AdmissionControlMetrics | None]  # mutable-ok: Callable params
+
+
+class AdmissionControlState:
+    """Per-process admission counters and the in-flight semaphore shared by one worker's requests."""
+
+    def __init__(self, metrics_factory: AdmissionControlMetricsFactory) -> None:
+        self._metrics_factory = metrics_factory
+        self._metrics: AdmissionControlMetrics | None = None
+        self._metrics_init_attempted = False
+        self._admitted = 0
+        self._queued = 0
+        self._rejected_total = 0
+        self._semaphore: asyncio.Semaphore | None = None
+        self._semaphore_loop: asyncio.AbstractEventLoop | None = None
+
+    def get_stats(self) -> AdmissionControlStats:
+        return AdmissionControlStats(
+            admitted=self._admitted,
+            queued=self._queued,
+            rejected_total=self._rejected_total,
+        )
+
+    def get_semaphore(self, max_in_flight_requests: int) -> asyncio.Semaphore:
+        loop: Final = asyncio.get_running_loop()
+        if self._semaphore_loop is not loop:
+            self._semaphore = asyncio.Semaphore(max_in_flight_requests)
+            self._semaphore_loop = loop
+        semaphore: Final = self._semaphore
+        if semaphore is None:
+            raise RuntimeError("Admission control semaphore was not initialized")
+        return semaphore
+
+    def record_admission(self) -> None:
+        self._admitted += 1
+        metrics: Final = self._get_metrics()
+        if metrics is not None:
+            metrics.admitted_gauge.inc()
+
+    def record_release(self) -> None:
+        self._admitted -= 1
+        metrics: Final = self._get_metrics()
+        if metrics is not None:
+            metrics.admitted_gauge.dec()
+
+    def record_queue(self) -> None:
+        self._queued += 1
+        metrics: Final = self._get_metrics()
+        if metrics is not None:
+            metrics.queued_gauge.inc()
+
+    def record_dequeue(self) -> None:
+        self._queued -= 1
+        metrics: Final = self._get_metrics()
+        if metrics is not None:
+            metrics.queued_gauge.dec()
+
+    def record_rejection(self, reason: str) -> None:
+        self._rejected_total += 1
+        metrics: Final = self._get_metrics()
+        if metrics is not None:
+            metrics.rejected_counter.labels(reason=reason).inc()
+
+    def _get_metrics(self) -> AdmissionControlMetrics | None:
+        if not self._metrics_init_attempted:
+            self._metrics_init_attempted = True
+            self._metrics = self._metrics_factory()
+        return self._metrics
+
+
+class AdmissionControlMiddleware:
+    def __init__(
+        self,
+        app: ASGIApp,
+        get_settings: AdmissionControlSettingsGetter,
+        state: AdmissionControlState,
+    ) -> None:
+        self.app = app
+        self.get_settings = get_settings
+        self.state = state
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        settings: Final = self.get_settings()
+        if settings is None or _get_route_path(scope) in _EXEMPT_PATHS:
+            await self.app(scope, receive, send)
+            return
+
+        state: Final = self.state
+        semaphore: Final = state.get_semaphore(settings.max_in_flight_requests)
+        if not semaphore.locked():
+            await semaphore.acquire()
+            state.record_admission()
+        elif state.get_stats().queued >= settings.max_queued_requests:
+            state.record_rejection("queue_full")
+            await _overloaded_response(state)(scope, receive, send)
+            return
+        else:
+            state.record_queue()
+            try:
+                await asyncio.wait_for(
+                    semaphore.acquire(),
+                    timeout=settings.queue_timeout_seconds,
+                )
+            except asyncio.TimeoutError:
+                state.record_dequeue()
+                state.record_rejection("queue_timeout")
+                await _overloaded_response(state)(scope, receive, send)
+                return
+            except asyncio.CancelledError:
+                state.record_dequeue()
+                raise
+            state.record_dequeue()
+            state.record_admission()
+
+        try:
+            await self.app(scope, receive, send)
+        finally:
+            semaphore.release()
+            state.record_release()
+
+
+def _get_route_path(scope: Scope) -> str:
+    """Strip the ASGI root_path (SERVER_ROOT_PATH) the same way Starlette does before route matching."""
+    path: Final[str] = scope["path"]
+    root_path: Final[str] = scope.get("root_path", "")
+    if not root_path or not path.startswith(root_path):
+        return path
+    if path == root_path:
+        return ""
+    if path[len(root_path)] == "/":
+        return path[len(root_path) :]
+    return path
+
+
+def _create_gauge(gauge_type: Callable[..., object], name: str, description: str) -> _Gauge:
+    metric: Final = (
+        gauge_type(name, description, multiprocess_mode="livesum")
+        if "PROMETHEUS_MULTIPROC_DIR" in os.environ
+        else gauge_type(name, description)
+    )
+    if not isinstance(metric, _Gauge):
+        raise TypeError("Admission gauge has an unexpected type")
+    return metric
+
+
+def create_prometheus_admission_metrics() -> AdmissionControlMetrics | None:
+    try:
+        from prometheus_client import Counter, Gauge
+
+        return AdmissionControlMetrics(
+            admitted_gauge=_create_gauge(
+                Gauge,
+                "litellm_admission_admitted_requests",
+                "Number of requests admitted by this worker",
+            ),
+            queued_gauge=_create_gauge(
+                Gauge,
+                "litellm_admission_queued_requests",
+                "Number of requests queued by this worker",
+            ),
+            rejected_counter=Counter(  # mutable-ok: Prometheus requires runtime Counter construction
+                "litellm_admission_rejected_requests_total",
+                "Number of requests rejected by this worker",
+                labelnames=("reason",),
+            ),
+        )
+    except (ImportError, ValueError):
+        return None
+
+
+admission_control_state: Final = AdmissionControlState(create_prometheus_admission_metrics)
+
+
+def get_admission_control_stats() -> AdmissionControlStats:
+    return admission_control_state.get_stats()
+
+
+_PositiveInt: TypeAlias = Annotated[int, Field(gt=0)]
+_NonNegativeInt: TypeAlias = Annotated[int, Field(ge=0)]
+_PositiveFloat: TypeAlias = Annotated[float, Field(gt=0)]
+_AdmissionControlRaw: TypeAlias = int | float | str | None
+
+
+def _hashable(value: object) -> _AdmissionControlRaw:
+    return value if value is None or isinstance(value, (int, float, str)) else repr(value)
+
+
+_POSITIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_PositiveInt)
+_NON_NEGATIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_NonNegativeInt)
+_POSITIVE_FLOAT_ADAPTER: Final[TypeAdapter[float]] = TypeAdapter(_PositiveFloat)
+
+
+@lru_cache(maxsize=16)
+def _parse_admission_control_settings(
+    max_in_flight_raw: _AdmissionControlRaw,
+    max_queued_raw: _AdmissionControlRaw,
+    queue_timeout_raw: _AdmissionControlRaw,
+) -> AdmissionControlSettings | None:
+    try:
+        max_in_flight: Final = _POSITIVE_INT_ADAPTER.validate_python(max_in_flight_raw)
+        max_queued: Final = (
+            max_in_flight if max_queued_raw is None else _NON_NEGATIVE_INT_ADAPTER.validate_python(max_queued_raw)
+        )
+        queue_timeout: Final = _POSITIVE_FLOAT_ADAPTER.validate_python(queue_timeout_raw)
+    except ValidationError as exc:
+        verbose_proxy_logger.error(
+            "Ignoring invalid admission control settings, per-worker admission control is disabled: %s",
+            exc,
+        )
+        return None
+    return AdmissionControlSettings(
+        max_in_flight_requests=max_in_flight,
+        max_queued_requests=max_queued,
+        queue_timeout_seconds=queue_timeout,
+    )
+
+
+def get_admission_control_settings(settings: Mapping[str, object]) -> AdmissionControlSettings | None:
+    max_in_flight_raw: Final = settings.get("max_in_flight_requests_per_worker")
+    if max_in_flight_raw is None:
+        return None
+    return _parse_admission_control_settings(
+        _hashable(max_in_flight_raw),
+        _hashable(settings.get("max_queued_requests_per_worker")),
+        _hashable(settings.get("admission_queue_timeout_seconds", 1.0)),
+    )
+
+
+def _overloaded_response(state: AdmissionControlState) -> JSONResponse:
+    stats: Final = state.get_stats()
+    return JSONResponse(
+        status_code=503,
+        headers={"retry-after": "1"},  # mutable-ok: Starlette expects a plain headers mapping
+        content={  # mutable-ok: Starlette serializes a plain response mapping
+            "error": {  # mutable-ok: nested response mapping
+                "message": (
+                    f"Worker at capacity: {stats.admitted} in-flight, {stats.queued} queued requests. Retry later."
+                ),
+                "type": "overloaded_error",
+                "code": "503",
+            }
+        },
+    )

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -9,10 +9,11 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
 from __future__ import annotations
 
 import hmac
+import inspect
 import json
 import os
 import re
-from collections.abc import Callable, Mapping
+from collections.abc import AsyncGenerator, Callable, Mapping
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Annotated, Final, cast
 
@@ -32,6 +33,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
 from litellm.llms.anthropic.common_utils import AnthropicModelInfo
 from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
 from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
+from litellm.passthrough.main import AsyncPassthroughStreamingResponse
 from litellm.proxy._types import *
 from litellm.proxy.auth.handle_jwt import JWTHandler
 from litellm.proxy.auth.route_checks import RouteChecks
@@ -40,6 +42,7 @@ from litellm.proxy.auth.user_api_key_auth import (
     user_api_key_auth,
     user_api_key_auth_websocket,
 )
+from litellm.proxy.common_request_processing import open_sse_before_first_byte
 from litellm.proxy.common_utils.http_parsing_utils import (
     _read_request_body,
     _safe_get_request_headers,
@@ -47,6 +50,9 @@ from litellm.proxy.common_utils.http_parsing_utils import (
     get_form_data,
     get_request_body,
 )
+from litellm.proxy.common_utils.sse_keepalive import (
+    wrap_passthrough_sse_bytes_with_keepalive_pings,
+)
 from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
 from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
     HttpPassThroughEndpointHelpers,
@@ -1478,6 +1484,74 @@ def is_azure_ai_search_service_level_index_create(method: str, endpoint: str) ->
     return path == "indexes" or path.endswith("/indexes")
 
 
+async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> AsyncGenerator[bytes, None]:
+    try:
+        async for chunk in upstream:
+            yield chunk
+    finally:
+        await upstream.aclose()
+
+
+async def _relay_azure_router_model(
+    llm_router: litellm.Router,
+    model: str,
+    endpoint: str,
+    request: Request,
+    request_body: Mapping[str, object],
+    is_streaming_request: bool,
+    user_api_key_dict: UserAPIKeyAuth,
+) -> Response:
+    result: Final = await llm_router.allm_passthrough_route(
+        model=model,
+        method=request.method,
+        endpoint=endpoint,
+        request_query_params=request.query_params,
+        request_headers=_safe_get_request_headers(request),
+        stream=is_streaming_request,
+        content=None,
+        data=None,
+        files=None,
+        json=(request_body if request.headers.get("content-type") == "application/json" else None),
+        params=None,
+        headers=None,
+        cookies=None,
+        litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict),
+    )
+
+    if not is_streaming_request:
+        upstream: Final = cast(httpx.Response, result)
+        return Response(
+            content=await upstream.aread(),
+            status_code=upstream.status_code,
+            headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None),
+        )
+
+    if inspect.isasyncgen(result):
+        sse_headers: Final = {"content-type": "text/event-stream"}
+        return StreamingResponse(
+            content=wrap_passthrough_sse_bytes_with_keepalive_pings(
+                stream=_relay_upstream_bytes(result),
+                ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
+                upstream_headers=sse_headers,
+            ),
+            status_code=200,
+            headers=sse_headers,
+        )
+
+    upstream_stream: Final = cast(AsyncPassthroughStreamingResponse, result)
+    return StreamingResponse(
+        content=wrap_passthrough_sse_bytes_with_keepalive_pings(
+            stream=_relay_upstream_bytes(upstream_stream),
+            ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
+            upstream_headers=upstream_stream.headers,
+        ),
+        status_code=upstream_stream.status_code,
+        headers=HttpPassThroughEndpointHelpers.get_response_headers(
+            headers=upstream_stream.headers, custom_headers=None
+        ),
+    )
+
+
 @router.api_route(
     "/azure_ai/{endpoint:path}",
     methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
@@ -1528,55 +1602,18 @@ async def azure_proxy_route(
             if is_router_model:
                 request_body = await get_request_body(request)
                 is_streaming_request = is_passthrough_request_streaming(request_body)
-                result = await llm_router.allm_passthrough_route(
-                    model=part,
-                    method=request.method,
-                    endpoint=endpoint,
-                    request_query_params=request.query_params,
-                    request_headers=_safe_get_request_headers(request),
-                    stream=is_streaming_request,
-                    content=None,
-                    data=None,
-                    files=None,
-                    json=(request_body if request.headers.get("content-type") == "application/json" else None),
-                    params=None,
-                    headers=None,
-                    cookies=None,
-                    litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict),
-                )
-
-                if is_streaming_request:
-                    # Check if result is an async generator (from _async_streaming)
-                    import inspect
-
-                    if inspect.isasyncgen(result):
-                        # Result is already an async generator, use it directly
-                        return StreamingResponse(
-                            content=result,
-                            status_code=200,
-                            headers={"content-type": "text/event-stream"},
-                        )
-                    else:
-                        # Result is an httpx.Response, use aiter_bytes()
-                        result = cast(httpx.Response, result)
-                        return StreamingResponse(
-                            content=result.aiter_bytes(),
-                            status_code=result.status_code,
-                            headers=HttpPassThroughEndpointHelpers.get_response_headers(
-                                headers=result.headers,
-                                custom_headers=None,
-                            ),
-                        )
-
-                # Non-streaming response
-                result = cast(httpx.Response, result)
-                content = await result.aread()
-                return Response(
-                    content=content,
-                    status_code=result.status_code,
-                    headers=HttpPassThroughEndpointHelpers.get_response_headers(
-                        headers=result.headers,
-                        custom_headers=None,
+                return await open_sse_before_first_byte(
+                    _relay_azure_router_model(
+                        llm_router=llm_router,
+                        model=part,
+                        endpoint=endpoint,
+                        request=request,
+                        request_body=request_body,
+                        is_streaming_request=is_streaming_request,
+                        user_api_key_dict=user_api_key_dict,
+                    ),
+                    ping_interval_seconds=(
+                        litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None
                     ),
                 )
             elif is_vector_store_index:
@@ -1659,6 +1696,12 @@ async def azure_proxy_route(
 
 from abc import ABC, abstractmethod
 
+_VERTEX_LOCATION_REQUIRED_DETAIL: Final = (
+    "No Vertex AI location for this request. Include /projects/<project>/locations/<location>/ in the "
+    "route, set vertex_location in default_vertex_config (or DEFAULT_VERTEXAI_LOCATION), or add the "
+    "model to model_list with use_in_pass_through: true."
+)
+
 
 class BaseVertexAIPassThroughHandler(ABC):
     @staticmethod
@@ -1666,29 +1709,18 @@ class BaseVertexAIPassThroughHandler(ABC):
     def get_default_base_target_url(vertex_location: str | None) -> str:
         pass
 
-    @staticmethod
-    @abstractmethod
-    def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
-        pass
-
 
 class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler):
     @staticmethod
     def get_default_base_target_url(vertex_location: str | None) -> str:
         return "https://discoveryengine.googleapis.com/"
 
-    @staticmethod
-    def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
-        return base_target_url
-
 
 class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler):
     @staticmethod
     def get_default_base_target_url(vertex_location: str | None) -> str:
-        return get_vertex_base_url(vertex_location)
-
-    @staticmethod
-    def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
+        if vertex_location is None:
+            raise HTTPException(status_code=400, detail=_VERTEX_LOCATION_REQUIRED_DETAIL)
         return get_vertex_base_url(vertex_location)
 
 
@@ -1911,10 +1943,8 @@ async def _prepare_vertex_auth_headers(
     router_credentials: LiteLLM_ManagedVectorStore | None,
     vertex_project: str | None,
     vertex_location: str | None,
-    base_target_url: str | None,
-    get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler,
     user_api_key_dict: UserAPIKeyAuth,
-) -> tuple[Mapping[str, str], str | None, bool, str | None, str | None]:
+) -> tuple[Mapping[str, str], bool, str | None, str | None]:
     """
     Prepare authentication headers for Vertex AI pass-through requests.
 
@@ -1924,15 +1954,12 @@ async def _prepare_vertex_auth_headers(
         router_credentials: Optional vector store credentials from registry
         vertex_project: Vertex project ID
         vertex_location: Vertex location
-        base_target_url: Base URL for the Vertex AI service
-        get_vertex_pass_through_handler: Handler for the specific Vertex AI service
         user_api_key_dict: The caller's resolved authentication, so only the secret that
             authenticated them is stripped on the credential-less branch
 
     Returns:
         tuple containing:
             - headers: dict - Authentication headers to use
-            - base_target_url: str | None - Updated base target URL
             - headers_passed_through: bool - Whether headers were passed through from request
             - vertex_project: str | None - Updated vertex project ID
             - vertex_location: str | None - Updated vertex location
@@ -1985,14 +2012,8 @@ async def _prepare_vertex_auth_headers(
         # Add the Authorization header with vendor credentials
         headers["Authorization"] = f"Bearer {auth_header}"
 
-        if base_target_url is not None:
-            base_target_url = get_vertex_pass_through_handler.update_base_target_url_with_credential_location(
-                base_target_url, vertex_location
-            )
-
     return (
         headers,
-        base_target_url,
         headers_passed_through,
         vertex_project,
         vertex_location,
@@ -2085,12 +2106,9 @@ async def _base_vertex_proxy_route(
         location=vertex_location,
     )
 
-    base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
-
     # Prepare authentication headers
     (
         headers,
-        base_target_url,
         headers_passed_through,
         vertex_project,
         vertex_location,
@@ -2100,13 +2118,10 @@ async def _base_vertex_proxy_route(
         router_credentials=router_credentials,
         vertex_project=vertex_project,
         vertex_location=vertex_location,
-        base_target_url=base_target_url,
-        get_vertex_pass_through_handler=get_vertex_pass_through_handler,
         user_api_key_dict=user_api_key_dict,
     )
 
-    if base_target_url is None:
-        base_target_url = get_vertex_base_url(vertex_location)
+    base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
 
     request_route: Final = encoded_endpoint
     verbose_proxy_logger.debug("request_route %s", request_route)

litellm/proxy/proxy_server.py

@@ -583,6 +583,11 @@ try:
 except ImportError:
     build_billing_metrics_recorder = None
     shutdown_billing_metrics_recorder = None
+from litellm.proxy.middleware.admission_control_middleware import (
+    AdmissionControlMiddleware,
+    admission_control_state,
+    get_admission_control_settings,
+)
 from litellm.proxy.middleware.in_flight_requests_middleware import (
     InFlightRequestsMiddleware,
 )
@@ -15265,20 +15270,33 @@ async def async_queue_request(
 
         if llm_router is None:
             raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value})
-
-        response: Final = await llm_router.schedule_acompletion(**data)
+        router: Final = llm_router
 
         if "stream" in data and data["stream"] is True:  # use generate_responses to stream responses
-            return StreamingResponse(
-                async_data_generator(
-                    user_api_key_dict=user_api_key_dict,
-                    response=response,
-                    request_data=data,
-                    request=request,
-                ),
-                media_type="text/event-stream",
+
+            async def produce_queue_stream() -> StreamingResponse:
+                return StreamingResponse(
+                    async_data_generator(
+                        user_api_key_dict=user_api_key_dict,
+                        response=await router.schedule_acompletion(**data),
+                        request_data=data,
+                        request=request,
+                    ),
+                    media_type="text/event-stream",
+                )
+
+            async def audit_late_failure(exc: Exception) -> HTTPException | None:
+                return await proxy_logging_obj.post_call_failure_hook(
+                    user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=data
+                )
+
+            return await open_sse_before_first_byte(
+                produce_queue_stream(),
+                ping_interval_seconds=ttft_keepalive_interval(data, router),
+                on_late_failure=audit_late_failure,
             )
 
+        response: Final = await router.schedule_acompletion(**data)
         fastapi_response.headers.update({"x-litellm-priority": str(data["priority"])})
         return response
     except Exception as e:
@@ -16534,6 +16552,9 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
     {
         "max_parallel_requests": "Integer",
         "global_max_parallel_requests": "Integer",
+        "max_in_flight_requests_per_worker": "Integer",
+        "max_queued_requests_per_worker": "Integer",
+        "admission_queue_timeout_seconds": "Float",
         "max_request_size_mb": "Integer",
         "max_batch_file_size_mb": "Integer",
         "max_file_size_mb": "Integer",
@@ -18209,6 +18230,11 @@ app.add_middleware(
     get_max_request_size_mb=lambda: general_settings.get("max_request_size_mb"),
     is_request_size_limit_enabled=lambda: premium_user is True,
 )
+app.add_middleware(
+    AdmissionControlMiddleware,
+    get_settings=lambda: get_admission_control_settings(general_settings),
+    state=admission_control_state,
+)
 
 
 async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse":

litellm/proxy/public_endpoints/public_v1/model_hub.py

@@ -3,7 +3,7 @@
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from types import MappingProxyType
-from typing import Annotated, Final, Protocol
+from typing import Annotated, Final, Literal, Protocol
 
 from fastapi import APIRouter, Depends, Request
 from typing_extensions import ReadOnly, TypedDict
@@ -20,10 +20,12 @@ from litellm.proxy.list_api.list_framework import (
     Scope,
     ScopeAll,
     SortKey,
+    handle_facet,
     handle_list,
 )
 from litellm.proxy.utils import PrismaClient
 from litellm.types.proxy.management_endpoints.management_v1 import (
+    FacetListResponse,
     ListResponse,
     ProblemDetail,
 )
@@ -95,16 +97,37 @@ class HealthEnricher:
         return tuple(_with_health(row, health.get(row.model_group)) for row in rows)
 
 
+FEATURE_PREFIX: Final = "supports_"
+
+
+def _features(row: ModelGroupInfoProxy) -> tuple[str, ...]:
+    """A row's capabilities as one repeated field, so selecting two of them matches either.
+
+    The hub's feature control has always been a multi-select over the `supports_*` flags.
+    One boolean filter per flag would AND them, which is the opposite of what it does.
+    """
+    return tuple(
+        sorted(
+            name.removeprefix(FEATURE_PREFIX)
+            for name, value in row.model_dump().items()
+            if name.startswith(FEATURE_PREFIX) and value is True
+        )
+    )
+
+
 def _cells(row: ModelGroupInfoProxy) -> Cells:
     return MappingProxyType(
         {
             "model_group": row.model_group,
             "mode": row.mode,
             "providers": tuple(row.providers),
+            "features": _features(row),
             "max_input_tokens": row.max_input_tokens,
             "max_output_tokens": row.max_output_tokens,
             "input_cost_per_token": row.input_cost_per_token,
             "output_cost_per_token": row.output_cost_per_token,
+            "rpm": row.rpm,
+            "tpm": row.tpm,
         }
     )
 
@@ -126,20 +149,28 @@ def _scope(_caller: UserAPIKeyAuth) -> Scope:
 MODEL_HUB_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType(
     {
         "mode": FilterSpec(type=str, ops=frozenset(("eq", "in"))),
-        "providers": FilterSpec(type=str, ops=frozenset(("contains",))),
+        "providers": FilterSpec(type=str, ops=frozenset(("contains", "in"))),
+        "features": FilterSpec(type=str, ops=frozenset(("in",))),
     }
 )
 
+MODEL_HUB_FACETS: Final[Mapping[str, str]] = MappingProxyType(
+    {"providers": "providers", "modes": "mode", "features": "features"}
+)
+
 MODEL_HUB_LIST_SPEC: Final[ListSpec[ModelGroupInfoProxy, ModelGroupInfoProxy]] = ListSpec(
     resource="model groups",
     sortable=frozenset(
         (
             "model_group",
             "mode",
+            "providers",
             "max_input_tokens",
             "max_output_tokens",
             "input_cost_per_token",
             "output_cost_per_token",
+            "rpm",
+            "tpm",
         )
     ),
     searchable=frozenset(("model_group",)),
@@ -153,6 +184,32 @@ MODEL_HUB_LIST_SPEC: Final[ListSpec[ModelGroupInfoProxy, ModelGroupInfoProxy]] =
 )
 
 
+def _published_rows() -> Sequence[ModelGroupInfoProxy]:
+    from litellm.proxy.proxy_server import (
+        _get_model_group_info,  # pyright: ignore[reportPrivateUsage]  # /public/model_hub imports it the same way
+        llm_router,
+    )
+
+    if llm_router is None:
+        raise ManagementProblem(
+            ProblemDetail(
+                type=f"{PROBLEM_TYPE_BASE}no-llm-router",
+                title="No models configured",
+                status=400,
+                detail=CommonProxyErrors.no_llm_router.value,
+            )
+        )
+    if litellm.public_model_groups is None:
+        return ()
+    return tuple(
+        _get_model_group_info(
+            llm_router=llm_router,
+            all_models_str=litellm.public_model_groups,
+            model_group=None,
+        )
+    )
+
+
 def _executor(
     rows: Sequence[ModelGroupInfoProxy],
     prisma_client: PrismaClient | None,
@@ -191,37 +248,11 @@ async def public_model_hub_list(
     ```
     """
     try:
-        from litellm.proxy.proxy_server import (
-            _get_model_group_info,  # pyright: ignore[reportPrivateUsage]  # /public/model_hub imports it the same way
-            llm_router,
-            prisma_client,
-        )
-
-        if llm_router is None:
-            raise ManagementProblem(
-                ProblemDetail(
-                    type=f"{PROBLEM_TYPE_BASE}no-llm-router",
-                    title="No models configured",
-                    status=400,
-                    detail=CommonProxyErrors.no_llm_router.value,
-                )
-            )
-
-        rows: Final[Sequence[ModelGroupInfoProxy]] = (
-            ()
-            if litellm.public_model_groups is None
-            else tuple(
-                _get_model_group_info(
-                    llm_router=llm_router,
-                    all_models_str=litellm.public_model_groups,
-                    model_group=None,
-                )
-            )
-        )
+        from litellm.proxy.proxy_server import prisma_client
 
         return await handle_list(
             spec=MODEL_HUB_LIST_SPEC,
-            executor=_executor(rows, prisma_client),
+            executor=_executor(_published_rows(), prisma_client),
             request=request,
             caller=user_api_key_dict,
         )
@@ -240,3 +271,53 @@ async def public_model_hub_list(
                 detail="Failed to list public model groups.",
             )
         )
+
+
+@router.get(
+    "/model_hub/{facet}",
+    tags=["public", "model management"],  # mutable-ok: fastapi types tags as list[str | Enum]
+    dependencies=(Depends(user_api_key_auth),),
+    response_model=FacetListResponse,
+)
+async def public_model_hub_facet(
+    request: Request,
+    facet: Literal["providers", "modes", "features"],
+    user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+) -> FacetListResponse:
+    """
+    The distinct providers, modes or features across the published model groups, for the
+    Model Hub's filter dropdowns. No authentication.
+
+    Carries the same filters and search as the list route, so a dropdown offers exactly
+    the values the table can show: asking for providers under `filter[mode][in]=chat`
+    lists only the providers that serve a chat model.
+
+    Example curl:
+    ```
+    curl --location --globoff \
+        'http://0.0.0.0:4000/public/v1/model_hub/providers?filter[mode][in]=chat&page_size=50'
+    ```
+    """
+    try:
+        return await handle_facet(
+            spec=MODEL_HUB_LIST_SPEC,
+            executor=InMemoryListExecutor(rows=_published_rows(), cells=_cells),
+            request=request,
+            caller=user_api_key_dict,
+            field=MODEL_HUB_FACETS[facet],
+        )
+
+    except ManagementProblem:
+        raise
+    except Exception as e:  # noqa: BLE001  # a router error answers as a problem document, not the OpenAI error shape
+        verbose_proxy_logger.exception(
+            "litellm.proxy.public_endpoints.public_v1.model_hub.public_model_hub_facet(): Exception occured - %s", e
+        )
+        raise ManagementProblem(
+            ProblemDetail(
+                type=f"{PROBLEM_TYPE_BASE}internal-server-error",
+                title="Internal server error",
+                status=500,
+                detail="Failed to list public model group values.",
+            )
+        )

litellm/proxy/rag_endpoints/endpoints.py

@@ -23,11 +23,14 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
 from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
     LiteLLM_ManagedVectorStore,
 )
-from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.proxy._types import *
 from litellm.proxy.auth.auth_utils import is_request_body_safe
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
-from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.proxy.common_request_processing import (
+    ProxyBaseLLMRequestProcessing,
+    open_sse_before_first_byte,
+    ttft_keepalive_interval,
+)
 from litellm.proxy.common_utils.http_parsing_utils import (
     _read_request_body,
     _safe_get_request_headers,
@@ -48,6 +51,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
     assert_user_can_access_vector_store_id,
 )
 from litellm.repositories.table_repositories import ManagedVectorStoresRepository
+from litellm.types.utils import ModelResponse
 
 if TYPE_CHECKING:
     from litellm.proxy.utils import PrismaClient
@@ -756,43 +760,53 @@ async def rag_query(
             merged_retrieval_config.get("custom_llm_provider"),
         )
 
-        # Call query
-        response: Final = await litellm.aquery(
-            model=model,
-            messages=messages,
-            retrieval_config=merged_retrieval_config,
-            vector_store_params=store_data,
-            rerank=rerank,
-            stream=stream,
-            router=llm_router,
-            **request_data,
-        )
-
-        hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
-        custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
-            user_api_key_dict=user_api_key_dict,
-            call_id=hidden_params.get("litellm_call_id", None) or "",
-            model_id=hidden_params.get("model_id", None) or "",
-            cache_key=hidden_params.get("cache_key", None) or "",
-            api_base=hidden_params.get("api_base", None) or "",
-            version=version,
-            response_cost=hidden_params.get("response_cost", None),
-            request_data=request_data,
-        )
-
-        if isinstance(response, CustomStreamWrapper):
-            return StreamingResponse(
-                select_data_generator(
-                    response=response,
-                    user_api_key_dict=user_api_key_dict,
-                    request_data=request_data,
-                    request=request,
-                ),
-                media_type="text/event-stream",
-                headers=custom_headers,
+        async def query() -> ModelResponse:
+            return await litellm.aquery(
+                model=model,
+                messages=messages,
+                retrieval_config=merged_retrieval_config,
+                vector_store_params=store_data,
+                rerank=rerank,
+                stream=stream,
+                router=llm_router,
+                **request_data,
             )
 
-        fastapi_response.headers.update(custom_headers)
+        def custom_headers_for(response: ModelResponse) -> Mapping[str, str]:
+            hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
+            return ProxyBaseLLMRequestProcessing.get_custom_headers(
+                user_api_key_dict=user_api_key_dict,
+                call_id=hidden_params.get("litellm_call_id", None) or "",
+                model_id=hidden_params.get("model_id", None) or "",
+                cache_key=hidden_params.get("cache_key", None) or "",
+                api_base=hidden_params.get("api_base", None) or "",
+                version=version,
+                response_cost=hidden_params.get("response_cost", None),
+                request_data=request_data,
+            )
+
+        if stream:
+
+            async def produce_stream() -> StreamingResponse:
+                response: Final = await query()
+                return StreamingResponse(
+                    select_data_generator(
+                        response=response,
+                        user_api_key_dict=user_api_key_dict,
+                        request_data=request_data,
+                        request=request,
+                    ),
+                    media_type="text/event-stream",
+                    headers=custom_headers_for(response),
+                )
+
+            return await open_sse_before_first_byte(
+                produce_stream(),
+                ping_interval_seconds=ttft_keepalive_interval(data, llm_router),
+            )
+
+        response: Final = await query()
+        fastapi_response.headers.update(custom_headers_for(response))
         return response
 
     except HTTPException:

litellm/proxy/spend_tracking/spend_management_endpoints.py

@@ -3,7 +3,8 @@ import collections
 import json
 import os
 from collections.abc import Mapping, Sequence
-from datetime import datetime, timedelta, timezone
+from datetime import date, datetime, timedelta, timezone
+from itertools import groupby
 from types import MappingProxyType
 from typing import (
     TYPE_CHECKING,
@@ -16,7 +17,6 @@ from typing import (
     TypeAlias,
     TypedDict,
     TypeVar,
-    cast,  # noqa: TID251  # prisma group_by returns untyped aggregate mappings
 )
 
 import fastapi
@@ -201,16 +201,12 @@ class _SessionSpendStats(NamedTuple):
 _SessionSpendMap: TypeAlias = Mapping[tuple[str, str], _SessionSpendStats]
 
 
-class _SpendSumAggregate(TypedDict, total=False):
-    spend: ReadOnly[float]
-
-
-class _SpendGroupByRow(TypedDict):
+class _SpendDailySummaryRow(TypedDict):
+    day: ReadOnly[str]
     api_key: ReadOnly[str]
     user: ReadOnly[str | None]
     model: ReadOnly[str]
-    startTime: ReadOnly[object]
-    _sum: ReadOnly[_SpendSumAggregate]
+    spend: ReadOnly[float]
 
 
 async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]:
@@ -251,6 +247,66 @@ def _verification_token_table(prisma_client: PrismaClient) -> _VerificationToken
     return VerificationTokenRepository(prisma_client).table
 
 
+def _spend_logs_daily_summary_sql(
+    *,
+    start_date_iso: str,
+    end_date_iso: str,
+    api_key: str | None,
+    request_id: str | None,
+    user_id: str | None,
+) -> tuple[str, tuple[object, ...]]:
+    filter_params: Final[tuple[tuple[str, object], ...]] = tuple(
+        (column, value)
+        for column, value in (
+            ("api_key", api_key),
+            ("request_id", request_id),
+            ('"user"', user_id),
+        )
+        if value is not None
+    )
+    filter_clauses: Final[tuple[str, ...]] = tuple(
+        f"AND {column} = ${index}" for index, (column, _) in enumerate(filter_params, start=3)
+    )
+    filter_sql: Final = "\n".join(filter_clauses)
+    sql_query: Final = f"""
+SELECT
+    to_char(date_trunc('day', "startTime"), 'YYYY-MM-DD') AS day,
+    api_key,
+    "user",
+    model,
+    SUM(spend) AS spend
+FROM "LiteLLM_SpendLogs"
+WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') AND "startTime" <= ($2::timestamptz AT TIME ZONE 'UTC')
+{filter_sql}
+GROUP BY 1, 2, 3, 4
+ORDER BY 1
+"""
+    params: Final[tuple[object, ...]] = (
+        start_date_iso,
+        end_date_iso,
+        *(value for _, value in filter_params),
+    )
+    return sql_query, params
+
+
+def _sum_spend_by(
+    rows: Sequence[_SpendDailySummaryRow], column: Literal["api_key", "user", "model"]
+) -> Mapping[str | None, float]:
+    keys: Final = frozenset(row[column] for row in rows)
+    return {key: sum(float(row["spend"]) for row in rows if row[column] == key) for key in keys}
+
+
+def _daily_summary_item(summary_date: date, rows: Sequence[_SpendDailySummaryRow]) -> Mapping[str, object]:
+    api_key_spend: Final = {key: value for key, value in _sum_spend_by(rows, "api_key").items() if key is not None}
+    return {
+        **api_key_spend,
+        "startTime": summary_date,
+        "spend": sum(float(row["spend"]) for row in rows),
+        "users": _sum_spend_by(rows, "user"),
+        "models": _sum_spend_by(rows, "model"),
+    }
+
+
 async def _find_spend_logs(
     prisma_client: PrismaClient,
     where: Mapping[str, object],
@@ -3266,18 +3322,22 @@ async def view_spend_logs(
             start_date_iso: Final = start_date_obj.isoformat()
             end_date_iso: Final = end_date_obj.isoformat()
 
-            filter_query: Final = {
+            filter_query: Final[
+                dict[str, object]
+            ] = {  # mutable-ok: legacy filters are extended for optional parameters
                 "startTime": {
                     "gte": start_date_iso,  # Greater than or equal to Start Date
                     "lte": end_date_iso,  # Less than or equal to End Date
                 }
             }
 
+            summary_api_key: Final[str | None] = (
+                prisma_client.hash_token(token=api_key)
+                if api_key is not None and api_key.startswith("sk-")
+                else api_key
+            )
             if api_key is not None and isinstance(api_key, str):
-                if api_key.startswith("sk-"):
-                    filter_query["api_key"] = prisma_client.hash_token(token=api_key)
-                else:
-                    filter_query["api_key"] = api_key
+                filter_query["api_key"] = summary_api_key
             if request_id is not None and isinstance(request_id, str):
                 filter_query["request_id"] = request_id
             if user_id is not None and isinstance(user_id, str):
@@ -3296,58 +3356,34 @@ async def view_spend_logs(
                 return data
 
             # Legacy behavior: return summarized data (when summarize=true)
-            # SQL query
-            response: Final = await SpendLogsRepository(prisma_client).table.group_by(
-                by=["api_key", "user", "model", "startTime"],
-                where=filter_query,
-                sum={
-                    "spend": True,
-                },
+            summary_sql_and_params: Final = _spend_logs_daily_summary_sql(
+                start_date_iso=start_date_iso,
+                end_date_iso=end_date_iso,
+                api_key=summary_api_key,
+                request_id=request_id,
+                user_id=user_id,
             )
+            sql_query, params = summary_sql_and_params
+            rows: Final[Sequence[_SpendDailySummaryRow]] = await _query_raw(prisma_client, sql_query, *params)
+            if len(rows) == 0:
+                return []  # pyright: ignore[reportUnknownVariableType]  # empty summary has no element type
 
-            if isinstance(response, list) and len(response) > 0 and isinstance(response[0], dict):
-                spend_rows: Final = cast(Sequence[_SpendGroupByRow], response)  # cast-ok: by/sum fix the shape
-                result: Final[dict] = {}
-                for record in spend_rows:
-                    dt_object = datetime.strptime(str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ")
-                    date = dt_object.date()
-                    if date not in result:
-                        result[date] = {"users": {}, "models": {}}
-                    api_key = record["api_key"]
-                    user_id = record["user"]
-                    model = record["model"]
-                    result[date]["spend"] = result[date].get("spend", 0) + record.get("_sum", {}).get("spend", 0)
-                    result[date][api_key] = result[date].get(api_key, 0) + record.get("_sum", {}).get("spend", 0)
-                    result[date]["users"][user_id] = result[date]["users"].get(user_id, 0) + record.get("_sum", {}).get(
-                        "spend", 0
-                    )
-                    result[date]["models"][model] = result[date]["models"].get(model, 0) + record.get("_sum", {}).get(
-                        "spend", 0
-                    )
-                return_list: Final = []
-                final_date = None
-                for k, v in sorted(result.items()):
-                    return_list.append({**v, "startTime": k})
-                    final_date = k
-
-                end_date_date: Final = end_date_obj.date()
-                if final_date is not None and final_date < end_date_date:
-                    current_date = final_date + timedelta(days=1)
-                    while current_date <= end_date_date:
-                        # Represent current_date as string because original response has it this way
-                        return_list.append(
-                            {
-                                "startTime": current_date,
-                                "spend": 0,
-                                "users": {},
-                                "models": {},
-                            }
-                        )  # If no data, will stay as zero
-                        current_date += timedelta(days=1)  # Move on to the next day
-
-                return return_list
-
-            return response
+            summary_items: Final = tuple(
+                _daily_summary_item(date.fromisoformat(day), tuple(day_rows))
+                for day, day_rows in groupby(rows, key=lambda row: row["day"])
+            )
+            final_date: Final = date.fromisoformat(rows[-1]["day"])
+            end_date_date: Final = end_date_obj.date()
+            padding: Final[tuple[Mapping[str, object], ...]] = tuple(
+                {
+                    "startTime": final_date + timedelta(days=offset),
+                    "spend": 0,
+                    "users": {},
+                    "models": {},
+                }
+                for offset in range(1, (end_date_date - final_date).days + 1)
+            )
+            return [*summary_items, *padding]
 
         else:
             scoped_filter: Final[dict[str, str]] = {}

litellm/router.py

@@ -197,6 +197,7 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import (
 from litellm.scheduler import FlowItem, Scheduler
 from litellm.types.llms.openai import (
     AllMessageValues,
+    ChatCompletionToolParam,
     FileTypes,
     OpenAIFileObject,
     OpenAIFilesPurpose,
@@ -9545,6 +9546,7 @@ class Router:
             public_model_name for _, public_model_name in self.team_model_to_deployment_indices
         )
 
+        self.pattern_router.remove_deployment(model_id)
         for team_id in list(self.team_pattern_routers.keys()):
             team_pattern_router = self.team_pattern_routers[team_id]
             team_pattern_router.remove_deployment(model_id)
@@ -11762,7 +11764,7 @@ class Router:
         self,
         messages: list[dict[str, str]] | None,
         input: str | list | None,
-        instructions: str | None = None,
+        request_kwargs: Mapping[str, object] | None = None,
     ) -> int:
         """
         Count input tokens for context-window pre-call checks.
@@ -11772,9 +11774,28 @@ class Router:
         The Responses payload is normalized to chat messages via the shared
         LiteLLMCompletionResponsesConfig transform so the same token_counter path covers
         both API surfaces and `instructions` tokens are included in the count.
+
+        Prompt content the message list never carries is read from `request_kwargs`:
+        `tools` (Chat Completions, Responses and Anthropic Messages shapes) and the
+        Anthropic Messages top-level `system` block.
         """
+        from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
+            anthropic_system_to_openai_message,
+        )
+
+        extras: Final = request_kwargs if request_kwargs is not None else MappingProxyType({})
+        raw_instructions: Final = extras.get("instructions")
+        instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None
+        raw_tools: Final = extras.get("tools")
+        tools: Final = (
+            cast(list[ChatCompletionToolParam], raw_tools)  # cast-ok: token_counter formats any tool dict shape
+            if isinstance(raw_tools, list) and raw_tools
+            else None
+        )
+        system_message: Final = anthropic_system_to_openai_message(extras.get("system"))
         if messages is not None:
-            return litellm.token_counter(messages=messages)
+            counted_messages: Final = (system_message, *messages) if system_message is not None else messages
+            return litellm.token_counter(messages=counted_messages, tools=tools)
         if input is not None:
             from openai.types.responses.response_create_params import ResponseInputParam
 
@@ -11787,7 +11808,10 @@ class Router:
                 input=typed_input,
                 responses_api_request={"instructions": instructions} if instructions is not None else {},
             )
-            return litellm.token_counter(messages=cast(list, input_messages))  # cast-ok: transformed chat messages
+            return litellm.token_counter(
+                messages=cast(list, input_messages),  # cast-ok: transformed chat messages
+                tools=tools,
+            )
         raise ValueError("Either messages or input must be provided to count tokens")
 
     def _deployment_max_input_tokens(self, model: str, deployment: Mapping[str, object]) -> int | None:
@@ -11833,14 +11857,13 @@ class Router:
         """
         if messages is None and input is None:
             return None
-        raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None
         try:
             if not self._pre_call_checks_need_token_count(model, healthy_deployments):
                 return None
             return await asyncify(self._count_pre_call_check_tokens)(
                 messages=cast(list[dict[str, str]] | None, messages),  # cast-ok: forwarded to the sync counter
                 input=cast(str | list | None, input),  # cast-ok: forwarded to the sync counter
-                instructions=raw_instructions if isinstance(raw_instructions, str) else None,
+                request_kwargs=request_kwargs,
             )
         except Exception as e:  # noqa: BLE001  # best-effort: an uncountable prompt must not fail the request
             verbose_router_logger.error(
@@ -11887,8 +11910,6 @@ class Router:
         _rate_limit_error = False
         parent_otel_span: Final = _get_parent_otel_span_from_kwargs(request_kwargs)
 
-        raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None
-        instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None
         has_countable_input: Final = messages is not None or input is not None
 
         ## get model group RPM ##
@@ -11919,7 +11940,7 @@ class Router:
                             return _returned_deployments
                         try:
                             input_tokens = self._count_pre_call_check_tokens(
-                                messages=messages, input=input, instructions=instructions
+                                messages=messages, input=input, request_kwargs=request_kwargs
                             )
                         except Exception as e:
                             verbose_router_logger.error(

litellm/router_strategy/complexity_router/README.md

@@ -190,6 +190,9 @@ model_list:
 
         # Let that replacement also override a kept session pin, for image turns only (default: false)
         modality_pin_override: true
+
+        # Refreshes on every pin reuse, so this is idle time rather than total session length (default: 3600)
+        session_affinity_ttl_seconds: 300
 ```
 
 ## Usage
@@ -240,6 +243,10 @@ affinity write happens upstream of the gate and stores the session's own model,
 turn replays the original pin and the override is never pinned in its place. It does nothing
 unless `modality_routing` is also on.
 
+### Session pin retention
+
+`session_affinity_ttl_seconds` is the idle window for both the model pin selected by session affinity and the deployment pin. Every request that reuses a pin refreshes its TTL, so a session actively sending requests stays pinned. After the window passes with no pin reuse, the next request classifies again and creates a fresh pin. Omit the setting to track the default of 3600 seconds.
+
 ### Heuristic-first chaining
 
 `classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM

litellm/rust_bridge/__init__.py

@@ -4,6 +4,7 @@ from litellm.rust_bridge.configuration import use_litellm_rust
 from litellm.rust_bridge.loader import (
     get_native_bridge,
     native_bridge_available,
+    reset_native_bridge_cache,
 )
 
-__all__ = ["get_native_bridge", "native_bridge_available", "use_litellm_rust"]
+__all__ = ["get_native_bridge", "native_bridge_available", "reset_native_bridge_cache", "use_litellm_rust"]

litellm/rust_bridge/loader.py

@@ -24,6 +24,12 @@ def get_native_bridge() -> ModuleType | None:
     return _native
 
 
+def reset_native_bridge_cache() -> None:
+    """Forget the cached extension so the next lookup reimports it from disk."""
+    global _cached_bridge
+    _cached_bridge = _BRIDGE_SENTINEL
+
+
 def native_bridge_available() -> bool:
     """Whether the packaged Rust extension is importable."""
     return get_native_bridge() is not None

litellm/secret_managers/get_azure_ad_token_provider.py

@@ -57,9 +57,11 @@ def get_azure_ad_token_provider(
     from azure import identity
     from azure.identity import (
         CertificateCredential,
+        ChainedTokenCredential,
         ClientSecretCredential,
         DefaultAzureCredential,
         ManagedIdentityCredential,
+        WorkloadIdentityCredential,
         get_bearer_token_provider,
     )
 
@@ -101,6 +103,28 @@ def get_azure_ad_token_provider(
         # DefaultAzureCredential doesn't require explicit environment variables
         # It automatically discovers credentials from the environment (managed identity, CLI, etc.)
         credential = DefaultAzureCredential()
+    elif cred == AzureCredentialType.DeploymentIdentityCredential:
+        # DefaultAzureCredential cannot express this: excluding its developer credentials still
+        # leaves one managed identity link, which AZURE_CLIENT_ID pins to a user assigned identity,
+        # so a host running as a system assigned identity never gets asked
+        workload_client_id: Final = os.environ.get("AZURE_CLIENT_ID")
+        workload_tenant_id: Final = os.environ.get("AZURE_TENANT_ID")
+        workload_token_file: Final = os.environ.get("AZURE_FEDERATED_TOKEN_FILE")
+        credential = ChainedTokenCredential(
+            *(
+                (
+                    WorkloadIdentityCredential(
+                        client_id=workload_client_id,
+                        tenant_id=workload_tenant_id,
+                        token_file_path=workload_token_file,
+                    ),
+                )
+                if workload_client_id and workload_tenant_id and workload_token_file
+                else ()
+            ),
+            *((ManagedIdentityCredential(client_id=workload_client_id),) if workload_client_id else ()),
+            ManagedIdentityCredential(),
+        )
     else:
         cred_cls: Final = getattr(identity, cred)
         credential = cred_cls()

litellm/types/caching.py

@@ -18,6 +18,11 @@ class LiteLLMCacheType(str, Enum):
     GCS = "gcs"
 
 
+class SemanticCacheScope(str, Enum):
+    KEY = "key"
+    END_USER = "end_user"
+
+
 CachingSupportedCallTypes = Literal[
     "completion",
     "acompletion",

litellm/types/management_endpoints/cache_settings_endpoints.py

@@ -187,6 +187,19 @@ CACHE_SETTINGS_FIELDS: Final[list[CacheSettingsField]] = [
         ui_field_name="Embedding Model",
         redis_type="semantic",
     ),
+    CacheSettingsField(
+        field_name="semantic_cache_scope",
+        field_type="String",
+        field_value=None,
+        field_description=(
+            "Isolation granularity for semantic cache hits. 'key' shares hits between all end users of a key/team/org."
+            " 'end_user' also isolates per end user; requests without an end user fall back to the key scope."
+        ),
+        field_default="key",
+        options=["key", "end_user"],
+        ui_field_name="Semantic Cache Scope",
+        redis_type="semantic",
+    ),
     # GCP IAM authentication fields
     CacheSettingsField(
         field_name="gcp_service_account",

litellm/types/proxy/management_endpoints/management_v1.py

@@ -1,5 +1,6 @@
 """Shared response shapes for the `/management/v1` control-plane surface."""
 
+from collections.abc import Sequence
 from typing import Generic, TypeVar
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -38,7 +39,7 @@ class PageMeta(BaseModel):
 class FacetListResponse(BaseModel):
     """The distinct values one column takes over a filtered query. `data` holds bare values, not entity rows."""
 
-    data: list[str]
+    data: Sequence[str]
     meta: PageMeta
     links: PageLinks
 

litellm/types/secret_managers/get_azure_ad_token_provider.py

@@ -6,3 +6,4 @@ class AzureCredentialType(str, Enum):
     ManagedIdentityCredential = "ManagedIdentityCredential"
     CertificateCredential = "CertificateCredential"
     DefaultAzureCredential = "DefaultAzureCredential"
+    DeploymentIdentityCredential = "DeploymentIdentityCredential"

ruff-strict-budget.json

@@ -96,7 +96,7 @@
     "limit": 10
   },
   "DTZ007": {
-    "limit": 17
+    "limit": 6
   },
   "DTZ011": {
     "limit": 3

tests/e2e/ui/tests/tables/tableScrolling.spec.ts

@@ -0,0 +1,218 @@
+import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test";
+import { ADMIN_STORAGE_PATH } from "../../constants";
+import { Page } from "../../fixtures/pages";
+import { navigateToPage } from "../../helpers/navigation";
+import { CHAT_MODEL_A, masterKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic";
+
+const VIEWPORT = { width: 1280, height: 720 };
+const SEED_ROWS = 40;
+const LOG_ROWS = 20;
+const BODY_SCROLL_PX = 500;
+const MAX_FOOTER_GAP_PX = 40;
+
+interface GeneratedKey {
+  key: string;
+}
+
+interface CreatedTeam {
+  team_id: string;
+}
+
+interface CreatedModel {
+  model_info: { id: string };
+}
+
+interface BoxMetrics {
+  top: number;
+  bottom: number;
+  scrollHeight: number;
+  clientHeight: number;
+  scrollWidth: number;
+  clientWidth: number;
+}
+
+const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
+
+const adminHeaders = (): Record<string, string> => ({ Authorization: `Bearer ${masterKey()}` });
+
+const appShellMain = (page: PlaywrightPage): Locator => page.locator("main").first();
+
+const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true });
+
+const visibleDataTable = (page: PlaywrightPage): Locator => visibleTestId(page, "data-table-root").first();
+
+const visibleRows = (page: PlaywrightPage): Locator => visibleDataTable(page).locator("tbody tr");
+
+const metrics = (locator: Locator): Promise<BoxMetrics> =>
+  locator.evaluate((el) => {
+    const rect = el.getBoundingClientRect();
+    return {
+      top: rect.top,
+      bottom: rect.bottom,
+      scrollHeight: el.scrollHeight,
+      clientHeight: el.clientHeight,
+      scrollWidth: el.scrollWidth,
+      clientWidth: el.clientWidth,
+    };
+  });
+
+async function postOk<T>(request: APIRequestContext, path: string, data: Record<string, unknown>): Promise<T> {
+  const res = await request.post(path, { headers: adminHeaders(), data });
+  expect(res.ok(), `POST ${path} failed (${res.status()}): ${await res.text()}`).toBe(true);
+  return (await res.json()) as T;
+}
+
+const oneAtATime = <T>(count: number, call: (index: number) => Promise<T>): Promise<readonly T[]> =>
+  Array.from({ length: count }, (_, i) => i).reduce<Promise<readonly T[]>>(
+    async (previous, i) => [...(await previous), await call(i)],
+    Promise.resolve([]),
+  );
+
+async function expectRowsAtLeast(page: PlaywrightPage, count: number): Promise<void> {
+  await expect.poll(() => visibleRows(page).count(), { timeout: 30_000 }).toBeGreaterThanOrEqual(count);
+}
+
+async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): Promise<void> {
+  await visibleTestId(page, "pagination-page-size").click();
+  await page.getByRole("option", { name: size, exact: true }).click();
+}
+
+async function expectBodyIsTheOnlyScroller(page: PlaywrightPage): Promise<void> {
+  const scroller = await metrics(appShellMain(page));
+  const body = visibleTestId(page, "data-table-scroller");
+  const bodyBefore = await metrics(body);
+  const headBefore = await metrics(visibleTestId(page, "data-table-head"));
+  const footer = await metrics(visibleDataTable(page));
+
+  expect(scroller.scrollHeight, "page scroller must not overflow vertically").toBe(scroller.clientHeight);
+  expect(scroller.scrollWidth, "page scroller must not overflow horizontally").toBe(scroller.clientWidth);
+  expect(bodyBefore.scrollHeight, "table body must be the element that scrolls").toBeGreaterThan(
+    bodyBefore.clientHeight,
+  );
+  expect(footer.bottom, "pagination footer must be inside the page").toBeLessThanOrEqual(scroller.bottom);
+  expect(scroller.bottom - footer.bottom, "pagination footer must sit at the bottom of the page").toBeLessThanOrEqual(
+    MAX_FOOTER_GAP_PX,
+  );
+
+  await body.evaluate((el, px) => {
+    el.scrollTop = px;
+  }, BODY_SCROLL_PX);
+  await expect.poll(() => body.evaluate((el) => el.scrollTop)).toBeGreaterThan(0);
+  const headAfter = await metrics(visibleTestId(page, "data-table-head"));
+  expect(Math.round(headAfter.top), "header must stay put while the body scrolls").toBe(Math.round(headBefore.top));
+}
+
+const rowsPaintingPastAnAncestor = (page: PlaywrightPage): Promise<string[]> =>
+  visibleDataTable(page)
+    .locator("table")
+    .evaluate((table) => {
+      const scrollsVertically = (el: Element): boolean =>
+        /auto|scroll/.test(getComputedStyle(el).overflowY) && el.scrollHeight > el.clientHeight + 1;
+      const boxesUpToTheScroller = (el: Element | null): Element[] =>
+        el === null || el === document.body || scrollsVertically(el)
+          ? []
+          : [el, ...boxesUpToTheScroller(el.parentElement)];
+      const describe = (el: Element): string =>
+        `<${el.tagName.toLowerCase()} class="${el.getAttribute("class") ?? ""}">`;
+      return Array.from(table.querySelectorAll("tbody tr")).flatMap((row, index) => {
+        const rowBottom = row.getBoundingClientRect().bottom;
+        return boxesUpToTheScroller(row.parentElement)
+          .filter((box) => rowBottom > box.getBoundingClientRect().bottom + 1)
+          .map(
+            (box) =>
+              `row ${index} bottom ${Math.round(rowBottom)} past ${describe(box)} bottom ${Math.round(box.getBoundingClientRect().bottom)}`,
+          );
+      });
+    });
+
+test.describe("Admin tables scroll inside the page", () => {
+  test.use({ storageState: ADMIN_STORAGE_PATH, viewport: VIEWPORT });
+
+  test("Virtual Keys: rows scroll under a sticky header and the page itself never scrolls", async ({
+    page,
+    request,
+  }) => {
+    const suffix = uniqueSuffix();
+    const keys = await oneAtATime(SEED_ROWS, (i) =>
+      postOk<GeneratedKey>(request, "/key/generate", { key_alias: `e2e-scroll-key-${suffix}-${i}` }),
+    );
+    try {
+      await navigateToPage(page, Page.ApiKeys);
+      await expectRowsAtLeast(page, SEED_ROWS);
+      await expectBodyIsTheOnlyScroller(page);
+    } finally {
+      await request.post("/key/delete", { headers: adminHeaders(), data: { keys: keys.map((k) => k.key) } });
+    }
+  });
+
+  test("Teams: rows scroll under a sticky header and the page itself never scrolls", async ({ page, request }) => {
+    const suffix = uniqueSuffix();
+    const teams = await oneAtATime(SEED_ROWS, (i) =>
+      postOk<CreatedTeam>(request, "/team/new", { team_alias: `e2e-scroll-team-${suffix}-${i}` }),
+    );
+    try {
+      await navigateToPage(page, Page.Teams);
+      await expectRowsAtLeast(page, SEED_ROWS);
+      await expectBodyIsTheOnlyScroller(page);
+    } finally {
+      await request.post("/team/delete", { headers: adminHeaders(), data: { team_ids: teams.map((t) => t.team_id) } });
+    }
+  });
+
+  test("Request Logs: rows scroll under a sticky header and the page itself never scrolls", async ({
+    page,
+    request,
+  }) => {
+    const suffix = uniqueSuffix();
+    const ids = await oneAtATime(LOG_ROWS, (i) =>
+      sendChatCompletion(request, { model: CHAT_MODEL_A, prompt: `scroll ${suffix} ${i}` }),
+    );
+    await waitForSpendLog(request, ids[ids.length - 1]);
+
+    await navigateToPage(page, Page.Logs);
+    await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 });
+    await setRowsPerPage(page, "25");
+    await expectRowsAtLeast(page, LOG_ROWS);
+    await expectBodyIsTheOnlyScroller(page);
+  });
+
+  test("Tags: no row paints past the box it lives in", async ({ page, request }) => {
+    const suffix = uniqueSuffix();
+    const names = Array.from({ length: SEED_ROWS }, (_, i) => `e2e-scroll-tag-${suffix}-${i}`);
+    await oneAtATime(SEED_ROWS, (i) =>
+      postOk<unknown>(request, "/tag/new", { name: names[i], description: "LIT-4738 scroll" }),
+    );
+    try {
+      await navigateToPage(page, Page.TagManagement);
+      await expectRowsAtLeast(page, SEED_ROWS);
+      expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
+    } finally {
+      await oneAtATime(SEED_ROWS, (i) =>
+        request.post("/tag/delete", { headers: adminHeaders(), data: { name: names[i] } }),
+      );
+    }
+  });
+
+  test("Model Hub: no row paints past the box it lives in", async ({ page, request }) => {
+    const suffix = uniqueSuffix();
+    const models = await oneAtATime(SEED_ROWS, (i) =>
+      postOk<CreatedModel>(request, "/model/new", {
+        model_name: `e2e-scroll-model-${suffix}-${i}`,
+        litellm_params: {
+          model: "openai/fake-gpt-4",
+          api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`,
+          api_key: "fake-key",
+        },
+      }),
+    );
+    try {
+      await navigateToPage(page, Page.ModelHubTable);
+      await expectRowsAtLeast(page, SEED_ROWS);
+      expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
+    } finally {
+      await oneAtATime(SEED_ROWS, (i) =>
+        request.post("/model/delete", { headers: adminHeaders(), data: { id: models[i].model_info.id } }),
+      );
+    }
+  });
+});

tests/e2e/ui/tests/users/searchUsers.spec.ts

@@ -17,7 +17,7 @@ test.describe("Internal Users Search", () => {
   test("narrows the table to the matching email, and restores it when cleared", async ({ page }) => {
     await goToInternalUsers(page);
 
-    const search = page.getByPlaceholder("Search by email…");
+    const search = page.getByPlaceholder("Search by email or ID…");
     await expect(search).toBeVisible();
 
     await search.fill("noteam@");

tests/guardrails_tests/test_tracing_guardrails.py

@@ -806,3 +806,91 @@ def test_guardrail_status_fields_computation():
     )
     assert status_fields_no_guardrail.get("llm_api_status") == "success"
     assert status_fields_no_guardrail.get("guardrail_status") == "not_run"
+
+
+@pytest.mark.parametrize(
+    "status, guardrail_information, expected_guardrail_status",
+    [
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": "success"},
+                {"guardrail_status": "guardrail_intervened"},
+            ],
+            "guardrail_intervened",
+            id="pre_call_success_before_blocker",
+        ),
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": "guardrail_intervened"},
+                {"guardrail_status": "success"},
+            ],
+            "guardrail_intervened",
+            id="blocker_before_success",
+        ),
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": "success"},
+                {"guardrail_status": "guardrail_failed_to_respond"},
+            ],
+            "guardrail_failed_to_respond",
+            id="failure_outranks_success",
+        ),
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": "guardrail_failed_to_respond"},
+                {"guardrail_status": "guardrail_intervened"},
+            ],
+            "guardrail_intervened",
+            id="intervention_outranks_failure",
+        ),
+        pytest.param(
+            "success",
+            [
+                {"guardrail_status": "success"},
+                {"guardrail_status": "success"},
+            ],
+            "success",
+            id="all_success_stays_success",
+        ),
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": "some_new_status"},
+                {"guardrail_status": "blocked"},
+            ],
+            "guardrail_intervened",
+            id="unknown_status_does_not_mask_blocker",
+        ),
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": {"unhashable": True}},
+                {"guardrail_status": "guardrail_intervened"},
+            ],
+            "guardrail_intervened",
+            id="unhashable_status_is_skipped",
+        ),
+    ],
+)
+def test_guardrail_status_fields_severity_across_entries(
+    status, guardrail_information, expected_guardrail_status
+):
+    """
+    A blocked request must never be reported as a guardrail success.
+
+    With multiple guardrails on one request (e.g. a pre_call mask that passes,
+    then a post_call guardrail that blocks), entries are recorded in execution
+    order, so the earlier "success" entry must not shadow the later
+    "guardrail_intervened" entry: the aggregate takes the most severe status,
+    regardless of entry order.
+    """
+    from litellm.litellm_core_utils.litellm_logging import _get_status_fields
+
+    fields = _get_status_fields(
+        status=status, guardrail_information=guardrail_information, error_str=None
+    )
+    assert fields.get("guardrail_status") == expected_guardrail_status

tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py

@@ -763,7 +763,7 @@ class _MigrateDeployHarness:
             "_resolve_specific_migration",
             staticmethod(self.resolved.append),
         )
-        monkeypatch.setattr(utils_module.subprocess, "run", self._fake_run)
+        monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run)
         monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None)
 
         self.baseline_succeeds = True

tests/load_tests/test_granian_admission_saturation.py

@@ -0,0 +1,150 @@
+import asyncio
+import os
+import socket
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("LITELLM_RUN_SATURATION_BENCHMARK") != "1",
+    reason="set LITELLM_RUN_SATURATION_BENCHMARK=1 to run the saturation benchmark",
+)
+
+
+def _free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
+        listener.bind(("127.0.0.1", 0))
+        return int(listener.getsockname()[1])
+
+
+def _percentile(values: list[float], percentile: float) -> float:
+    return sorted(values)[min(int(len(values) * percentile), len(values) - 1)]
+
+
+@pytest.mark.asyncio
+async def test_granian_admission_control_saturation(tmp_path: Path) -> None:
+    fake_port: Final = _free_port()
+    proxy_port: Final = _free_port()
+    fake_script: Final = Path(__file__).parents[1] / "_fake_openai_endpoint_server.py"
+    config_path: Final = tmp_path / "saturation_config.yaml"
+    config_path.write_text(
+        f"""model_list:
+  - model_name: slow-endpoint
+    litellm_params:
+      model: openai/slow-endpoint
+      api_base: http://127.0.0.1:{fake_port}/v1
+general_settings:
+  master_key: sk-saturation
+  max_in_flight_requests_per_worker: 8
+  max_queued_requests_per_worker: 8
+  admission_queue_timeout_seconds: 0.5
+"""
+    )
+    fake_process: Final = subprocess.Popen(
+        [sys.executable, str(fake_script), "--host", "127.0.0.1", "--port", str(fake_port)],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    try:
+        proxy_process: Final = subprocess.Popen(
+            [
+                sys.executable,
+                "-m",
+                "litellm.proxy.proxy_cli",
+                "--config",
+                str(config_path),
+                "--run_granian",
+                "--num_workers",
+                "1",
+                "--port",
+                str(proxy_port),
+            ],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        try:
+            async with httpx.AsyncClient(base_url=f"http://127.0.0.1:{proxy_port}") as client:
+                deadline: Final = time.monotonic() + 60
+                while time.monotonic() < deadline:
+                    try:
+                        response: Final = await client.get("/health/liveliness", timeout=2)
+                        if response.status_code == 200:
+                            break
+                    except httpx.HTTPError:
+                        pass
+                    await asyncio.sleep(0.25)
+                else:
+                    raise AssertionError("Granian proxy did not become healthy")
+
+                liveness_latencies: Final[list[float]] = []
+                stop_sampling: Final = asyncio.Event()
+
+                async def sample_liveness() -> None:
+                    while not stop_sampling.is_set():
+                        start: Final = time.perf_counter()
+                        try:
+                            response = await client.get("/health/liveliness", timeout=2)
+                            response.raise_for_status()
+                            liveness_latencies.append(time.perf_counter() - start)
+                        except httpx.HTTPError:
+                            pass
+                        await asyncio.sleep(0.05)
+
+                async def send_completion() -> tuple[int, float, bool]:
+                    start: Final = time.perf_counter()
+                    response = await client.post(
+                        "/chat/completions",
+                        headers={"Authorization": "Bearer sk-saturation"},
+                        json={
+                            "model": "slow-endpoint",
+                            "messages": [{"role": "user", "content": "hello"}],
+                        },
+                        timeout=10,
+                    )
+                    return response.status_code, time.perf_counter() - start, "retry-after" in response.headers
+
+                sampler: Final = asyncio.create_task(sample_liveness())
+                results: Final = await asyncio.gather(*(send_completion() for _ in range(200)))
+                stop_sampling.set()
+                await sampler
+
+            statuses: Final = [result[0] for result in results]
+            latencies: Final = [result[1] for result in results]
+            rejected: Final = [result for result in results if result[0] == 503]
+            assert set(statuses) <= {200, 503}
+            assert rejected
+            assert all(result[2] for result in rejected)
+            assert _percentile(latencies, 0.99) < 5
+            assert liveness_latencies
+            assert _percentile(liveness_latencies, 0.95) < 0.5
+
+            duration: Final = max(latencies)
+            print(
+                "\nmetric          value\n"
+                f"rps             {len(results) / duration:.2f}\n"
+                f"200 count       {statuses.count(200)}\n"
+                f"503 count       {statuses.count(503)}\n"
+                f"p50             {_percentile(latencies, 0.50):.3f}s\n"
+                f"p95             {_percentile(latencies, 0.95):.3f}s\n"
+                f"p99             {_percentile(latencies, 0.99):.3f}s\n"
+                f"liveness p95    {_percentile(liveness_latencies, 0.95):.3f}s"
+            )
+        finally:
+            proxy_process.terminate()
+            try:
+                proxy_process.wait(timeout=10)
+            except subprocess.TimeoutExpired:
+                proxy_process.kill()
+                proxy_process.wait()
+    finally:
+        fake_process.terminate()
+        try:
+            fake_process.wait(timeout=10)
+        except subprocess.TimeoutExpired:
+            fake_process.kill()
+            fake_process.wait()

tests/rust-python-harness/AGENTS.md

@@ -3,42 +3,75 @@
 ```text
 tests/rust-python-harness/
 ├── __main__.py
+├── cli/
+│   ├── __init__.py
+│   ├── catalog.py
+│   └── commands.py
 │
 ├── strategies/
 │   ├── e2e_parity/
-│   │   ├── runner.py
+│   │   ├── __init__.py
+│   │   ├── reporting.py
 │   │   ├── sdk/
-│   │   │   ├── ocr/
-│   │   │   ├── messages/
-│   │   │   ├── chat_completions/
-│   │   │   └── responses/
-│   │   └── gateway/
+│   │   │   └── ocr/
 │   │
 │   ├── trace_parity/
-│   │   ├── runner.py
-│   │   ├── sdk/
-│   │   └── gateway/
+│   │   ├── __init__.py
+│   │   ├── models.py
+│   │   ├── reporting.py
+│   │   └── sdk/
+│   │       ├── chat_completions/
+│   │       ├── messages/
+│   │       ├── ocr/
+│   │       └── transcription/
 │   │
-│   └── unit_tests/
-│       ├── runner.py
-│       ├── mapping_validator.py
-│       ├── python_runner.py
-│       └── rust_runner.py
+│   ├── unit_tests_mapping/
+│   │   ├── __init__.py
+│   │   ├── contracts.py
+│   │   ├── cases/
+│   │   │   └── ocr.py
+│   │   ├── mapping_report.py
+│   │   ├── mappings.py
+│   │   ├── mapping_validator.py
+│   │   ├── reporting.py
+│   │   └── runner.py
+│   │
+│   ├── unit_tests_parity/
+│   │   ├── __init__.py
+│   │   ├── reporting.py
+│   │   └── runner.py
+│   │
+│   └── unit_tests_rust/
+│       ├── __init__.py
+│       ├── reporting.py
+│       └── runner.py
 │
 └── shared/
     ├── parity/
     ├── tracing/
-    └── reporting/
+    ├── reporting/
+    │   └── strategy.py
+    └── unit_runners/
+        └── suite_runner.py
 ```
 
+- A strategy is a folder under `strategies/` with a one-line `AGENTS.md` and an `__init__.py` exporting exactly one `STRATEGY: StrategyDefinition`; its id must equal the folder name
+- `shared/reporting/strategy.py` is the contract: runnable module/suite specs, not-implemented/skipped specs, the runner protocol, and `StrategyDefinition`
+- Every `STRATEGY` explicitly classifies every SDK function; surface-aware strategies declare their surfaces and classify the complete surface-by-function matrix
 - Run locally only; no CI integration
-- `__main__.py` selects strategies and combines their reports; each strategy also runs independently
+- `python -m tests.rust-python-harness run <strategy>|all` runs the selected strategy; `--function` is common, while each strategy exposes only its supported options
+- Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr`
+- `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases
 - `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses
-- `trace_parity/` compares mapped operations, call counts, and required execution ordering
-- E2E and trace runners share orchestration across `sdk/` and `gateway/`; surface-specific execution lives in those folders
-- `unit_tests/runner.py` combines mapping validation, Python test runs, and native Rust test runs
-- `mapping_validator.py` matches Python/Rust tests by agreed names or annotations and reports missing or ambiguous counterparts
-- `python_runner.py` runs existing Python tests with Rust disabled and enabled in separate processes, verifies backend selection, and compares results
-- `rust_runner.py` runs Cargo tests; native Rust unit tests stay beside their implementation
-- `shared/` contains reusable parity, tracing, and reporting machinery
+- `trace_parity/` compares mapped operations, call counts, and required execution ordering; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`)
+- E2E and trace strategies load their registered module cases and run surface-specific execution from their folders
+- `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest
+- `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report
+- `unit_tests_parity/runner.py` runs each contract's `unit_parity_scope` with `LITELLM_RUST=0` and `LITELLM_RUST=1` in separate processes and requires matching outcomes, including failures; exclusions require a reason in the contract
+- `unit_tests_rust/runner.py` runs each contract's focused Cargo test suite; native Rust unit tests stay beside their implementation
+- `shared/unit_runners/suite_runner.py` runs typed suites registered in code with nodeids of the form `suite:<strategy_id>:<function>:<suite>`
+- Every strategy declares its report sections and presentation in its own `reporting.py`; shared reporting code only provides reusable models and cell-formatting primitives
+- `shared/` contains reusable parity, tracing, reporting primitives, and unit-runner machinery
 - Keep fixtures with their owning API and existing Python tests in their current locations
+- Each strategy folder carries an `AGENTS.md` one-liner stating what it should be doing
+- Run the harness's own checks with `uv run pytest -o consider_namespace_packages=true tests/rust-python-harness/shared tests/rust-python-harness/cli tests/rust-python-harness/strategies/unit_tests_mapping tests/rust-python-harness/strategies/unit_tests_parity tests/rust-python-harness/strategies/unit_tests_rust tests/test_rust_python_harness.py -q`

tests/rust-python-harness/README.md

@@ -1,105 +0,0 @@
-# Rust/Python migration harness
-
-This local harness follows [the agreed structure](AGENTS.md). The root command selects strategies and combines their reports. Each strategy has an independent entry point
-
-```text
-strategies/
-  e2e_parity/runner.py
-    sdk/ocr/fixtures/
-    sdk/messages/
-    sdk/chat_completions/
-    sdk/responses/
-    gateway/
-  existing_e2e_test_sdk/runner.py
-  trace_parity/runner.py
-    sdk/
-    gateway/
-  unit_tests/
-    runner.py
-    mapping_validator.py
-    python_runner.py
-    rust_runner.py
-shared/
-  parity/
-  tracing/
-  reporting/
-```
-
-## Run locally
-
-```bash
-uv run python -m tests.rust-python-harness --list
-uv run python -m tests.rust-python-harness --function ocr --plain
-uv run python -m tests.rust-python-harness --strategy e2e_parity --surface sdk --function ocr --plain
-uv run python -m tests.rust-python-harness.strategies.e2e_parity.runner --function ocr --plain
-uv run python -m tests.rust-python-harness.strategies.trace_parity.runner --plain
-uv run python -m tests.rust-python-harness.strategies.unit_tests.runner --plain
-uv run python -m tests.rust-python-harness.strategies.existing_e2e_test_sdk.runner --function transcription --plain
-```
-
-Use `--interactive` for strategy and function selection, `--pytest-arg=-x` to stop pytest on its first failure, and `--coverage` to write Python coverage under `target/rust-python-harness/`. The harness enables pytest namespace-package discovery only for its own invocations
-
-This harness has no CI execution. A configured test that fails or disappears makes the command fail. An unconfigured strategy cell remains planned and contributes no passing evidence. Interruptions and collection errors stop execution; ordinary test failures remain in the combined report while later strategies run
-
-## Strategy responsibilities
-
-E2E parity compares SDK objects, exceptions, callbacks, streams, and provider requests. Gateway tests compare HTTP responses. Both surfaces use the same strategy runner and keep execution details and fixtures in their own folders. OCR has recorded sync/async SDK coverage; the existing Messages and Responses bridge checks remain partial
-
-Trace parity compares operation names through an explicit Python/Rust mapping, call counts, and required completion-before-start ordering with `shared/tracing/compare.py`. Surface tests supply captured operation intervals. No production trace instrumentation or trace case is configured yet
-
-Unit testing combines test mapping validation, separate Python processes with Rust disabled and enabled, backend verification, result comparison, and native Cargo tests. Native tests stay beside their Rust implementation. Existing Python tests stay at their original paths. No complete Python/native unit mapping is configured yet, so these cells remain planned
-
-The existing E2E SDK strategy retains the live provider tests configured upstream. It runs OCR, Chat Completions, and Transcription checks from their existing paths and reports them separately from parity tests. These tests require provider credentials
-
-## Configure cases
-
-Each strategy has a `strategy.json`. Its `functions` object defines SDK cases for OCR, Messages, Responses, Count Tokens, Chat Completions, and Transcription. E2E and trace manifests also accept a `gateway` object keyed by API name. A case has `coverage`, `selectors`, and an optional `note`
-
-```json
-{
-  "coverage": "partial",
-  "selectors": ["tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py"]
-}
-```
-
-Selectors use pytest file or node syntax. A selector ending in `/` includes tests recursively from that directory
-
-Use `planned` with no selectors until an executable contract exists, `partial` for incomplete coverage, `complete` for the full contract, and `not_applicable` when a strategy does not apply. The dashboard shows passing evidence separately from coverage completeness and LOC coverage
-
-Unit cases use `unit_suite` instead of `selectors`, pointing to a repository-relative JSON file with this shape:
-
-```json
-{
-  "python_selectors": ["tests/test_api.py::test_decode"],
-  "cargo_manifest": "litellm-rust/Cargo.toml",
-  "cargo_package": "litellm-core",
-  "cargo_filter": "ocr::",
-  "backend": {
-    "environment_variable": "LITELLM_USE_RUST_OCR",
-    "probe": "tests.rust-python-harness.strategies.unit_tests.python_runner:ocr_backend"
-  },
-  "mappings": [{"python": "tests/test_api.py::test_decode", "rust": "ocr::test_decode"}]
-}
-```
-
-Names match automatically when the collected Python and Rust test names agree. Explicit `mappings` handle different names, class names, and parametrized cases. Missing or ambiguous counterparts fail validation in either direction. The Cargo filter must select the same behavior as the Python selectors
-
-The backend probe returns `python` or `rust` and runs at startup and before every test call, after fixtures have run. The OCR probe verifies the dispatch flag and native extension availability. Surface tests must also assert that calls reach their intended implementation to catch per-call fallback. Python outcomes must agree, and failed runs remain failures even if both backends fail identically
-
-## OCR fixtures
-
-Fixtures, provider configuration, input strategies, and recording commands live in [the OCR package](strategies/e2e_parity/sdk/ocr/fixtures/README.md). Record with provider credentials:
-
-```bash
-uv run python -m tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.fixtures.record --examples 1000
-```
-
-`LITELLM_OCR_FIXTURE_DIR` and `--fixture-dir` override the default directory. Shared recording, replay, comparison, streaming, and cassette persistence live in `shared/parity/`
-
-Run the harness's own checks locally:
-
-```bash
-uv run pytest -o consider_namespace_packages=true tests/rust-python-harness/shared tests/rust-python-harness/strategies/unit_tests tests/test_rust_python_harness.py -q
-```
-
-Existing OCR parity gaps remain visible: invalid-model provider errors differ, Reducto lacks a native contract, and the expanded Azure corpus exposes duplicate Content-Type headers. Moving the harness does not change provider responses or weaken assertions

tests/rust-python-harness/__init__.py

@@ -1,5 +1,4 @@
-"""Interactive Rust/Python SDK parity test harness."""
+from .cli import main
+from .cli.catalog import load_catalog
 
-from .catalog import load_catalog
-
-__all__ = ["load_catalog"]
+__all__ = ["load_catalog", "main"]

tests/rust-python-harness/catalog.py

@@ -1,75 +0,0 @@
-from __future__ import annotations
-
-import json
-from pathlib import Path
-from typing import Final
-
-from pydantic import BaseModel, ConfigDict, ValidationError
-
-from .shared.reporting.models import Coverage, HarnessCase, SDK_FUNCTIONS, Strategy
-
-STRATEGIES_ROOT: Final = Path(__file__).parent / "strategies"
-
-
-class CaseSpec(BaseModel):
-    model_config = ConfigDict(frozen=True, extra="forbid")
-
-    coverage: Coverage
-    selectors: tuple[str, ...] = ()
-    note: str = ""
-    unit_suite: str | None = None
-
-
-class StrategySpec(BaseModel):
-    model_config = ConfigDict(frozen=True, extra="forbid")
-
-    order: int
-    id: str
-    label: str
-    description: str
-    functions: dict[str, CaseSpec]
-    gateway: dict[str, CaseSpec] = {}
-
-
-def _load_strategy(source: Path) -> Strategy:
-    data: Final = StrategySpec.model_validate_json(source.read_text(encoding="utf-8"))
-    if set(data.functions) != set(SDK_FUNCTIONS):
-        raise ValueError(f"{source}: functions must exactly match {SDK_FUNCTIONS}")
-    cases: Final = tuple(
-        HarnessCase(
-            strategy_id=data.id,
-            strategy_label=data.label,
-            sdk_function=name,
-            coverage=case.coverage,
-            selectors=case.selectors,
-            note=case.note,
-            surface=surface,
-            unit_suite=case.unit_suite,
-        )
-        for surface, functions in (("sdk", data.functions), ("gateway", data.gateway))
-        for name in (SDK_FUNCTIONS if surface == "sdk" else functions)
-        for case in (functions[name],)
-    )
-    for case in cases:
-        if case.coverage in {Coverage.PLANNED, Coverage.NOT_APPLICABLE} and (case.selectors or case.unit_suite):
-            raise ValueError(f"{source}: {case.coverage.value} case {case.key} cannot configure tests")
-        if any(not selector.strip() for selector in case.selectors):
-            raise ValueError(f"{source}: empty selector in {case.key}")
-        if data.id == "unit_tests" and case.selectors:
-            raise ValueError(f"{source}: unit_tests must configure unit_suite instead of pytest selectors")
-        if data.id != "unit_tests" and case.unit_suite:
-            raise ValueError(f"{source}: unit_suite is only valid for unit_tests")
-    return Strategy(data.order, data.id, data.label, data.description, source.parent, cases)
-
-
-def load_catalog(root: Path = STRATEGIES_ROOT) -> tuple[Strategy, ...]:
-    sources: Final = tuple(sorted(root.glob("*/strategy.json")))
-    if not sources:
-        raise ValueError(f"No strategy manifests found below {root}")
-    try:
-        strategies: Final = tuple(sorted((_load_strategy(source) for source in sources), key=lambda item: item.order))
-    except (ValidationError, json.JSONDecodeError) as error:
-        raise ValueError(str(error)) from error
-    if len({strategy.id for strategy in strategies}) != len(strategies):
-        raise ValueError(f"Duplicate strategy id in {root}")
-    return strategies

tests/rust-python-harness/cli.py

@@ -1,245 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import importlib.util
-from collections.abc import Sequence
-from pathlib import Path
-
-from .catalog import load_catalog
-from .shared.reporting.models import SDK_FUNCTIONS, HarnessCase, Strategy
-from .shared.reporting.orchestration import StrategyRunner, run_strategies
-from .shared.reporting.ui import make_dashboard
-from .strategies.e2e_parity.runner import run as run_e2e
-from .strategies.existing_e2e_test_sdk.runner import run as run_existing
-from .strategies.trace_parity.runner import run as run_trace
-from .strategies.unit_tests.mapping_validator import FunctionReport, build_function_report
-from .strategies.unit_tests.runner import run as run_units
-
-REPO_ROOT = Path(__file__).resolve().parents[2]
-COVERAGE_ROOT = REPO_ROOT / "target" / "rust-python-harness"
-
-
-def _parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(
-        prog="rust-python-harness",
-        description="Run Rust/Python parity tests with a live strategy-by-SDK-function dashboard.",
-    )
-    parser.add_argument(
-        "-i",
-        "--interactive",
-        action="store_true",
-        help="pick strategies and SDK functions in a guided terminal menu",
-    )
-    parser.add_argument(
-        "--list", action="store_true", help="show the catalog without running tests"
-    )
-    parser.add_argument(
-        "--strategy",
-        action="append",
-        default=[],
-        metavar="ID",
-        help="run only this strategy",
-    )
-    parser.add_argument(
-        "--function",
-        action="append",
-        default=[],
-        dest="sdk_functions",
-        choices=SDK_FUNCTIONS,
-        help="run only this SDK function",
-    )
-    parser.add_argument("--surface", choices=("sdk", "gateway"), help="run only this API surface")
-    parser.add_argument(
-        "--validate-ledger",
-        action="store_true",
-        help=(
-            "report Python<->Rust test-parity ledger gaps and drift instead of "
-            "running the dashboard; narrow with --function"
-        ),
-    )
-    parser.add_argument(
-        "--plain",
-        action="store_true",
-        help="disable the interactive terminal dashboard",
-    )
-    parser.add_argument(
-        "--coverage",
-        action="store_true",
-        help="write Python reference LOC reports (HTML, JSON, and XML)",
-    )
-    parser.add_argument(
-        "--pytest-arg",
-        action="append",
-        default=[],
-        metavar="ARG",
-        help="append an argument to pytest (repeatable, for example --pytest-arg=-x)",
-    )
-    return parser
-
-
-def _coverage_pytest_args(output_root: Path = COVERAGE_ROOT) -> tuple[str, ...]:
-    output_root.mkdir(parents=True, exist_ok=True)
-    return (
-        "--cov=litellm",
-        "--cov-context=test",
-        f"--cov-report=json:{output_root / 'python.json'}",
-        f"--cov-report=xml:{output_root / 'python.xml'}",
-        f"--cov-report=html:{output_root / 'python-html'}",
-    )
-
-
-def _pick_values(
-    title: str, options: Sequence[tuple[str, str]], input_fn=input
-) -> set[str]:
-    print(f"\n{title} (Enter = all)")
-    for index, (value, label) in enumerate(options, start=1):
-        print(f"  {index:>2}. {label}  [{value}]")
-    while True:
-        answer = input_fn("Choose numbers, comma-separated: ").strip()
-        if not answer:
-            return set()
-        try:
-            indexes = {int(part.strip()) for part in answer.split(",")}
-        except ValueError:
-            print("Please enter numbers separated by commas.")
-            continue
-        if indexes and all(1 <= index <= len(options) for index in indexes):
-            return {options[index - 1][0] for index in indexes}
-        print(f"Choose values from 1 to {len(options)}.")
-
-
-def _interactive_filters(strategies: Sequence[Strategy]) -> tuple[set[str], set[str]]:
-    strategy_ids = _pick_values(
-        "Testing strategies", [(strategy.id, strategy.label) for strategy in strategies]
-    )
-    sdk_functions = _pick_values(
-        "SDK functions",
-        [(name, name) for name in SDK_FUNCTIONS],
-    )
-    return strategy_ids, sdk_functions
-
-
-def _select(
-    strategies: Sequence[Strategy], strategy_ids: set[str], sdk_functions: set[str]
-) -> tuple[HarnessCase, ...]:
-    known_ids = {strategy.id for strategy in strategies}
-    unknown = strategy_ids - known_ids
-    if unknown:
-        raise ValueError(f"Unknown strategy: {', '.join(sorted(unknown))}")
-    return tuple(
-        case
-        for strategy in strategies
-        if not strategy_ids or strategy.id in strategy_ids
-        for case in strategy.cases
-        if not sdk_functions or case.sdk_function in sdk_functions
-    )
-
-
-def _print_catalog(strategies: Sequence[Strategy]) -> None:
-    for strategy in strategies:
-        print(f"{strategy.id:20} {strategy.label}")
-        for case in strategy.cases:
-            selectors = (
-                ", ".join(case.selectors) if case.selectors else case.unit_suite or "no test configured"
-            )
-            print(f"  {case.surface}/{case.sdk_function:12} {case.coverage.value:14} {selectors}")
-
-
-def _print_function_report(report: FunctionReport) -> None:
-    print(f"\n{report.sdk_function}")
-    if report.ledger is None or report.audit is None:
-        print("  no ledger yet")
-        return
-    ledger, audit = report.ledger, report.audit
-    print(
-        f"  {ledger.mapped_count}/{ledger.total_count} python tests mapped to rust "
-        f"({ledger.percentage}%)"
-    )
-    print(f"  {len(ledger.rust_only_tests)} rust-only tests with no python counterpart")
-    if audit.is_clean:
-        print("  ledger is in sync with the live test files")
-        return
-    for label, items in (
-        ("ledger references a python test that no longer exists", audit.missing_python_tests),
-        ("python test exists but is not tracked in the ledger", audit.stale_python_tests),
-        ("ledger references a rust test that no longer exists", audit.missing_rust_tests),
-        ("rust test exists but is not tracked in the ledger", audit.stale_rust_tests),
-    ):
-        for item in items:
-            print(f"  {label}: {item}")
-
-
-def _validate_ledger(sdk_functions: set[str]) -> int:
-    functions = sdk_functions or set(SDK_FUNCTIONS)
-    reports = tuple(build_function_report(function) for function in sorted(functions))
-    for report in reports:
-        _print_function_report(report)
-    return 0 if all(report.is_clean for report in reports) else 1
-
-
-def _resolve_runner(strategy_id: str) -> StrategyRunner:
-    match strategy_id:
-        case "e2e_parity":
-            return run_e2e
-        case "trace_parity":
-            return run_trace
-        case "unit_tests":
-            return run_units
-        case "existing_e2e_test_sdk":
-            return run_existing
-        case _:
-            raise ValueError(f"Unknown strategy: {strategy_id}")
-
-
-def main(argv: Sequence[str] | None = None, *, strategy_id: str | None = None) -> int:
-    args = _parser().parse_args(argv)
-    if args.coverage and importlib.util.find_spec("pytest_cov") is None:
-        _parser().error(
-            "--coverage requires the project's pytest-cov dependency; run with "
-            "`poetry run python -m tests.rust-python-harness --coverage`"
-        )
-    if args.validate_ledger:
-        return _validate_ledger(set(args.sdk_functions))
-    catalog = load_catalog()
-    strategies = tuple(strategy for strategy in catalog if strategy_id is None or strategy.id == strategy_id)
-    if args.list:
-        _print_catalog(strategies)
-        return 0
-
-    strategy_ids = set(args.strategy)
-    sdk_functions = set(args.sdk_functions)
-    if args.interactive:
-        picked_strategies, picked_functions = _interactive_filters(strategies)
-        strategy_ids = strategy_ids or picked_strategies
-        sdk_functions = sdk_functions or picked_functions
-
-    try:
-        selected = _select(strategies, strategy_ids, sdk_functions)
-        cases = tuple(case for case in selected if args.surface is None or case.surface == args.surface)
-    except ValueError as exc:
-        _parser().error(str(exc))
-    selected_strategy_ids = {case.strategy_id for case in cases}
-    visible_strategies = tuple(
-        strategy for strategy in strategies if strategy.id in selected_strategy_ids
-    )
-    dashboard = make_dashboard(
-        visible_strategies,
-        plain=args.plain,
-        confidence_strategies=strategies,
-    )
-    pytest_args = [*args.pytest_arg]
-    if args.coverage:
-        pytest_args.extend(_coverage_pytest_args())
-    with dashboard:
-        exit_code, run = run_strategies(
-            cases=cases,
-            repo_root=REPO_ROOT,
-            on_update=dashboard.update,
-            pytest_args=pytest_args,
-            resolve_runner=_resolve_runner,
-        )
-        dashboard.finish(run, exit_code)
-    if args.coverage and (COVERAGE_ROOT / "python.json").exists():
-        print(f"Python LOC heatmap: {COVERAGE_ROOT / 'python-html' / 'index.html'}")
-        print(f"Machine-readable coverage: {COVERAGE_ROOT / 'python.json'}")
-    return exit_code

tests/rust-python-harness/cli/__init__.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import sys
+from collections.abc import Sequence
+from typing import Final, cast
+
+import click
+
+from ..shared.reporting.models import SDK_FUNCTIONS, SdkFunction, Strategy, Surface
+from .catalog import load_catalog
+from .commands import run_command, select_cases
+
+__all__ = ["load_catalog", "main"]
+
+_INTERRUPTED_EXIT_CODE: Final = 130
+
+
+def _function_option() -> click.Option:
+    return click.Option(
+        ("--function", "sdk_functions"),
+        type=click.Choice(SDK_FUNCTIONS),
+        multiple=True,
+        help="run only this SDK function; repeat to select more than one",
+    )
+
+
+def _run_all_command(strategies: Sequence[Strategy]) -> click.Command:
+    def run_all(sdk_functions: tuple[str, ...]) -> int:
+        selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions))
+        cases: Final = select_cases(strategies, selected_functions)
+        return run_command(strategies, cases)
+
+    return click.Command(
+        "all",
+        params=[_function_option()],
+        callback=run_all,
+        help="run every strategy",
+    )
+
+
+def _strategy_command(strategy: Strategy) -> click.Command:
+    params: list[click.Parameter] = [_function_option()]
+    if strategy.definition.surfaces:
+        params.append(
+            click.Option(
+                ("--surface",),
+                type=click.Choice(strategy.definition.surfaces),
+                help="run only this API surface; omit to run every surface",
+            )
+        )
+    runner_argument: Final = strategy.definition.runner_argument
+    if runner_argument is not None:
+        params.append(
+            click.Option(
+                (runner_argument.option, "runner_args"),
+                multiple=True,
+                metavar=runner_argument.metavar,
+                help=runner_argument.help,
+            )
+        )
+
+    def run_strategy(
+        sdk_functions: tuple[str, ...],
+        surface: str | None = None,
+        runner_args: tuple[str, ...] = (),
+    ) -> int:
+        selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions))
+        selected_surface: Final = cast(Surface | None, surface)
+        cases: Final = select_cases((strategy,), selected_functions, selected_surface)
+        return run_command((strategy,), cases, runner_args)
+
+    return click.Command(
+        strategy.id,
+        params=params,
+        callback=run_strategy,
+        help=strategy.description,
+    )
+
+
+def _build_cli(strategies: Sequence[Strategy]) -> click.Group:
+    root: Final = click.Group(
+        "rust-python-harness",
+        help="Run Rust/Python parity tests with raw progress and strategy reports.",
+    )
+    run: Final = click.Group("run", help="run one strategy or the complete harness")
+    run.add_command(_run_all_command(strategies))
+    for strategy in strategies:
+        run.add_command(_strategy_command(strategy))
+    root.add_command(run)
+    return root
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    try:
+        strategies: Final = load_catalog()
+        result: Final = _build_cli(strategies).main(
+            args=None if argv is None else list(argv),
+            prog_name="rust-python-harness",
+            standalone_mode=False,
+        )
+        exit_code: Final = result if isinstance(result, int) else 0
+    except click.ClickException as error:
+        error.show()
+        return error.exit_code
+    except click.Abort:
+        click.echo("Aborted!", err=True)
+        return 1
+    except KeyboardInterrupt:
+        sys.stderr.write("\nInterrupted\n")
+        return _INTERRUPTED_EXIT_CODE
+    if exit_code == _INTERRUPTED_EXIT_CODE:
+        sys.stderr.write("Interrupted\n")
+    return exit_code

tests/rust-python-harness/cli/catalog.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import hashlib
+import importlib
+import importlib.util
+import pkgutil
+import sys
+from pathlib import Path
+from types import ModuleType
+from typing import Final
+
+from .. import strategies as _strategies_package
+from ..shared.reporting.models import SDK_FUNCTIONS, SURFACES, CaseDisposition, HarnessCase, Strategy
+from ..shared.reporting.strategy import StrategyDefinition
+
+_STRATEGIES_PACKAGE: Final = _strategies_package
+STRATEGIES_ROOT: Final = Path(_STRATEGIES_PACKAGE.__path__[0])
+
+
+def _load_strategy_module(name: str, folder: Path, prefix: str | None) -> ModuleType:
+    if prefix is not None:
+        return importlib.import_module(f"{prefix}.{name}")
+    module_name: Final = _synthetic_module_name(folder)
+    spec: Final = importlib.util.spec_from_file_location(
+        module_name, folder / "__init__.py"
+    )
+    if spec is None or spec.loader is None:
+        raise ValueError(f"{folder}: cannot load strategy package")
+    module: Final = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    try:
+        spec.loader.exec_module(module)
+    except Exception as error:
+        del sys.modules[module_name]
+        raise ValueError(f"{folder}: cannot import strategy package: {error}") from error
+    return module
+
+
+def _synthetic_module_name(folder: Path) -> str:
+    digest: Final = hashlib.sha1(str(folder.resolve()).encode()).hexdigest()[:8]
+    return f"_harness_strategy_{folder.name}_{digest}"
+
+
+def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy:
+    module: Final = _load_strategy_module(name, folder, prefix)
+    definition: Final = getattr(module, "STRATEGY", None)
+    if not isinstance(definition, StrategyDefinition):
+        raise ValueError(f"{folder}: __init__.py must export STRATEGY: StrategyDefinition")
+    if definition.id != name:
+        raise ValueError(f"{folder}: strategy id {definition.id!r} must match folder name {name!r}")
+    if definition.directory.resolve() != folder.resolve():
+        raise ValueError(f"{folder}: strategy directory must be {folder}")
+    if len(set(definition.surfaces)) != len(definition.surfaces) or any(
+        surface not in SURFACES for surface in definition.surfaces
+    ):
+        raise ValueError(f"{folder}: invalid strategy surfaces: {definition.surfaces}")
+    keys: Final = tuple((case.surface, case.sdk_function) for case in definition.cases)
+    duplicates: Final = tuple(sorted(key for key in set(keys) if keys.count(key) > 1))
+    if duplicates:
+        raise ValueError(f"{folder}: duplicate strategy cases: {duplicates}")
+    expected: Final = frozenset(
+        (surface, function)
+        for surface in (definition.surfaces or (None,))
+        for function in SDK_FUNCTIONS
+    )
+    actual: Final = frozenset(keys)
+    if actual != expected:
+        missing: Final = tuple(sorted(expected - actual))
+        extra: Final = tuple(sorted(actual - expected))
+        raise ValueError(
+            f"{folder}: strategy cases must exactly match its declared matrix; missing={missing}, extra={extra}"
+        )
+    incompatible: Final = tuple(
+        (case.surface, case.sdk_function)
+        for case in definition.cases
+        if case.spec.disposition is CaseDisposition.RUNNABLE
+        and not isinstance(case.spec, definition.runnable_spec)
+    )
+    if incompatible:
+        raise ValueError(f"{folder}: runnable cases do not match {definition.runnable_spec.__name__}: {incompatible}")
+    cases: Final = tuple(
+        HarnessCase(
+            strategy_id=definition.id,
+            strategy_label=definition.label,
+            sdk_function=case.sdk_function,
+            spec=case.spec,
+            surface=case.surface,
+        )
+        for case in definition.cases
+    )
+    return Strategy(
+        definition.order,
+        definition.id,
+        definition.label,
+        definition.description,
+        definition.directory,
+        cases,
+        definition,
+    )
+
+
+def load_catalog(root: Path | None = None) -> tuple[Strategy, ...]:
+    resolved: Final = STRATEGIES_ROOT if root is None else root
+    prefix: Final = _STRATEGIES_PACKAGE.__name__ if resolved == STRATEGIES_ROOT else None
+    folders: Final = tuple(
+        info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg
+    )
+    if not folders:
+        raise ValueError(f"No strategy packages found below {resolved}")
+    strategies: Final = tuple(
+        _load_strategy(name, resolved / name, prefix) for name in sorted(folders)
+    )
+    ids: Final = [strategy.id for strategy in strategies]
+    if len(set(ids)) != len(ids):
+        raise ValueError(f"Duplicate strategy id in {resolved}")
+    return tuple(sorted(strategies, key=lambda strategy: (strategy.order, strategy.id)))

tests/rust-python-harness/cli/commands.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from collections.abc import Sequence, Set
+from dataclasses import replace
+from pathlib import Path
+from typing import Final
+
+from ..shared.reporting.models import HarnessCase, SdkFunction, Strategy, Surface
+from ..shared.reporting.orchestration import run_strategies
+from ..shared.reporting.ui import make_dashboard
+
+REPO_ROOT: Final = Path(__file__).resolve().parents[3]
+
+
+def select_cases(
+    strategies: Sequence[Strategy],
+    sdk_functions: Set[SdkFunction],
+    surface: Surface | None = None,
+) -> tuple[HarnessCase, ...]:
+    return tuple(
+        case
+        for strategy in strategies
+        for case in strategy.cases
+        if (not sdk_functions or case.sdk_function in sdk_functions)
+        and (surface is None or case.surface == surface)
+    )
+
+
+def run_command(
+    strategies: Sequence[Strategy],
+    cases: Sequence[HarnessCase],
+    runner_args: Sequence[str] = (),
+) -> int:
+    grouped: Final = {
+        strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id)
+        for strategy in strategies
+    }
+    visible: Final = tuple(strategy for strategy in strategies if grouped[strategy.id])
+    runners: Final = tuple(replace(strategy, cases=grouped[strategy.id]) for strategy in visible)
+    dashboard: Final = make_dashboard(visible)
+    with dashboard:
+        exit_code, run = run_strategies(runners, REPO_ROOT, dashboard.update, runner_args)
+        if exit_code != 130:
+            dashboard.finish(run, exit_code)
+    return exit_code

tests/rust-python-harness/cli/test_cli.py

@@ -0,0 +1,458 @@
+from __future__ import annotations
+
+import importlib
+from collections.abc import Callable, Sequence
+from dataclasses import replace
+from pathlib import Path
+from typing import Final
+
+import pytest
+
+from ..shared.reporting.models import (
+    SDK_FUNCTIONS,
+    SURFACES,
+    CaseDisposition,
+    HarnessCase,
+    HarnessRun,
+    RunStatus,
+    Strategy,
+)
+from ..shared.reporting.strategy import NotImplementedCaseSpec, SkippedCaseSpec, StrategyDefinition
+from ..shared.reporting.ui import PlainDashboard, final_report, make_dashboard
+from ..strategies.unit_tests_mapping.mappings import UNIT_TEST_CONTRACTS
+from ..strategies.unit_tests_parity import UNIT_PARITY_SUITES
+from ..strategies.unit_tests_rust import RUST_SUITES
+from . import main
+from .catalog import STRATEGIES_ROOT, load_catalog
+from .commands import REPO_ROOT, select_cases
+
+
+def _strategy_source(
+    *,
+    strategy_id: str = "example",
+    surfaces: tuple[str, ...] = (),
+    drop: tuple[str | None, str] | None = None,
+    duplicate: tuple[str | None, str] | None = None,
+    incompatible: tuple[str | None, str] | None = None,
+) -> str:
+    cells: Final = tuple(
+        (surface, function)
+        for surface in (surfaces or (None,))
+        for function in SDK_FUNCTIONS
+        if (surface, function) != drop
+    )
+    definitions: Final = tuple(
+        (
+            f"strategy.CaseDefinition({function!r}, "
+            "strategy.ModuleCaseSpec(coverage=models.Coverage.COMPLETE, module='tests.example'), "
+            f"surface={surface!r})"
+            if (surface, function) == incompatible
+            else (
+                f"strategy.CaseDefinition({function!r}, "
+                "strategy.NotImplementedCaseSpec(reason='Not implemented yet'), "
+                f"surface={surface!r})"
+            )
+        )
+        for surface, function in (*cells, *((duplicate,) if duplicate is not None else ()))
+    )
+    return (
+        "import importlib\n"
+        "from pathlib import Path\n"
+        "strategy = importlib.import_module('tests.rust-python-harness.shared.reporting.strategy')\n"
+        "models = importlib.import_module('tests.rust-python-harness.shared.reporting.models')\n"
+        "runner = importlib.import_module('tests.rust-python-harness.strategies.trace_parity.runner')\n"
+        "rendering = importlib.import_module('tests.rust-python-harness.shared.reporting.rendering')\n"
+        "def render(results):\n"
+        "    return (rendering.ReportSection('Example outcomes', "
+        "tuple(rendering.render_case_outcome(r) for r in results)),)\n"
+        f"CASES = ({','.join(definitions)},)\n"
+        "STRATEGY = strategy.StrategyDefinition("
+        f"id={strategy_id!r}, order=1, label='Example strategy', description='Example description', "
+        "directory=Path(__file__).parent, runnable_spec=strategy.SuiteCaseSpec, cases=CASES, "
+        f"run=runner.run_trace_cases, render=render, surfaces={surfaces!r})\n"
+    )
+
+
+def _write_strategy_folder(
+    root: Path,
+    name: str = "example",
+    *,
+    init_source: str | None = None,
+) -> Path:
+    folder: Final = root / name
+    folder.mkdir(parents=True)
+    (folder / "__init__.py").write_text(init_source or _strategy_source(), encoding="utf-8")
+    return folder
+
+
+def test_should_load_surface_aware_and_function_only_strategies() -> None:
+    strategies: Final = load_catalog()
+
+    assert [strategy.id for strategy in strategies] == [
+        "e2e_parity",
+        "trace_parity",
+        "unit_tests_mapping",
+        "unit_tests_parity",
+        "unit_tests_rust",
+    ]
+    for strategy in strategies:
+        expected: Final = tuple(
+            (surface, function) for surface in (strategy.definition.surfaces or (None,)) for function in SDK_FUNCTIONS
+        )
+        assert tuple((case.surface, case.sdk_function) for case in strategy.cases) == expected
+
+
+def test_unit_strategies_use_function_only_cases() -> None:
+    strategies: Final = {
+        strategy.id: strategy
+        for strategy in load_catalog()
+        if strategy.id in {"unit_tests_mapping", "unit_tests_parity", "unit_tests_rust"}
+    }
+
+    for sdk_function in SDK_FUNCTIONS:
+        cases: Final = tuple(
+            case for strategy in strategies.values() for case in strategy.cases if case.sdk_function == sdk_function
+        )
+        assert len(cases) == 3
+        assert all(case.surface is None for case in cases)
+        expected_mapping: Final = (
+            CaseDisposition.RUNNABLE if sdk_function in UNIT_TEST_CONTRACTS else CaseDisposition.NOT_IMPLEMENTED
+        )
+        assert cases[0].spec.disposition is expected_mapping
+        expected_parity: Final = (
+            CaseDisposition.RUNNABLE if sdk_function in UNIT_PARITY_SUITES else CaseDisposition.NOT_IMPLEMENTED
+        )
+        expected_rust: Final = (
+            CaseDisposition.RUNNABLE if sdk_function in RUST_SUITES else CaseDisposition.NOT_IMPLEMENTED
+        )
+        assert cases[1].spec.disposition is expected_parity
+        assert cases[2].spec.disposition is expected_rust
+
+
+def test_raw_dashboard_is_always_the_default() -> None:
+    assert isinstance(make_dashboard(load_catalog()), PlainDashboard)
+
+
+def test_every_strategy_folder_complies() -> None:
+    strategies: Final = load_catalog()
+    folders: Final = {
+        path.name for path in STRATEGIES_ROOT.iterdir() if path.is_dir() and (path / "__init__.py").exists()
+    }
+
+    assert folders == {strategy.id for strategy in strategies}
+    for strategy in strategies:
+        definition: Final = strategy.definition
+        assert isinstance(definition, StrategyDefinition)
+        assert definition.directory == strategy.directory
+        assert not (strategy.directory / "strategy.json").exists()
+        assert (strategy.directory / "AGENTS.md").exists()
+        for case in strategy.cases:
+            if case.spec.disposition is CaseDisposition.RUNNABLE:
+                assert isinstance(case.spec, definition.runnable_spec)
+
+
+@pytest.mark.parametrize("surfaces", ((), SURFACES))
+def test_should_reject_a_registry_missing_a_declared_matrix_cell(tmp_path: Path, surfaces: tuple[str, ...]) -> None:
+    surface: Final = surfaces[0] if surfaces else None
+    _write_strategy_folder(
+        tmp_path,
+        init_source=_strategy_source(surfaces=surfaces, drop=(surface, "count_tokens")),
+    )
+
+    with pytest.raises(ValueError, match="must exactly match its declared matrix"):
+        load_catalog(tmp_path)
+
+
+def test_should_reject_a_duplicate_matrix_cell(tmp_path: Path) -> None:
+    _write_strategy_folder(tmp_path, init_source=_strategy_source(duplicate=(None, "ocr")))
+
+    with pytest.raises(ValueError, match="duplicate strategy cases"):
+        load_catalog(tmp_path)
+
+
+def test_should_reject_invalid_declared_surfaces(tmp_path: Path) -> None:
+    _write_strategy_folder(tmp_path, init_source=_strategy_source(surfaces=("sdk", "sdk")))
+
+    with pytest.raises(ValueError, match="invalid strategy surfaces"):
+        load_catalog(tmp_path)
+
+
+def test_should_reject_a_folder_without_a_strategy_definition(tmp_path: Path) -> None:
+    folder: Final = tmp_path / "example"
+    folder.mkdir()
+    (folder / "__init__.py").write_text("VALUE = 1\n", encoding="utf-8")
+
+    with pytest.raises(ValueError, match="STRATEGY"):
+        load_catalog(tmp_path)
+
+
+def test_should_reject_a_strategy_id_that_differs_from_its_folder(tmp_path: Path) -> None:
+    _write_strategy_folder(tmp_path, init_source=_strategy_source(strategy_id="other"))
+
+    with pytest.raises(ValueError, match="must match folder name"):
+        load_catalog(tmp_path)
+
+
+def test_should_reject_a_runnable_case_incompatible_with_the_strategy(tmp_path: Path) -> None:
+    _write_strategy_folder(tmp_path, init_source=_strategy_source(incompatible=(None, "ocr")))
+
+    with pytest.raises(ValueError, match="runnable cases do not match SuiteCaseSpec"):
+        load_catalog(tmp_path)
+
+
+@pytest.mark.parametrize("case_type", (NotImplementedCaseSpec, SkippedCaseSpec))
+def test_should_reject_an_unavailable_case_with_a_blank_reason(
+    case_type: type[NotImplementedCaseSpec] | type[SkippedCaseSpec],
+) -> None:
+    with pytest.raises(ValueError, match="at least 1 character"):
+        case_type(reason=" ")
+
+
+def test_should_select_functions_and_surfaces() -> None:
+    strategy: Final = next(strategy for strategy in load_catalog() if strategy.id == "e2e_parity")
+
+    assert tuple(case.key for case in select_cases((strategy,), {"messages"})) == (
+        "e2e_parity:messages",
+        "e2e_parity:gateway:messages",
+    )
+    assert tuple(case.display_name for case in select_cases((strategy,), {"ocr"}, "gateway")) == ("gateway/ocr",)
+
+
+def _assert_unavailable_cell(strategy: Strategy, case: HarnessCase, section_title: str) -> None:
+    spec: Final = case.spec
+    assert isinstance(spec, (NotImplementedCaseSpec, SkippedCaseSpec))
+    scoped: Final = replace(strategy, cases=(case,))
+    exit_code, run = strategy.definition.run((case,), REPO_ROOT, lambda _: None)
+    result: Final = run.results[case.key]
+    expected: Final = (
+        RunStatus.NOT_IMPLEMENTED if spec.disposition is CaseDisposition.NOT_IMPLEMENTED else RunStatus.SKIPPED
+    )
+    report: Final = final_report(run, exit_code, (scoped,))
+
+    assert exit_code == 0
+    assert result.status is expected
+    assert spec.reason in report
+    assert section_title in report
+    expected_result: Final = "NOT RUN" if expected is RunStatus.NOT_IMPLEMENTED else "SKIPPED"
+    expected_implemented: Final = 0 if expected is RunStatus.NOT_IMPLEMENTED else 1
+    assert f"Result: {expected_result}" in report
+    assert f"Harness support: {expected_implemented}/1 cases implemented" in report
+
+
+def test_every_unavailable_case_finishes_and_explains_itself() -> None:
+    section_titles: Final = {
+        "e2e_parity": "End-to-end parity outcomes",
+        "trace_parity": "trace comparisons",
+        "unit_tests_mapping": "Python/Rust unit-test mappings",
+        "unit_tests_parity": "Python backend parity outcomes",
+        "unit_tests_rust": "Native Rust unit-test outcomes",
+    }
+    unavailable: Final = tuple(
+        (strategy, case)
+        for strategy in load_catalog()
+        for case in strategy.cases
+        if case.spec.disposition is not CaseDisposition.RUNNABLE
+    )
+
+    for strategy, case in unavailable:
+        _assert_unavailable_cell(strategy, case, section_titles[strategy.id])
+
+
+@pytest.mark.parametrize(
+    ("strategy_id", "present", "absent"),
+    (
+        ("e2e_parity", "--surface", "--pytest-arg"),
+        ("trace_parity", "--surface", "--pytest-arg"),
+        ("unit_tests_parity", "--pytest-arg", "--surface"),
+        ("unit_tests_mapping", "--detail", "--surface"),
+        ("unit_tests_rust", "--function", "--surface"),
+    ),
+)
+def test_strategy_help_only_lists_supported_options(
+    strategy_id: str,
+    present: str,
+    absent: str,
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    exit_code: Final = main(["run", strategy_id, "--help"])
+    captured: Final = capsys.readouterr()
+
+    assert exit_code == 0
+    assert present in captured.out
+    assert absent not in captured.out
+
+
+def test_run_help_lists_all_and_every_strategy(capsys: pytest.CaptureFixture[str]) -> None:
+    exit_code: Final = main(["run", "--help"])
+    captured: Final = capsys.readouterr()
+
+    assert exit_code == 0
+    for command in (
+        "all",
+        "e2e_parity",
+        "trace_parity",
+        "unit_tests_mapping",
+        "unit_tests_parity",
+        "unit_tests_rust",
+    ):
+        assert command in captured.out
+
+
+@pytest.mark.parametrize(
+    "argv",
+    (
+        ("list",),
+        ("check",),
+        ("run", "--strategy", "unit_tests_parity"),
+        ("run", "unit_tests_parity", "--surface", "sdk"),
+        ("run", "unit_tests_parity", "--plain"),
+        ("run", "unit_tests_parity", "--runner-arg=-x"),
+        ("run", "all", "--pytest-arg=-x"),
+    ),
+)
+def test_removed_commands_and_options_are_rejected(argv: tuple[str, ...], capsys: pytest.CaptureFixture[str]) -> None:
+    exit_code: Final = main(argv)
+    captured: Final = capsys.readouterr()
+
+    assert exit_code == 2
+    assert captured.err
+
+
+def test_strategy_command_forwards_repeated_filters_and_runner_arguments(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    cli: Final = importlib.import_module("tests.rust-python-harness.cli")
+    captured: list[tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]]] = []
+
+    def capture_run(
+        strategies: Sequence[Strategy],
+        cases: Sequence[HarnessCase],
+        runner_args: Sequence[str] = (),
+    ) -> int:
+        captured.append(
+            (
+                tuple(strategy.id for strategy in strategies),
+                tuple(case.display_name for case in cases),
+                tuple(runner_args),
+            )
+        )
+        return 0
+
+    monkeypatch.setattr(cli, "run_command", capture_run)
+
+    assert (
+        main(
+            [
+                "run",
+                "unit_tests_parity",
+                "--function",
+                "ocr",
+                "--function",
+                "messages",
+                "--pytest-arg=-x",
+            ]
+        )
+        == 0
+    )
+    assert captured == [
+        (("unit_tests_parity",), ("ocr", "messages"), ("-x",)),
+    ]
+
+
+def test_omitted_surface_selects_every_strategy_surface(monkeypatch: pytest.MonkeyPatch) -> None:
+    cli: Final = importlib.import_module("tests.rust-python-harness.cli")
+    selected: list[str] = []
+
+    def capture_run(
+        strategies: Sequence[Strategy],
+        cases: Sequence[HarnessCase],
+        runner_args: Sequence[str] = (),
+    ) -> int:
+        del strategies, runner_args
+        selected.extend(case.display_name for case in cases)
+        return 0
+
+    monkeypatch.setattr(cli, "run_command", capture_run)
+
+    assert main(["run", "e2e_parity", "--function", "ocr"]) == 0
+    assert selected == ["sdk/ocr", "gateway/ocr"]
+
+
+def test_run_all_selects_every_declared_case_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    cli: Final = importlib.import_module("tests.rust-python-harness.cli")
+    selected: list[HarnessCase] = []
+
+    def capture_run(
+        strategies: Sequence[Strategy],
+        cases: Sequence[HarnessCase],
+        runner_args: Sequence[str] = (),
+    ) -> int:
+        del strategies, runner_args
+        selected.extend(cases)
+        return 0
+
+    monkeypatch.setattr(cli, "run_command", capture_run)
+
+    assert main(["run", "all", "--function", "ocr"]) == 0
+    assert len(selected) == 7
+    assert sum(case.surface is None for case in selected) == 3
+    assert sum(case.surface is not None for case in selected) == 4
+
+
+def test_run_reports_not_implemented_surface_as_not_run(
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    exit_code: Final = main(["run", "trace_parity", "--surface", "gateway", "--function", "ocr"])
+    captured: Final = capsys.readouterr()
+
+    assert exit_code == 0
+    assert "Result: NOT RUN" in captured.out
+    assert "Harness support: 0/1 cases implemented" in captured.out
+    assert "Cases: 1 selected, 1 not implemented, 0 skipped" in captured.out
+    assert "Not implemented" in captured.out
+    assert "No gateway OCR trace-parity case is registered." in captured.out
+
+
+def test_keyboard_interrupt_exits_cleanly(
+    monkeypatch: pytest.MonkeyPatch,
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    cli: Final = importlib.import_module("tests.rust-python-harness.cli")
+
+    def interrupt() -> tuple[object, ...]:
+        raise KeyboardInterrupt
+
+    monkeypatch.setattr(cli, "load_catalog", interrupt)
+
+    exit_code: Final = main(["run", "all"])
+    captured: Final = capsys.readouterr()
+
+    assert exit_code == 130
+    assert captured.out == ""
+    assert captured.err == "\nInterrupted\n"
+
+
+def test_runner_interrupt_skips_the_completion_report(
+    monkeypatch: pytest.MonkeyPatch,
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    commands: Final = importlib.import_module("tests.rust-python-harness.cli.commands")
+
+    def interrupt_run(
+        strategies: Sequence[Strategy],
+        repo_root: Path,
+        on_update: Callable[[HarnessRun], None],
+        runner_args: Sequence[str] = (),
+    ) -> tuple[int, HarnessRun]:
+        del repo_root, on_update, runner_args
+        run: Final = HarnessRun.from_cases(case for strategy in strategies for case in strategy.cases)
+        return 130, run
+
+    monkeypatch.setattr(commands, "run_strategies", interrupt_run)
+
+    exit_code: Final = main(["run", "trace_parity", "--surface", "gateway"])
+    captured: Final = capsys.readouterr()
+
+    assert exit_code == 130
+    assert "Rust <-> Python parity report" not in captured.out
+    assert captured.err == "Interrupted\n"

tests/rust-python-harness/conftest.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+from typing import Final
+
+import pytest
+
+HARNESS_ROOT: Final = Path(__file__).resolve().parents[2]
+
+
+@pytest.fixture(autouse=True)
+def subprocess_test_environment(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("PYTHONPATH", str(HARNESS_ROOT))
+    monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1")
+
+
+@pytest.fixture
+def cargo_project(tmp_path: Path) -> Callable[[str, str], Path]:
+    def create(package: str, source: str) -> Path:
+        manifest: Final = tmp_path / "Cargo.toml"
+        manifest.write_text(
+            f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n'
+        )
+        (tmp_path / "src").mkdir()
+        (tmp_path / "src/lib.rs").write_text(source)
+        return manifest
+
+    return create

tests/rust-python-harness/shared/native_build.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import importlib.util
+import os
+import subprocess
+import sys
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Final
+
+from litellm.rust_bridge import get_native_bridge, reset_native_bridge_cache
+
+MATURIN_SPEC: Final = "maturin==1.15.0"
+BRIDGE_FEATURE: Final = "trace-parity"
+_RUST_ROOT: Final = "litellm-rust"
+_LOCKFILE: Final = "Cargo.lock"
+_SOURCE_SUFFIXES: Final = frozenset({".rs", ".toml"})
+_FAILURE_OUTPUT_LINES: Final = 15
+
+
+def needs_rebuild(native_mtime: float | None, newest_source_mtime: float | None) -> bool:
+    if native_mtime is None:
+        return True
+    if newest_source_mtime is None:
+        return False
+    return newest_source_mtime > native_mtime
+
+
+def _source_files(rust_root: Path) -> Iterator[Path]:
+    for path in rust_root.rglob("*"):
+        relative: Final = path.relative_to(rust_root)
+        if "target" in relative.parts or not path.is_file():
+            continue
+        if path.name == _LOCKFILE or path.suffix in _SOURCE_SUFFIXES:
+            yield path
+
+
+def _newest_source_mtime(repo_root: Path) -> float | None:
+    rust_root: Final = repo_root / _RUST_ROOT
+    if not rust_root.is_dir():
+        return None
+    return max((path.stat().st_mtime for path in _source_files(rust_root)), default=None)
+
+
+def _native_module_path() -> Path | None:
+    try:
+        spec: Final = importlib.util.find_spec("litellm.rust_bridge._native")
+    except (ImportError, ValueError):
+        return None
+    origin: Final = getattr(spec, "origin", None)
+    return Path(origin) if origin else None
+
+
+def _drop_imported_bridge() -> None:
+    reset_native_bridge_cache()
+    for name in tuple(sys.modules):
+        if name.startswith("litellm.rust_bridge._native"):
+            del sys.modules[name]
+
+
+def _rebuild(repo_root: Path) -> tuple[bool, str]:
+    command: Final = ("uvx", "--from", MATURIN_SPEC, "maturin", "develop", "--features", BRIDGE_FEATURE)
+    completed: Final = subprocess.run(
+        command,
+        cwd=repo_root,
+        env={**os.environ, "VIRTUAL_ENV": sys.prefix},
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    output: Final = f"{completed.stdout}\n{completed.stderr}".strip()
+    lines: Final = tuple(output.splitlines())
+    return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:])
+
+
+def ensure_trace_bridge(repo_root: Path) -> str | None:
+    native_path: Final = _native_module_path()
+    native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None
+    if needs_rebuild(native_mtime, _newest_source_mtime(repo_root)):
+        print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True)
+        succeeded: Final
+        output: Final
+        succeeded, output = _rebuild(repo_root)
+        if not succeeded:
+            return f"native Rust bridge rebuild failed:\n{output}"
+        _drop_imported_bridge()
+    bridge: Final = get_native_bridge()
+    if bridge is None:
+        return "native Rust bridge is not importable"
+    if getattr(bridge, "_trace", None) is None:
+        return f"native Rust bridge does not expose _trace; it must be built with the {BRIDGE_FEATURE} feature"
+    return None

tests/rust-python-harness/shared/parity/__init__.py

@@ -1,3 +0,0 @@
-import pytest
-
-pytest.register_assert_rewrite("tests.rust-python-harness.shared.parity.compare")