Merge branch 'litellm_internal_staging' of github.com:BerriAI/litellm into litellm_/cranky-hamilton-21b5d0

.dockerignore

@@ -49,6 +49,8 @@ build/
 *.egg-info/
 .DS_Store
 **/node_modules
+ui/litellm-dashboard/.next
+ui/litellm-dashboard/out
 litellm-rust/target/
 litellm/rust_bridge/_native*.so
 *.log

.github/pull_request_template.md

@@ -1,17 +1,17 @@
 ## Relevant issues
 
-<!-- e.g. "Fixes #000" -->
+<!-- e.g., "Fixes #000" -->
 
 ## Linear ticket
 
-<!-- if you are an internal contributor (e.g., your username is postfixed with -berri or -berriai), add "Resolves " followed by the Linear ticket e.g. "Resolves LIT-1234" to magically link the Linear ticket to the GitHub PR -->
+<!-- if you are an internal contributor (e.g., your username is postfixed with -berri or -berriai), add "Resolves " followed by the Linear ticket e.g., "Resolves LIT-1234" to magically link the Linear ticket to the GitHub PR -->
 
 ## Pre-Submission checklist
 
 **Please complete all items before asking a LiteLLM maintainer to review your PR**
 
 - [ ] I have added meaningful tests
-- [ ] My PR passes all unit tests on [`make test-unit`](https://docs.litellm.ai/docs/extras/contributing_code)
+- [ ] My PR passes all CI/CD checks (e.g., lint, format, unit tests)
 - [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem
 - [ ] I have requested a Greptile review by commenting `@greptileai` and received a **Confidence Score of at least 4/5** before requesting a maintainer review
 
@@ -19,29 +19,13 @@
 
 If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slack (#pr-review)](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA).
 
-## CI (LiteLLM team)
-
-> **CI status guideline:**
->
-> - 50-55 passing tests: main is stable with minor issues.
-> - 45-49 passing tests: acceptable but needs attention
-> - <= 40 passing tests: unstable; be careful with your merges and assess the risk.
-
-- [ ] **Branch creation CI run**  
-       Link:
-
-- [ ] **CI run for the last commit**  
-       Link:
-
-- [ ] **Merge / cherry-pick CI run**  
-       Links:
-
 ## Screenshots / Proof of Fix
 
-<!-- Include screenshots, screen recordings, or log output demonstrating that your changes work as expected.
-     For bug fixes: show reproduction before the fix and passing behavior after.
-     For new features: show the feature working end-to-end.
-     For UI changes: include before/after screenshots. -->
+<!-- Include screenshots, screen recordings, or command (e.g., curl) + output demonstrating that your changes work as expected
+     The proof must be completely e2e with no mocks, using, for example, actual LLM calls costing real $. `pytest` commands are not enough
+     For bug fixes: show reproduction before the fix and passing behavior after
+     For new features: show the feature working end-to-end
+     For UI changes: include before/after screenshots -->
 
 ## Type
 

Dockerfile

@@ -1,12 +1,33 @@
+# syntax=docker/dockerfile:1.7
+
 # Base image for building
 ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
 
 # Runtime image
 ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
+# Pinned by digest like the other base images; bump explicitly on Node upgrades.
+ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6
 
 FROM $UV_IMAGE AS uvbin
 
+# Admin UI builder. Pinned to the build platform so the architecture-independent
+# Next.js static export compiles once natively even in a multi-arch build,
+# instead of once per target arch under QEMU.
+FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder
+
+ENV NEXT_TELEMETRY_DISABLED=1 \
+    npm_config_fund=false \
+    npm_config_audit=false
+
+WORKDIR /ui
+
+COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./
+RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline
+
+COPY ui/litellm-dashboard/ ./
+RUN npm run build
+
 # Builder stage
 FROM $LITELLM_BUILD_IMAGE AS builder
 
@@ -48,7 +69,13 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
 # Copy full source tree
 COPY . .
 
-# Build Admin UI before final sync
+# Replace the committed UI bundle with the one built from this exact source.
+# Clearing first drops the committed bundle's content-hashed chunks that COPY
+# would otherwise leave behind alongside the fresh ones.
+RUN rm -rf litellm/proxy/_experimental/out
+COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/
+
+# Build Admin UI before final sync (applies the enterprise color override when present)
 RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
 
 # Install project and workspace packages (fast - deps already cached)

docker/Dockerfile.database

@@ -1,12 +1,33 @@
+# syntax=docker/dockerfile:1.7
+
 # Base image for building
 ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
 
 # Runtime image
 ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
+# Pinned by digest like the other base images; bump explicitly on Node upgrades.
+ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6
 
 FROM $UV_IMAGE AS uvbin
 
+# Admin UI builder. Pinned to the build platform so the architecture-independent
+# Next.js static export compiles once natively even in a multi-arch build,
+# instead of once per target arch under QEMU.
+FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder
+
+ENV NEXT_TELEMETRY_DISABLED=1 \
+    npm_config_fund=false \
+    npm_config_audit=false
+
+WORKDIR /ui
+
+COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./
+RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline
+
+COPY ui/litellm-dashboard/ ./
+RUN npm run build
+
 FROM $LITELLM_BUILD_IMAGE AS builder
 
 WORKDIR /app
@@ -46,7 +67,13 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
 # Copy full source tree
 COPY . .
 
-# Build Admin UI before final sync
+# Replace the committed UI bundle with the one built from this exact source.
+# Clearing first drops the committed bundle's content-hashed chunks that COPY
+# would otherwise leave behind alongside the fresh ones.
+RUN rm -rf litellm/proxy/_experimental/out
+COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/
+
+# Build Admin UI before final sync (applies the enterprise color override when present)
 RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
 
 # Install project and workspace packages (fast - deps already cached)

docker/Dockerfile.non_root

@@ -1,11 +1,32 @@
+# syntax=docker/dockerfile:1.7
+
 # Base images
 ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
 ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
 ARG PROXY_EXTRAS_SOURCE=published
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
+# Pinned by digest like the other base images; bump explicitly on Node upgrades.
+ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6
 
 FROM $UV_IMAGE AS uvbin
 
+# Admin UI builder. Pinned to the build platform so the architecture-independent
+# Next.js static export compiles once natively even in a multi-arch build,
+# instead of once per target arch under QEMU.
+FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder
+
+ENV NEXT_TELEMETRY_DISABLED=1 \
+    npm_config_fund=false \
+    npm_config_audit=false
+
+WORKDIR /ui
+
+COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./
+RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline
+
+COPY ui/litellm-dashboard/ ./
+RUN npm run build
+
 FROM $LITELLM_BUILD_IMAGE AS builder
 ARG PROXY_EXTRAS_SOURCE
 WORKDIR /app
@@ -53,6 +74,12 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
 # Copy full source tree
 COPY . .
 
+# Replace the committed UI bundle with the one built from this exact source.
+# Clearing first drops the committed bundle's content-hashed chunks that COPY
+# would otherwise leave behind alongside the fresh ones.
+RUN rm -rf litellm/proxy/_experimental/out
+COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/
+
 # Set non-root flag for build time consistency
 ENV LITELLM_NON_ROOT=true
 

litellm-rust/Cargo.lock

@@ -673,6 +673,7 @@ dependencies = [
  "serde",
  "serde_json",
  "thiserror 2.0.18",
+ "tokio",
 ]
 
 [[package]]

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

@@ -25,7 +25,7 @@ futures-util.workspace = true
 serde_json.workspace = true
 base64.workspace = true
 axum = { workspace = true, features = ["ws"], optional = true }
-serde = { workspace = true, optional = true }
+serde.workspace = true
 subtle = { workspace = true, optional = true }
 # sha2 hashes the master key into user_api_key_hash (matches the proxy's
 # SHA-256 hash_token) so the plaintext credential never enters a log payload.
@@ -34,7 +34,7 @@ pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
 
 [features]
 default = []
-server = ["dep:axum", "dep:subtle", "dep:serde", "dep:sha2"]
+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"]

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

@@ -26,4 +26,5 @@ pub(crate) const DEFAULT_MAX_BATCH_SIZE: usize = 256;
 pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500;
 
 /// Provider attributed to realtime sessions in the logging payload.
+#[cfg(feature = "server")]
 pub(crate) const DEFAULT_PROVIDER: &str = "openai";

litellm-rust/crates/ai-gateway/src/integrations/README.md

@@ -0,0 +1,127 @@
+# LiteLLM Rust integrations
+
+This directory contains Rust-native equivalents of LiteLLM integration hooks.
+The first supported surfaces are terminal custom loggers and pre/during-call
+custom guardrails.
+
+## File layout
+
+Every integration is a folder:
+
+- `mod.rs` contains the implementation, trait, runner, or adapter
+- `types.rs` contains the integration-local request, response, error, and future
+  types
+
+Do not add new flat integration files such as `custom_logger.rs`. Shared wire
+contracts that are used by multiple integrations can stay in
+`integrations/types.rs`.
+
+Call ordering and lifecycle timing live in `litellm-core/src/call_lifecycle`.
+Call-type modules, such as OCR, adapt their request and response shapes into
+that generic lifecycle runner.
+
+## CustomLogger
+
+Implement `CustomLogger` when Rust code needs to observe terminal success or
+failure events. Method names intentionally match Python `CustomLogger` names.
+
+```rust
+use litellm_ai_gateway::integrations::custom_logger::{
+    CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
+};
+
+struct RecordingLogger;
+
+impl CustomLogger for RecordingLogger {
+    fn async_log_success_event<'a>(
+        &'a self,
+        model_call_details: &'a ModelCallDetails,
+        response_obj: &'a CallbackValue,
+        timing: CallbackTiming,
+    ) -> LogFuture<'a> {
+        Box::pin(async move {
+            let model = &model_call_details.model;
+            let provider = &model_call_details.custom_llm_provider;
+            let call_type = model_call_details.call_type.to_string();
+            let request_id = model_call_details.request_id.as_deref();
+            let response_object = &response_obj.object;
+            let duration = timing.end_time - timing.start_time;
+            let standard_payload = model_call_details.standard_logging_payload.as_ref();
+
+            Ok(())
+        })
+    }
+
+    fn async_log_failure_event<'a>(
+        &'a self,
+        model_call_details: &'a ModelCallDetails,
+        response_obj: Option<&'a CallbackValue>,
+        timing: CallbackTiming,
+    ) -> LogFuture<'a> {
+        Box::pin(async move {
+            let error = model_call_details.failure_error.as_ref();
+            let response_object = response_obj.map(|value| value.object.as_str());
+            let duration = timing.end_time - timing.start_time;
+
+            Ok(())
+        })
+    }
+}
+```
+
+Use `CustomLoggerRunner` to fan out terminal events to configured loggers. The
+runner is a no-op when no loggers are configured, which is the expected fast
+path for requests without callbacks.
+
+## CustomGuardrail
+
+Implement `CustomGuardrail` when Rust code needs to run pre-call or native
+during-call checks. Method names intentionally match Python `CustomGuardrail`
+entrypoints inherited from Python `CustomLogger`.
+
+```rust
+use litellm_ai_gateway::integrations::custom_guardrail::{
+    CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailEventHook,
+    GuardrailFuture, GuardrailRequest,
+};
+
+struct BlocklistedPromptGuardrail;
+
+impl CustomGuardrail for BlocklistedPromptGuardrail {
+    fn guardrail_name(&self) -> &str {
+        "blocklisted-prompt"
+    }
+
+    fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
+        &[GuardrailEventHook::PreCall]
+    }
+
+    fn async_pre_call_hook<'a>(
+        &'a self,
+        _context: &'a GuardrailContext,
+        request: GuardrailRequest,
+    ) -> GuardrailFuture<'a> {
+        Box::pin(async move {
+            if request.data.to_string().contains("blocked phrase") {
+                return Ok(GuardrailDecision::Block(
+                    litellm_ai_gateway::integrations::custom_guardrail::GuardrailError::blocked(
+                        "blocked phrase detected",
+                    ),
+                ));
+            }
+            Ok(GuardrailDecision::Allow(request))
+        })
+    }
+}
+```
+
+Use `CustomGuardrailRunner::run_pre_call` for `pre_call` guardrails and
+`CustomGuardrailRunner::run_during_call` for `during_call` guardrails. A
+`GuardrailDecision::Mask` continues with modified request data.
+`GuardrailDecision::Block` short-circuits the provider call.
+
+## Current boundary
+
+These are Rust-only primitives. Python callback and guardrail adapters are a
+separate layer that should implement these Rust traits instead of changing the
+runner interfaces.

litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs

@@ -0,0 +1,468 @@
+//! Rust mirror of Python `CustomGuardrail` entrypoints used by the proxy.
+//!
+//! This module is intentionally Rust-only: Python/PyO3 adapters are a later
+//! layer that should implement this trait rather than changing the runner.
+
+use std::future::Future;
+use std::sync::Arc;
+
+use crate::integrations::custom_logger::{
+    CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
+};
+
+pub mod types;
+
+pub use types::{
+    GuardrailContext, GuardrailDecision, GuardrailDispatchReport, GuardrailError,
+    GuardrailEventHook, GuardrailFuture, GuardrailRequest,
+};
+
+pub trait CustomGuardrail: Send + Sync {
+    fn guardrail_name(&self) -> &str;
+
+    fn supported_event_hooks(&self) -> &[GuardrailEventHook];
+
+    /// Python 1:1 name: `async_pre_call_hook(user_api_key_dict, cache, data, call_type)`.
+    fn async_pre_call_hook<'a>(
+        &'a self,
+        _context: &'a GuardrailContext,
+        request: GuardrailRequest,
+    ) -> GuardrailFuture<'a> {
+        Box::pin(async move { Ok(GuardrailDecision::Allow(request)) })
+    }
+
+    /// Python 1:1 name: `async_moderation_hook(data, user_api_key_dict, call_type)`.
+    fn async_moderation_hook<'a>(
+        &'a self,
+        _context: &'a GuardrailContext,
+        request: GuardrailRequest,
+    ) -> GuardrailFuture<'a> {
+        Box::pin(async move { Ok(GuardrailDecision::Allow(request)) })
+    }
+}
+
+pub struct CustomGuardrailRunner {
+    guardrails: Vec<Arc<dyn CustomGuardrail>>,
+}
+
+impl CustomGuardrailRunner {
+    pub fn new(guardrails: Vec<Arc<dyn CustomGuardrail>>) -> Self {
+        Self { guardrails }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.guardrails.is_empty()
+    }
+
+    pub async fn run_pre_call(
+        &self,
+        context: &GuardrailContext,
+        request: GuardrailRequest,
+    ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
+        self.run_hook(GuardrailEventHook::PreCall, context, request)
+            .await
+    }
+
+    pub async fn run_during_call(
+        &self,
+        context: &GuardrailContext,
+        request: GuardrailRequest,
+    ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
+        self.run_hook(GuardrailEventHook::DuringCall, context, request)
+            .await
+    }
+
+    pub async fn run_before_provider<F, Fut, T>(
+        &self,
+        event_hook: GuardrailEventHook,
+        context: &GuardrailContext,
+        request: GuardrailRequest,
+        provider: F,
+    ) -> Result<T, GuardrailError>
+    where
+        F: FnOnce(GuardrailRequest) -> Fut,
+        Fut: Future<Output = Result<T, GuardrailError>>,
+    {
+        let (request, _) = self.run_hook(event_hook, context, request).await?;
+        provider(request).await
+    }
+
+    pub async fn run_pre_call_with_failure_logging(
+        &self,
+        context: &GuardrailContext,
+        request: GuardrailRequest,
+        logger_runner: &CustomLoggerRunner,
+        model_call_details: &ModelCallDetails,
+        timing: CallbackTiming,
+    ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
+        match self.run_pre_call(context, request).await {
+            Ok(result) => Ok(result),
+            Err(error) => {
+                let failure_details = model_call_details.clone().with_failure_error(LoggingError {
+                    message: error.message.clone(),
+                    kind: error.kind.clone(),
+                });
+                let response_obj = CallbackValue::new(
+                    "guardrail_error",
+                    serde_json::json!({
+                        "message": error.message,
+                        "kind": error.kind,
+                    }),
+                );
+                logger_runner
+                    .async_log_failure_event(&failure_details, Some(&response_obj), timing)
+                    .await;
+                Err(error)
+            }
+        }
+    }
+
+    async fn run_hook(
+        &self,
+        event_hook: GuardrailEventHook,
+        context: &GuardrailContext,
+        mut request: GuardrailRequest,
+    ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
+        if self.guardrails.is_empty() {
+            return Ok((request, GuardrailDispatchReport::default()));
+        }
+
+        let mut report = GuardrailDispatchReport::default();
+        for guardrail in &self.guardrails {
+            if !self.should_run(guardrail.as_ref(), event_hook, context) {
+                continue;
+            }
+
+            report.invoked += 1;
+            let decision = match event_hook {
+                GuardrailEventHook::PreCall => {
+                    guardrail
+                        .async_pre_call_hook(context, request.clone())
+                        .await?
+                }
+                GuardrailEventHook::DuringCall => {
+                    guardrail
+                        .async_moderation_hook(context, request.clone())
+                        .await?
+                }
+            };
+            match decision.into_request() {
+                Ok(next_request) => request = next_request,
+                Err(error) => return Err(error),
+            }
+        }
+
+        Ok((request, report))
+    }
+
+    fn should_run(
+        &self,
+        guardrail: &dyn CustomGuardrail,
+        event_hook: GuardrailEventHook,
+        context: &GuardrailContext,
+    ) -> bool {
+        let supports_hook = guardrail.supported_event_hooks().contains(&event_hook);
+        let selected = context.selected_guardrails.is_empty()
+            || context
+                .selected_guardrails
+                .iter()
+                .any(|name| name == guardrail.guardrail_name());
+        supports_hook && selected
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::integrations::custom_logger::{CallType, CallbackValue, CustomLogger, LogFuture};
+    use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload};
+    use serde_json::json;
+    use std::sync::Mutex;
+
+    #[derive(Clone)]
+    enum TestDecision {
+        Allow,
+        Mask,
+        Block,
+    }
+
+    struct RecordingCustomGuardrail {
+        name: String,
+        hooks: Vec<GuardrailEventHook>,
+        decision: TestDecision,
+        calls: Mutex<Vec<&'static str>>,
+    }
+
+    impl RecordingCustomGuardrail {
+        fn new(name: &str, hooks: Vec<GuardrailEventHook>, decision: TestDecision) -> Self {
+            Self {
+                name: name.to_string(),
+                hooks,
+                decision,
+                calls: Mutex::new(Vec::new()),
+            }
+        }
+
+        fn calls(&self) -> Vec<&'static str> {
+            self.calls.lock().unwrap().clone()
+        }
+
+        fn decision(&self, mut request: GuardrailRequest) -> GuardrailDecision {
+            match self.decision {
+                TestDecision::Allow => GuardrailDecision::Allow(request),
+                TestDecision::Mask => {
+                    request.data["masked"] = json!(true);
+                    GuardrailDecision::Mask(request)
+                }
+                TestDecision::Block => {
+                    GuardrailDecision::Block(GuardrailError::blocked("blocked by guardrail"))
+                }
+            }
+        }
+    }
+
+    impl CustomGuardrail for RecordingCustomGuardrail {
+        fn guardrail_name(&self) -> &str {
+            &self.name
+        }
+
+        fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
+            &self.hooks
+        }
+
+        fn async_pre_call_hook<'a>(
+            &'a self,
+            _context: &'a GuardrailContext,
+            request: GuardrailRequest,
+        ) -> GuardrailFuture<'a> {
+            Box::pin(async move {
+                self.calls.lock().unwrap().push("async_pre_call_hook");
+                Ok(self.decision(request))
+            })
+        }
+
+        fn async_moderation_hook<'a>(
+            &'a self,
+            _context: &'a GuardrailContext,
+            request: GuardrailRequest,
+        ) -> GuardrailFuture<'a> {
+            Box::pin(async move {
+                self.calls.lock().unwrap().push("async_moderation_hook");
+                Ok(self.decision(request))
+            })
+        }
+    }
+
+    #[tokio::test]
+    async fn pre_call_dispatches_to_async_pre_call_hook() {
+        let guardrail = Arc::new(RecordingCustomGuardrail::new(
+            "pre",
+            vec![GuardrailEventHook::PreCall],
+            TestDecision::Allow,
+        ));
+        let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]);
+        let context =
+            GuardrailContext::new(CallType::Ocr).with_selected_guardrails(vec!["pre".to_string()]);
+        let request = GuardrailRequest::new(json!({"messages": ["hello"]}));
+
+        let (result, report) = runner
+            .run_pre_call(&context, request)
+            .await
+            .expect("guardrail allows request");
+
+        assert_eq!(report.invoked, 1);
+        assert_eq!(result.data["messages"], json!(["hello"]));
+        assert_eq!(guardrail.calls(), vec!["async_pre_call_hook"]);
+    }
+
+    #[tokio::test]
+    async fn during_call_dispatches_to_async_moderation_hook() {
+        let guardrail = Arc::new(RecordingCustomGuardrail::new(
+            "during",
+            vec![GuardrailEventHook::DuringCall],
+            TestDecision::Allow,
+        ));
+        let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]);
+        let context = GuardrailContext::new(CallType::Completion)
+            .with_selected_guardrails(vec!["during".to_string()]);
+        let request = GuardrailRequest::new(json!({"prompt": "hello"}));
+
+        let (_result, report) = runner
+            .run_during_call(&context, request)
+            .await
+            .expect("guardrail allows request");
+
+        assert_eq!(report.invoked, 1);
+        assert_eq!(guardrail.calls(), vec!["async_moderation_hook"]);
+    }
+
+    #[tokio::test]
+    async fn mask_decision_continues_with_updated_request() {
+        let guardrail = Arc::new(RecordingCustomGuardrail::new(
+            "masker",
+            vec![GuardrailEventHook::PreCall],
+            TestDecision::Mask,
+        ));
+        let runner = CustomGuardrailRunner::new(vec![guardrail]);
+        let context = GuardrailContext::new(CallType::Ocr);
+        let request = GuardrailRequest::new(json!({"document": "secret"}));
+
+        let (result, report) = runner
+            .run_pre_call(&context, request)
+            .await
+            .expect("mask continues");
+
+        assert_eq!(report.invoked, 1);
+        assert_eq!(result.data["masked"], json!(true));
+    }
+
+    #[tokio::test]
+    async fn block_decision_short_circuits_and_logs_failure() {
+        struct RecordingFailureLogger {
+            errors: Mutex<Vec<String>>,
+        }
+
+        impl CustomLogger for RecordingFailureLogger {
+            fn async_log_failure_event<'a>(
+                &'a self,
+                model_call_details: &'a ModelCallDetails,
+                _response_obj: Option<&'a CallbackValue>,
+                _timing: CallbackTiming,
+            ) -> LogFuture<'a> {
+                Box::pin(async move {
+                    self.errors.lock().unwrap().push(
+                        model_call_details
+                            .failure_error
+                            .as_ref()
+                            .map(|error| error.kind.clone())
+                            .unwrap_or_default(),
+                    );
+                    Ok(())
+                })
+            }
+        }
+
+        let guardrail = Arc::new(RecordingCustomGuardrail::new(
+            "blocker",
+            vec![GuardrailEventHook::PreCall],
+            TestDecision::Block,
+        ));
+        let guardrail_runner = CustomGuardrailRunner::new(vec![guardrail]);
+        let logger = Arc::new(RecordingFailureLogger {
+            errors: Mutex::new(Vec::new()),
+        });
+        let logger_runner = CustomLoggerRunner::new(vec![logger.clone()]);
+        let context = GuardrailContext::new(CallType::Ocr);
+        let details = ModelCallDetails::from_standard_logging_payload(StandardLoggingPayload {
+            id: "req_ocr".to_string(),
+            litellm_call_id: "req_ocr".to_string(),
+            call_type: "ocr".to_string(),
+            model: "mistral-ocr-latest".to_string(),
+            custom_llm_provider: "mistral".to_string(),
+            response_cost: 0.0,
+            prompt_tokens: 0,
+            completion_tokens: 0,
+            total_tokens: 0,
+            start_time: 1.0,
+            end_time: 1.0,
+            stream: false,
+            metadata: StandardLoggingMetadata::default(),
+            messages: None,
+        });
+
+        let err = guardrail_runner
+            .run_pre_call_with_failure_logging(
+                &context,
+                GuardrailRequest::new(json!({"document": "bad"})),
+                &logger_runner,
+                &details,
+                CallbackTiming::new(1.0, 2.0),
+            )
+            .await
+            .expect_err("guardrail blocks request");
+
+        assert_eq!(err.kind, "GuardrailBlocked");
+        assert_eq!(
+            logger.errors.lock().unwrap().as_slice(),
+            ["GuardrailBlocked"]
+        );
+    }
+
+    #[tokio::test]
+    async fn block_decision_short_circuits_later_guardrails_and_provider_work() {
+        let blocking_guardrail = Arc::new(RecordingCustomGuardrail::new(
+            "blocker",
+            vec![GuardrailEventHook::PreCall],
+            TestDecision::Block,
+        ));
+        let later_guardrail = Arc::new(RecordingCustomGuardrail::new(
+            "later",
+            vec![GuardrailEventHook::PreCall],
+            TestDecision::Allow,
+        ));
+        let runner =
+            CustomGuardrailRunner::new(vec![blocking_guardrail.clone(), later_guardrail.clone()]);
+        let provider_called = Arc::new(Mutex::new(false));
+        let provider_called_for_closure = provider_called.clone();
+
+        let result = runner
+            .run_before_provider(
+                GuardrailEventHook::PreCall,
+                &GuardrailContext::new(CallType::Completion),
+                GuardrailRequest::new(json!({"prompt": "blocked"})),
+                move |_request| async move {
+                    *provider_called_for_closure.lock().unwrap() = true;
+                    Ok("provider response")
+                },
+            )
+            .await;
+
+        assert!(result.is_err());
+        assert_eq!(blocking_guardrail.calls(), vec!["async_pre_call_hook"]);
+        assert_eq!(later_guardrail.calls(), Vec::<&'static str>::new());
+        assert!(!*provider_called.lock().unwrap());
+    }
+
+    #[tokio::test]
+    async fn run_before_provider_returns_provider_guardrail_error_directly() {
+        let guardrail = Arc::new(RecordingCustomGuardrail::new(
+            "allow",
+            vec![GuardrailEventHook::PreCall],
+            TestDecision::Allow,
+        ));
+        let runner = CustomGuardrailRunner::new(vec![guardrail]);
+
+        let result = runner
+            .run_before_provider(
+                GuardrailEventHook::PreCall,
+                &GuardrailContext::new(CallType::Completion),
+                GuardrailRequest::new(json!({"prompt": "allowed"})),
+                |_request| async move {
+                    Err::<&'static str, GuardrailError>(GuardrailError::blocked(
+                        "provider-side guardrail error",
+                    ))
+                },
+            )
+            .await;
+
+        let err = result.expect_err("provider error is returned directly");
+        assert_eq!(err.kind, "GuardrailBlocked");
+        assert_eq!(err.message, "provider-side guardrail error");
+    }
+
+    #[tokio::test]
+    async fn no_guardrails_fast_path_dispatches_nothing() {
+        let runner = CustomGuardrailRunner::new(Vec::new());
+        let context = GuardrailContext::new(CallType::Ocr);
+        let request = GuardrailRequest::new(json!({"document": "ok"}));
+
+        let (result, report) = runner
+            .run_pre_call(&context, request)
+            .await
+            .expect("no guardrails allow request");
+
+        assert!(runner.is_empty());
+        assert_eq!(report, GuardrailDispatchReport::default());
+        assert_eq!(result.data["document"], json!("ok"));
+    }
+}

litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs

@@ -0,0 +1,110 @@
+use std::collections::HashMap;
+use std::future::Future;
+use std::pin::Pin;
+
+use serde_json::Value;
+
+use crate::integrations::custom_logger::CallType;
+
+pub type GuardrailFuture<'a> =
+    Pin<Box<dyn Future<Output = Result<GuardrailDecision, GuardrailError>> + Send + 'a>>;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum GuardrailEventHook {
+    PreCall,
+    DuringCall,
+}
+
+impl GuardrailEventHook {
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            Self::PreCall => "pre_call",
+            Self::DuringCall => "during_call",
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct GuardrailError {
+    pub message: String,
+    pub kind: String,
+}
+
+impl GuardrailError {
+    pub fn blocked(message: impl Into<String>) -> Self {
+        Self {
+            message: message.into(),
+            kind: "GuardrailBlocked".to_string(),
+        }
+    }
+}
+
+impl std::fmt::Display for GuardrailError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}: {}", self.kind, self.message)
+    }
+}
+
+impl std::error::Error for GuardrailError {}
+
+#[derive(Clone, Debug)]
+pub struct GuardrailContext {
+    pub call_type: CallType,
+    pub selected_guardrails: Vec<String>,
+    pub metadata: HashMap<String, Value>,
+    pub user_api_key_hash: Option<String>,
+    pub user_api_key_user_id: Option<String>,
+    pub user_api_key_team_id: Option<String>,
+    pub trace_parent: Option<String>,
+}
+
+impl GuardrailContext {
+    pub fn new(call_type: CallType) -> Self {
+        Self {
+            call_type,
+            selected_guardrails: Vec::new(),
+            metadata: HashMap::new(),
+            user_api_key_hash: None,
+            user_api_key_user_id: None,
+            user_api_key_team_id: None,
+            trace_parent: None,
+        }
+    }
+
+    pub fn with_selected_guardrails(mut self, selected_guardrails: Vec<String>) -> Self {
+        self.selected_guardrails = selected_guardrails;
+        self
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct GuardrailRequest {
+    pub data: Value,
+}
+
+impl GuardrailRequest {
+    pub fn new(data: Value) -> Self {
+        Self { data }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum GuardrailDecision {
+    Allow(GuardrailRequest),
+    Mask(GuardrailRequest),
+    Block(GuardrailError),
+}
+
+impl GuardrailDecision {
+    pub(super) fn into_request(self) -> Result<GuardrailRequest, GuardrailError> {
+        match self {
+            Self::Allow(request) | Self::Mask(request) => Ok(request),
+            Self::Block(error) => Err(error),
+        }
+    }
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct GuardrailDispatchReport {
+    pub invoked: usize,
+}

litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs

@@ -1,24 +0,0 @@
-//! The `CustomLogger` trait — the Rust mirror of Python
-//! `litellm/integrations/custom_logger.py::CustomLogger`.
-//!
-//! Synchronous (no `async_trait`): callbacks are O(1) enqueue-and-return so the
-//! realtime splice never blocks on a logger. Default bodies are no-ops so a
-//! logger can implement only the events it cares about.
-
-use crate::integrations::types::{LogError, LoggingError, StandardLoggingPayload};
-
-pub trait CustomLogger: Send + Sync {
-    /// Record a successful call. Default: no-op.
-    fn log_success_event(&self, _payload: &StandardLoggingPayload) -> Result<(), LogError> {
-        Ok(())
-    }
-
-    /// Record a failed call. Default: no-op.
-    fn log_failure_event(
-        &self,
-        _payload: &StandardLoggingPayload,
-        _error: &LoggingError,
-    ) -> Result<(), LogError> {
-        Ok(())
-    }
-}

litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs

@@ -0,0 +1,317 @@
+//! The `CustomLogger` trait — the Rust mirror of Python
+//! `litellm/integrations/custom_logger.py::CustomLogger`.
+//!
+//! The Python-named async terminal methods are the public Rust callback shape.
+
+use std::sync::Arc;
+
+pub mod types;
+
+pub use types::{
+    CallType, CallbackDispatchReport, CallbackTiming, CallbackValue, LogError, LogFuture,
+    LoggingError, ModelCallDetails,
+};
+
+pub trait CustomLogger: Send + Sync {
+    /// Python 1:1 name: `async_log_success_event(model_call_details, response_obj, start_time, end_time)`.
+    fn async_log_success_event<'a>(
+        &'a self,
+        _model_call_details: &'a ModelCallDetails,
+        _response_obj: &'a CallbackValue,
+        _timing: CallbackTiming,
+    ) -> LogFuture<'a> {
+        Box::pin(async { Ok(()) })
+    }
+
+    /// Python 1:1 name: `async_log_failure_event(model_call_details, response_obj, start_time, end_time)`.
+    fn async_log_failure_event<'a>(
+        &'a self,
+        _model_call_details: &'a ModelCallDetails,
+        _response_obj: Option<&'a CallbackValue>,
+        _timing: CallbackTiming,
+    ) -> LogFuture<'a> {
+        Box::pin(async { Ok(()) })
+    }
+}
+
+pub struct CustomLoggerRunner {
+    loggers: Vec<Arc<dyn CustomLogger>>,
+}
+
+impl CustomLoggerRunner {
+    pub fn new(loggers: Vec<Arc<dyn CustomLogger>>) -> Self {
+        Self { loggers }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.loggers.is_empty()
+    }
+
+    pub async fn async_log_success_event(
+        &self,
+        model_call_details: &ModelCallDetails,
+        response_obj: &CallbackValue,
+        timing: CallbackTiming,
+    ) -> CallbackDispatchReport {
+        if self.loggers.is_empty() {
+            return CallbackDispatchReport::default();
+        }
+
+        let mut report = CallbackDispatchReport::default();
+        for logger in &self.loggers {
+            report.invoked += 1;
+            if let Err(err) = logger
+                .async_log_success_event(model_call_details, response_obj, timing)
+                .await
+            {
+                report.dropped += 1;
+                eprintln!("litellm-ai-gateway: async_log_success_event dropped: {err}");
+            }
+        }
+        report
+    }
+
+    pub async fn async_log_failure_event(
+        &self,
+        model_call_details: &ModelCallDetails,
+        response_obj: Option<&CallbackValue>,
+        timing: CallbackTiming,
+    ) -> CallbackDispatchReport {
+        if self.loggers.is_empty() {
+            return CallbackDispatchReport::default();
+        }
+
+        let mut report = CallbackDispatchReport::default();
+        for logger in &self.loggers {
+            report.invoked += 1;
+            if let Err(err) = logger
+                .async_log_failure_event(model_call_details, response_obj, timing)
+                .await
+            {
+                report.dropped += 1;
+                eprintln!("litellm-ai-gateway: async_log_failure_event dropped: {err}");
+            }
+        }
+        report
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload};
+    use serde_json::json;
+    use std::sync::Mutex;
+
+    #[derive(Clone, Debug, PartialEq)]
+    struct RecordedEvent {
+        hook: &'static str,
+        model: String,
+        provider: String,
+        call_type: String,
+        request_id: Option<String>,
+        litellm_call_id: Option<String>,
+        user_id: Option<String>,
+        response_object: Option<String>,
+        error_kind: Option<String>,
+        start_time: f64,
+        end_time: f64,
+        standard_logging_model: Option<String>,
+    }
+
+    #[derive(Default)]
+    struct RecordingCustomLogger {
+        events: Mutex<Vec<RecordedEvent>>,
+    }
+
+    impl RecordingCustomLogger {
+        fn events(&self) -> Vec<RecordedEvent> {
+            self.events.lock().unwrap().clone()
+        }
+    }
+
+    impl CustomLogger for RecordingCustomLogger {
+        fn async_log_success_event<'a>(
+            &'a self,
+            model_call_details: &'a ModelCallDetails,
+            response_obj: &'a CallbackValue,
+            timing: CallbackTiming,
+        ) -> LogFuture<'a> {
+            Box::pin(async move {
+                self.events.lock().unwrap().push(RecordedEvent {
+                    hook: "async_log_success_event",
+                    model: model_call_details.model.clone(),
+                    provider: model_call_details.custom_llm_provider.clone(),
+                    call_type: model_call_details.call_type.to_string(),
+                    request_id: model_call_details.request_id.clone(),
+                    litellm_call_id: model_call_details.litellm_call_id.clone(),
+                    user_id: model_call_details.metadata.user_api_key_user_id.clone(),
+                    response_object: Some(response_obj.object.clone()),
+                    error_kind: None,
+                    start_time: timing.start_time,
+                    end_time: timing.end_time,
+                    standard_logging_model: model_call_details
+                        .standard_logging_payload
+                        .as_ref()
+                        .map(|payload| payload.model.clone()),
+                });
+                Ok(())
+            })
+        }
+
+        fn async_log_failure_event<'a>(
+            &'a self,
+            model_call_details: &'a ModelCallDetails,
+            response_obj: Option<&'a CallbackValue>,
+            timing: CallbackTiming,
+        ) -> LogFuture<'a> {
+            Box::pin(async move {
+                self.events.lock().unwrap().push(RecordedEvent {
+                    hook: "async_log_failure_event",
+                    model: model_call_details.model.clone(),
+                    provider: model_call_details.custom_llm_provider.clone(),
+                    call_type: model_call_details.call_type.to_string(),
+                    request_id: model_call_details.request_id.clone(),
+                    litellm_call_id: model_call_details.litellm_call_id.clone(),
+                    user_id: model_call_details.metadata.user_api_key_user_id.clone(),
+                    response_object: response_obj.map(|value| value.object.clone()),
+                    error_kind: model_call_details
+                        .failure_error
+                        .as_ref()
+                        .map(|error| error.kind.clone()),
+                    start_time: timing.start_time,
+                    end_time: timing.end_time,
+                    standard_logging_model: model_call_details
+                        .standard_logging_payload
+                        .as_ref()
+                        .map(|payload| payload.model.clone()),
+                });
+                Ok(())
+            })
+        }
+    }
+
+    fn payload(call_type: &str, model: &str, provider: &str) -> StandardLoggingPayload {
+        StandardLoggingPayload {
+            id: format!("req_{call_type}"),
+            litellm_call_id: format!("call_{call_type}"),
+            call_type: call_type.to_string(),
+            model: model.to_string(),
+            custom_llm_provider: provider.to_string(),
+            response_cost: 0.25,
+            prompt_tokens: 3,
+            completion_tokens: 4,
+            total_tokens: 7,
+            start_time: 10.0,
+            end_time: 11.5,
+            stream: false,
+            metadata: StandardLoggingMetadata {
+                user_api_key_hash: Some("hash".to_string()),
+                user_api_key_user_id: Some("user".to_string()),
+                user_api_key_team_id: Some("team".to_string()),
+                ..Default::default()
+            },
+            messages: Some(json!([{"role": "user", "content": "read this"}])),
+        }
+    }
+
+    #[tokio::test]
+    async fn rust_custom_logger_reads_success_payload_for_ocr() {
+        let logger = Arc::new(RecordingCustomLogger::default());
+        let runner = CustomLoggerRunner::new(vec![logger.clone()]);
+        let details = ModelCallDetails::from_standard_logging_payload(payload(
+            "ocr",
+            "mistral-ocr-latest",
+            "mistral",
+        ));
+        let response = CallbackValue::new("ocr", json!({"pages": [{"markdown": "ok"}]}));
+        let report = runner
+            .async_log_success_event(&details, &response, CallbackTiming::new(10.0, 11.5))
+            .await;
+
+        assert_eq!(report.invoked, 1);
+        assert_eq!(report.dropped, 0);
+        assert_eq!(
+            logger.events(),
+            vec![RecordedEvent {
+                hook: "async_log_success_event",
+                model: "mistral-ocr-latest".to_string(),
+                provider: "mistral".to_string(),
+                call_type: "ocr".to_string(),
+                request_id: Some("req_ocr".to_string()),
+                litellm_call_id: Some("call_ocr".to_string()),
+                user_id: Some("user".to_string()),
+                response_object: Some("ocr".to_string()),
+                error_kind: None,
+                start_time: 10.0,
+                end_time: 11.5,
+                standard_logging_model: Some("mistral-ocr-latest".to_string()),
+            }]
+        );
+    }
+
+    #[tokio::test]
+    async fn rust_custom_logger_reads_failure_payload_for_non_ocr_call_type() {
+        let logger = Arc::new(RecordingCustomLogger::default());
+        let runner = CustomLoggerRunner::new(vec![logger.clone()]);
+        let details = ModelCallDetails::from_standard_logging_payload(payload(
+            "acompletion",
+            "gpt-4.1-mini",
+            "openai",
+        ))
+        .with_failure_error(LoggingError {
+            message: "provider failed".to_string(),
+            kind: "ProviderError".to_string(),
+        });
+        let response = CallbackValue::new("error", json!({"message": "provider failed"}));
+        let report = runner
+            .async_log_failure_event(&details, Some(&response), CallbackTiming::new(2.0, 3.0))
+            .await;
+
+        assert_eq!(report.invoked, 1);
+        assert_eq!(report.dropped, 0);
+        assert_eq!(
+            logger.events(),
+            vec![RecordedEvent {
+                hook: "async_log_failure_event",
+                model: "gpt-4.1-mini".to_string(),
+                provider: "openai".to_string(),
+                call_type: "acompletion".to_string(),
+                request_id: Some("req_acompletion".to_string()),
+                litellm_call_id: Some("call_acompletion".to_string()),
+                user_id: Some("user".to_string()),
+                response_object: Some("error".to_string()),
+                error_kind: Some("ProviderError".to_string()),
+                start_time: 2.0,
+                end_time: 3.0,
+                standard_logging_model: Some("gpt-4.1-mini".to_string()),
+            }]
+        );
+    }
+
+    #[tokio::test]
+    async fn no_callback_fast_path_dispatches_nothing() {
+        let runner = CustomLoggerRunner::new(Vec::new());
+        let details = ModelCallDetails::new("mistral-ocr-latest", "mistral", CallType::Ocr);
+        let response = CallbackValue::new("ocr", json!({}));
+
+        let report = runner
+            .async_log_success_event(&details, &response, CallbackTiming::new(1.0, 1.5))
+            .await;
+
+        assert!(runner.is_empty());
+        assert_eq!(report, CallbackDispatchReport::default());
+    }
+
+    #[test]
+    fn with_standard_logging_payload_keeps_top_level_fields_in_sync() {
+        let details = ModelCallDetails::new("old-model", "old-provider", CallType::Completion)
+            .with_standard_logging_payload(payload("ocr", "mistral-ocr-latest", "mistral"));
+
+        assert_eq!(details.model, "mistral-ocr-latest");
+        assert_eq!(details.custom_llm_provider, "mistral");
+        assert_eq!(details.call_type, CallType::Ocr);
+        assert_eq!(details.request_id, Some("req_ocr".to_string()));
+        assert_eq!(details.litellm_call_id, Some("call_ocr".to_string()));
+    }
+}

litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs

@@ -0,0 +1,194 @@
+use std::collections::HashMap;
+use std::future::Future;
+use std::pin::Pin;
+
+use serde_json::Value;
+
+use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload};
+
+pub type LogFuture<'a> = Pin<Box<dyn Future<Output = Result<(), LogError>> + Send + 'a>>;
+
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct CallbackDispatchReport {
+    pub invoked: usize,
+    pub dropped: usize,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum CallType {
+    Ocr,
+    Realtime,
+    Completion,
+    Acompletion,
+    ChatCompletion,
+    Other(String),
+}
+
+impl CallType {
+    pub fn as_str(&self) -> &str {
+        match self {
+            Self::Ocr => "ocr",
+            Self::Realtime => "realtime",
+            Self::Completion => "completion",
+            Self::Acompletion => "acompletion",
+            Self::ChatCompletion => "chat_completion",
+            Self::Other(value) => value.as_str(),
+        }
+    }
+}
+
+impl From<&str> for CallType {
+    fn from(value: &str) -> Self {
+        match value {
+            "ocr" => Self::Ocr,
+            "realtime" => Self::Realtime,
+            "completion" => Self::Completion,
+            "acompletion" => Self::Acompletion,
+            "chat_completion" => Self::ChatCompletion,
+            other => Self::Other(other.to_string()),
+        }
+    }
+}
+
+impl std::fmt::Display for CallType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.as_str())
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct CallbackTiming {
+    pub start_time: f64,
+    pub end_time: f64,
+}
+
+impl CallbackTiming {
+    pub fn new(start_time: f64, end_time: f64) -> Self {
+        Self {
+            start_time,
+            end_time,
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct CallbackValue {
+    pub object: String,
+    pub value: Value,
+}
+
+impl CallbackValue {
+    pub fn new(object: impl Into<String>, value: Value) -> Self {
+        Self {
+            object: object.into(),
+            value,
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct ModelCallDetails {
+    pub model: String,
+    pub custom_llm_provider: String,
+    pub call_type: CallType,
+    pub metadata: StandardLoggingMetadata,
+    pub extra_metadata: HashMap<String, Value>,
+    pub request_id: Option<String>,
+    pub litellm_call_id: Option<String>,
+    pub response_cost: Option<f64>,
+    pub standard_logging_payload: Option<StandardLoggingPayload>,
+    pub failure_error: Option<LoggingError>,
+}
+
+impl ModelCallDetails {
+    pub fn new(
+        model: impl Into<String>,
+        custom_llm_provider: impl Into<String>,
+        call_type: CallType,
+    ) -> Self {
+        Self {
+            model: model.into(),
+            custom_llm_provider: custom_llm_provider.into(),
+            call_type,
+            metadata: StandardLoggingMetadata::default(),
+            extra_metadata: HashMap::new(),
+            request_id: None,
+            litellm_call_id: None,
+            response_cost: None,
+            standard_logging_payload: None,
+            failure_error: None,
+        }
+    }
+
+    pub fn from_standard_logging_payload(payload: StandardLoggingPayload) -> Self {
+        let request_id = Some(payload.id.clone());
+        let litellm_call_id = Some(payload.litellm_call_id.clone());
+        let response_cost = Some(payload.response_cost);
+        let metadata = payload.metadata.clone();
+        Self {
+            model: payload.model.clone(),
+            custom_llm_provider: payload.custom_llm_provider.clone(),
+            call_type: CallType::from(payload.call_type.as_str()),
+            metadata,
+            extra_metadata: HashMap::new(),
+            request_id,
+            litellm_call_id,
+            response_cost,
+            standard_logging_payload: Some(payload),
+            failure_error: None,
+        }
+    }
+
+    pub fn with_standard_logging_payload(mut self, payload: StandardLoggingPayload) -> Self {
+        self.model = payload.model.clone();
+        self.custom_llm_provider = payload.custom_llm_provider.clone();
+        self.call_type = CallType::from(payload.call_type.as_str());
+        self.request_id = Some(payload.id.clone());
+        self.litellm_call_id = Some(payload.litellm_call_id.clone());
+        self.response_cost = Some(payload.response_cost);
+        self.metadata = payload.metadata.clone();
+        self.standard_logging_payload = Some(payload);
+        self
+    }
+
+    pub fn with_failure_error(mut self, error: LoggingError) -> Self {
+        self.failure_error = Some(error);
+        self
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct LoggingError {
+    pub message: String,
+    pub kind: String,
+}
+
+#[derive(Clone, Debug)]
+pub struct LogError {
+    pub message: String,
+    pub kind: String,
+}
+
+impl LogError {
+    pub fn channel_full() -> Self {
+        Self {
+            message: "logging channel is full; dropping record".to_string(),
+            kind: "ChannelFull".to_string(),
+        }
+    }
+
+    pub fn channel_closed() -> Self {
+        Self {
+            message: "logging channel is closed; worker has shut down".to_string(),
+            kind: "ChannelClosed".to_string(),
+        }
+    }
+}
+
+impl std::fmt::Display for LogError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}: {}", self.kind, self.message)
+    }
+}
+
+impl std::error::Error for LogError {}

litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs

@@ -1,7 +1,8 @@
 //! A `CustomLogger` that ships finished events to the LiteLLM Python proxy's
 //! `/v1/rust_control_plane/logs` endpoint.
 //!
-//! The callback path is non-blocking: `log_success_event` / `log_failure_event`
+//! The callback path is non-blocking: `async_log_success_event` /
+//! `async_log_failure_event`
 //! build a `LogRecord` and `try_send` it onto a bounded channel, returning a
 //! `LogError` (never panicking, never awaiting) if the channel is full or the
 //! worker has gone away. A spawned background worker drains the channel, batches
@@ -15,54 +16,14 @@ use reqwest::Client;
 use tokio::sync::mpsc::{self, Receiver, Sender};
 use tokio::time::interval;
 
-use crate::constants::{
-    DEFAULT_CHANNEL_CAPACITY, DEFAULT_FLUSH_INTERVAL_MS, DEFAULT_MAX_BATCH_SIZE,
-    DEFAULT_PROXY_BASE_URL, RUST_CONTROL_PLANE_LOGS_PATH,
-};
-use crate::integrations::custom_logger::CustomLogger;
-use crate::integrations::types::{
-    CallbackLogsRequest, LogError, LogRecord, LoggingError, StandardLoggingPayload,
+use crate::constants::{DEFAULT_PROXY_BASE_URL, RUST_CONTROL_PLANE_LOGS_PATH};
+use crate::integrations::custom_logger::{
+    CallbackTiming, CallbackValue, CustomLogger, LogError, LogFuture, LoggingError,
+    ModelCallDetails,
 };
+use types::{CallbackLogsRequest, EgressTunables, LogRecord};
 
-/// Egress worker tunables. Each field defaults to the matching `DEFAULT_*` const
-/// in `crate::constants` and is overridable via an env var (read once at logger
-/// construction).
-struct EgressTunables {
-    channel_capacity: usize,
-    max_batch_size: usize,
-    flush_interval: Duration,
-}
-
-impl EgressTunables {
-    fn from_env() -> Self {
-        Self {
-            channel_capacity: env_positive(
-                "LITELLM_LOG_CHANNEL_CAPACITY",
-                DEFAULT_CHANNEL_CAPACITY,
-            ),
-            max_batch_size: env_positive("LITELLM_LOG_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE),
-            flush_interval: Duration::from_millis(env_positive(
-                "LITELLM_LOG_FLUSH_INTERVAL_MS",
-                DEFAULT_FLUSH_INTERVAL_MS,
-            )),
-        }
-    }
-}
-
-/// Parse a positive integer env var, falling back to `default` on missing,
-/// unparseable, or non-positive values. Generic over the integer type so one
-/// helper serves both the `usize` capacities and the `u64` interval.
-fn env_positive<T>(name: &str, default: T) -> T
-where
-    T: std::str::FromStr + PartialOrd + From<u8>,
-{
-    let zero = T::from(0u8);
-    std::env::var(name)
-        .ok()
-        .and_then(|value| value.trim().parse::<T>().ok())
-        .filter(|n| *n > zero)
-        .unwrap_or(default)
-}
+pub mod types;
 
 /// Ships realtime logging events to the LiteLLM Python proxy.
 pub struct LiteLLMPythonProxyAPILogger {
@@ -118,23 +79,50 @@ impl LiteLLMPythonProxyAPILogger {
 }
 
 impl CustomLogger for LiteLLMPythonProxyAPILogger {
-    fn log_success_event(&self, payload: &StandardLoggingPayload) -> Result<(), LogError> {
-        self.enqueue(LogRecord {
-            status: "success".to_string(),
-            payload: payload.clone(),
-            error: None,
+    fn async_log_success_event<'a>(
+        &'a self,
+        model_call_details: &'a ModelCallDetails,
+        _response_obj: &'a CallbackValue,
+        _timing: CallbackTiming,
+    ) -> LogFuture<'a> {
+        Box::pin(async move {
+            if let Some(payload) = &model_call_details.standard_logging_payload {
+                self.enqueue(LogRecord {
+                    status: "success".to_string(),
+                    payload: payload.clone(),
+                    error: None,
+                })?;
+            }
+            Ok(())
         })
     }
 
-    fn log_failure_event(
-        &self,
-        payload: &StandardLoggingPayload,
-        error: &LoggingError,
-    ) -> Result<(), LogError> {
-        self.enqueue(LogRecord {
-            status: "failure".to_string(),
-            payload: payload.clone(),
-            error: Some(format!("{}: {}", error.kind, error.message)),
+    fn async_log_failure_event<'a>(
+        &'a self,
+        model_call_details: &'a ModelCallDetails,
+        _response_obj: Option<&'a CallbackValue>,
+        _timing: CallbackTiming,
+    ) -> LogFuture<'a> {
+        Box::pin(async move {
+            if let Some(payload) = &model_call_details.standard_logging_payload {
+                let fallback_error;
+                let error = match &model_call_details.failure_error {
+                    Some(error) => error,
+                    None => {
+                        fallback_error = LoggingError {
+                            message: "callback failure event".to_string(),
+                            kind: "CallbackFailure".to_string(),
+                        };
+                        &fallback_error
+                    }
+                };
+                self.enqueue(LogRecord {
+                    status: "failure".to_string(),
+                    payload: payload.clone(),
+                    error: Some(format!("{}: {}", error.kind, error.message)),
+                })?;
+            }
+            Ok(())
         })
     }
 }

litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs

@@ -0,0 +1,72 @@
+use std::time::Duration;
+
+use serde::Serialize;
+
+use crate::constants::{
+    DEFAULT_CHANNEL_CAPACITY, DEFAULT_FLUSH_INTERVAL_MS, DEFAULT_MAX_BATCH_SIZE,
+};
+use crate::integrations::types::StandardLoggingPayload;
+
+#[derive(Serialize)]
+pub struct CallbackLogsRequest {
+    pub records: Vec<CallbackLogRecord>,
+}
+
+#[derive(Serialize)]
+pub struct CallbackLogRecord {
+    pub status: String,
+    pub standard_logging_payload: StandardLoggingPayload,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub error: Option<String>,
+}
+
+#[derive(Clone, Debug)]
+pub struct LogRecord {
+    pub status: String,
+    pub payload: StandardLoggingPayload,
+    pub error: Option<String>,
+}
+
+impl LogRecord {
+    pub fn into_callback_record(self) -> CallbackLogRecord {
+        CallbackLogRecord {
+            status: self.status,
+            standard_logging_payload: self.payload,
+            error: self.error,
+        }
+    }
+}
+
+pub(super) struct EgressTunables {
+    pub channel_capacity: usize,
+    pub max_batch_size: usize,
+    pub flush_interval: Duration,
+}
+
+impl EgressTunables {
+    pub fn from_env() -> Self {
+        Self {
+            channel_capacity: env_positive(
+                "LITELLM_LOG_CHANNEL_CAPACITY",
+                DEFAULT_CHANNEL_CAPACITY,
+            ),
+            max_batch_size: env_positive("LITELLM_LOG_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE),
+            flush_interval: Duration::from_millis(env_positive(
+                "LITELLM_LOG_FLUSH_INTERVAL_MS",
+                DEFAULT_FLUSH_INTERVAL_MS,
+            )),
+        }
+    }
+}
+
+fn env_positive<T>(name: &str, default: T) -> T
+where
+    T: std::str::FromStr + PartialOrd + From<u8>,
+{
+    let zero = T::from(0u8);
+    std::env::var(name)
+        .ok()
+        .and_then(|value| value.trim().parse::<T>().ok())
+        .filter(|n| *n > zero)
+        .unwrap_or(default)
+}

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

@@ -1,10 +1,12 @@
 //! Pure-Rust logging integrations. Names map 1:1 to Python
 //! `litellm/integrations/`:
+//!   - [`custom_guardrail::CustomGuardrail`] — the guardrail callback trait
 //!   - [`custom_logger::CustomLogger`]  — the callback trait
 //!   - [`litellm_python_proxy_api::LiteLLMPythonProxyAPILogger`] — ships events
-//!     to the Python proxy's `/v1/callbacks/logs` endpoint
+//!     to the Python proxy's `/v1/rust_control_plane/logs` endpoint
 //!   - [`types`] — the typed `StandardLoggingPayload` wire contract
 
+pub mod custom_guardrail;
 pub mod custom_logger;
 pub mod litellm_python_proxy_api;
 pub mod types;

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

@@ -28,68 +28,6 @@ pub struct RequestMetadata {
     pub user_api_key_team_id: Option<String>,
 }
 
-/// A logging-callback failure (e.g. a custom logger raised). Mirrors the Python
-/// failure-event shape: a message plus an exception kind/class name.
-#[derive(Clone, Debug)]
-pub struct LoggingError {
-    pub message: String,
-    pub kind: String,
-}
-
-/// A non-fatal error returned by a `CustomLogger` when it cannot enqueue an
-/// event (channel full or the background worker has shut down).
-#[derive(Clone, Debug)]
-pub struct LogError {
-    pub message: String,
-    pub kind: String,
-}
-
-impl LogError {
-    pub fn channel_full() -> Self {
-        Self {
-            message: "logging channel is full; dropping record".to_string(),
-            kind: "ChannelFull".to_string(),
-        }
-    }
-
-    pub fn channel_closed() -> Self {
-        Self {
-            message: "logging channel is closed; worker has shut down".to_string(),
-            kind: "ChannelClosed".to_string(),
-        }
-    }
-}
-
-impl std::fmt::Display for LogError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}: {}", self.kind, self.message)
-    }
-}
-
-impl std::error::Error for LogError {}
-
-/// Batch wrapper — the top-level request body.
-/// Matches Python `CallbackLogsRequest { records: list[CallbackLogRecord] }`.
-#[derive(Serialize)]
-pub struct CallbackLogsRequest {
-    pub records: Vec<CallbackLogRecord>,
-}
-
-/// One finished logging event.
-/// Matches `CallbackLogRecord { status, standard_logging_payload, error? }`.
-#[derive(Serialize)]
-pub struct CallbackLogRecord {
-    /// "success" | "failure". On "failure", `error` (or payload.error_str)
-    /// becomes the replayed exception string.
-    pub status: String,
-
-    pub standard_logging_payload: StandardLoggingPayload,
-
-    /// Only meaningful when status == "failure". Omitted on success.
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub error: Option<String>,
-}
-
 /// The self-describing payload. Field names are the EXACT JSON keys the Python
 /// replay path + spend-logs builder read.
 #[derive(Clone, Debug, Serialize)]
@@ -143,22 +81,3 @@ pub struct StandardLoggingMetadata {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub spend_logs_metadata: Option<HashMap<String, Value>>,
 }
-
-/// The unit handed to a `CustomLogger` sink: a finished payload plus its status
-/// and (on failure) the replayed error string.
-#[derive(Clone, Debug)]
-pub struct LogRecord {
-    pub status: String,
-    pub payload: StandardLoggingPayload,
-    pub error: Option<String>,
-}
-
-impl LogRecord {
-    pub fn into_callback_record(self) -> CallbackLogRecord {
-        CallbackLogRecord {
-            status: self.status,
-            standard_logging_payload: self.payload,
-            error: self.error,
-        }
-    }
-}

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

@@ -1,406 +1 @@
-//! End-to-end OCR orchestration.
-//!
-//! Owns supported OCR provider calls so the Python side stays a thin bridge:
-//! resolve the API key, build the URL + body via the pure transforms, POST it,
-//! and normalize the response. The HTTP client is built once and reused.
-
-use std::sync::OnceLock;
-use std::time::Duration;
-
-use litellm_core::error::CoreError;
-use litellm_core::ocr::transformation::{OcrAuthStrategy, OcrResponseHandling};
-use litellm_core::CoreResult;
-use serde_json::{Map, Value};
-
-mod common_utils;
-
-use common_utils::{
-    convert_document_url_to_data_uri, has_header, ocr_provider_config, poll_document_intelligence,
-    string_headers, truncate_error_body,
-};
-
-/// OCR over large documents can take a while; bound it generously rather than
-/// hanging forever on an unresponsive upstream. The client-level limit is the
-/// outer ceiling; callers can tighten it per request via ``run_ocr``'s ``timeout``.
-const OCR_TIMEOUT_SECS: u64 = 600;
-
-/// Process-wide async HTTP client (connection pool + TLS reused across calls).
-///
-/// The Python fallback path uses LiteLLM's standard `BaseLLMHTTPHandler`. This
-/// Rust path is opt-in and owns end-to-end OCR I/O, so it cannot call the
-/// Python handler directly; keep this route-scoped until litellm-rust has a
-/// shared HTTP abstraction.
-fn http_client() -> &'static reqwest::Client {
-    static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
-    CLIENT.get_or_init(|| {
-        reqwest::Client::builder()
-            .timeout(Duration::from_secs(OCR_TIMEOUT_SECS))
-            .build()
-            .expect("failed to build reqwest client")
-    })
-}
-
-fn upstream_headers(
-    headers: &[(String, String)],
-    auth_strategy: OcrAuthStrategy,
-    api_key: Option<&str>,
-) -> Vec<(String, String)> {
-    let auth_header = api_key.map(|api_key| match auth_strategy {
-        OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")),
-        OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()),
-    });
-    auth_header
-        .into_iter()
-        .chain(headers.iter().cloned())
-        .collect()
-}
-
-pub struct OcrRequest<'a> {
-    pub model: &'a str,
-    pub document: Value,
-    pub api_key: Option<&'a str>,
-    pub api_base: Option<&'a str>,
-    pub custom_llm_provider: &'a str,
-    pub extra_headers: Option<Map<String, Value>>,
-    pub optional_params: Map<String, Value>,
-    pub timeout: Option<Duration>,
-}
-
-/// Perform an OCR call end to end and return the normalized response as
-/// JSON (the shape the Python `OCRResponse` model expects).
-///
-/// Async: intended to be awaited directly by the Python bridge's async entrypoint.
-pub async fn ocr(request: OcrRequest<'_>) -> CoreResult<Value> {
-    let model = request.model;
-    let config = ocr_provider_config(request.custom_llm_provider, model)
-        .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.to_string()))?;
-    let env_lookup = |key: &str| std::env::var(key).ok();
-
-    let headers = string_headers(request.extra_headers)?;
-    let auth_strategy = config.auth_strategy();
-    let api_key = (!has_header(&headers, auth_strategy.header_name()))
-        .then(|| config.resolve_api_key(request.api_key, &env_lookup))
-        .transpose()?;
-    let url = config.complete_url(
-        request.api_base,
-        model,
-        &request.optional_params,
-        &env_lookup,
-    )?;
-    let filtered_params = config.map_ocr_params(&request.optional_params);
-    let document = if config.requires_data_uri_document() {
-        convert_document_url_to_data_uri(request.document).await?
-    } else {
-        request.document
-    };
-    let body = config
-        .transform_ocr_request(model, document, filtered_params)?
-        .data;
-    let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref());
-
-    let mut request_builder = http_client().post(&url).json(&body);
-    for (key, value) in &upstream_headers {
-        request_builder = request_builder.header(key, value);
-    }
-    if let Some(duration) = request.timeout {
-        request_builder = request_builder.timeout(duration);
-    }
-
-    let response = request_builder
-        .send()
-        .await
-        .map_err(|err| CoreError::Network(err.to_string()))?;
-
-    let status = response.status();
-    if config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
-        && status.as_u16() == 202
-    {
-        let operation_url = response
-            .headers()
-            .get("operation-location")
-            .and_then(|value| value.to_str().ok())
-            .map(str::to_string)
-            .ok_or_else(|| {
-                CoreError::InvalidResponse(
-                    "Azure Document Intelligence returned 202 but no Operation-Location header found"
-                        .to_string(),
-                )
-            })?;
-        let response_json =
-            poll_document_intelligence(&operation_url, &url, &upstream_headers, request.timeout)
-                .await?;
-        return Ok(config
-            .transform_ocr_response(model, response_json)?
-            .into_json());
-    }
-
-    let text = response
-        .text()
-        .await
-        .map_err(|err| CoreError::Network(err.to_string()))?;
-
-    if !status.is_success() {
-        return Err(CoreError::Http {
-            status: status.as_u16(),
-            body: truncate_error_body(&text),
-        });
-    }
-
-    let response_json: Value = serde_json::from_str(&text)
-        .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
-
-    Ok(config
-        .transform_ocr_response(model, response_json)?
-        .into_json())
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use serde_json::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();
-        let mut buffer = [0_u8; 1024];
-        loop {
-            let n = socket.read(&mut buffer).await.expect("reads request");
-            if n == 0 {
-                break;
-            }
-            request.extend_from_slice(&buffer[..n]);
-            if request.windows(4).any(|window| window == b"\r\n\r\n") {
-                break;
-            }
-        }
-        String::from_utf8(request).expect("request is utf8")
-    }
-
-    #[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 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"));
-    }
-
-    #[tokio::test]
-    async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() {
-        let listener = TcpListener::bind("127.0.0.1:0")
-            .await
-            .expect("test listener binds");
-        let addr = listener.local_addr().expect("listener has local addr");
-
-        let server = tokio::spawn(async move {
-            let (mut socket, _) = listener.accept().await.expect("accepts one request");
-            let request = read_http_headers(&mut socket).await;
-
-            let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
-            let response = format!(
-                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
-                response_body.len(),
-                response_body
-            );
-            socket
-                .write_all(response.as_bytes())
-                .await
-                .expect("writes response");
-            request
-        });
-
-        let mut headers = Map::new();
-        headers.insert(
-            "Authorization".to_string(),
-            Value::String("Bearer sk-from-python".to_string()),
-        );
-        headers.insert(
-            "x-trace-id".to_string(),
-            Value::String("trace-1".to_string()),
-        );
-
-        let response = ocr(OcrRequest {
-            model: "mistral-ocr-latest",
-            document: json!({
-                "type": "document_url",
-                "document_url": "https://example.com/doc.pdf"
-            }),
-            api_key: Some("sk-for-rust-fallback"),
-            api_base: Some(&format!("http://{addr}")),
-            custom_llm_provider: "mistral",
-            extra_headers: Some(headers),
-            optional_params: Map::new(),
-            timeout: Some(Duration::from_secs(5)),
-        })
-        .await
-        .expect("ocr request succeeds");
-
-        assert_eq!(response["pages"][0]["markdown"], "ok");
-
-        let request = server.await.expect("server task completes");
-        let authorization_count = request
-            .lines()
-            .filter(|line| line.to_ascii_lowercase().starts_with("authorization:"))
-            .count();
-        assert_eq!(authorization_count, 1, "{request}");
-        assert!(
-            request.contains("authorization: Bearer sk-from-python")
-                || request.contains("Authorization: Bearer sk-from-python"),
-            "{request}"
-        );
-    }
-
-    #[tokio::test]
-    async fn document_intelligence_poll_uses_resolved_subscription_key() {
-        let listener = TcpListener::bind("127.0.0.1:0")
-            .await
-            .expect("test listener binds");
-        let addr = listener.local_addr().expect("listener has local addr");
-        let operation_url = format!("http://{addr}/operations/1");
-
-        let server = tokio::spawn(async move {
-            let (mut post_socket, _) = listener.accept().await.expect("accepts post request");
-            let post_request = read_http_headers(&mut post_socket).await;
-            let post_response = format!(
-                "HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n"
-            );
-            post_socket
-                .write_all(post_response.as_bytes())
-                .await
-                .expect("writes post response");
-
-            let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request");
-            let poll_request = read_http_headers(&mut poll_socket).await;
-            let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#;
-            let poll_response = format!(
-                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
-                response_body.len(),
-                response_body
-            );
-            poll_socket
-                .write_all(poll_response.as_bytes())
-                .await
-                .expect("writes poll response");
-            (post_request, poll_request)
-        });
-
-        let response = ocr(OcrRequest {
-            model: "prebuilt-read",
-            document: json!({
-                "type": "document_url",
-                "document_url": "https://example.com/doc.pdf"
-            }),
-            api_key: Some("di-key"),
-            api_base: Some(&format!("http://{addr}")),
-            custom_llm_provider: "azure_ai/doc-intelligence",
-            extra_headers: None,
-            optional_params: Map::new(),
-            timeout: Some(Duration::from_secs(5)),
-        })
-        .await
-        .expect("document intelligence request succeeds");
-
-        assert_eq!(response["pages"][0]["markdown"], "ok");
-
-        let (post_request, poll_request) = server.await.expect("server task completes");
-        assert!(
-            post_request
-                .to_ascii_lowercase()
-                .contains("ocp-apim-subscription-key: di-key"),
-            "{post_request}"
-        );
-        assert!(
-            poll_request
-                .to_ascii_lowercase()
-                .contains("ocp-apim-subscription-key: di-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,
-            CoreError::InvalidRequest(
-                "OCR extra_headers.x-retry-count must be a string, got number".to_string()
-            )
-        );
-    }
-}
+pub use crate::ocr::{ocr, OcrRequest};

litellm-rust/crates/ai-gateway/src/io/realtime.rs

@@ -1,8 +1,8 @@
 //! End-to-end OpenAI realtime invocation.
 //!
-//! The host-facing entry point, mirroring `crate::io::ocr::run_ocr`: open the
-//! WebSocket to OpenAI, then splice a client realtime stream to the upstream,
-//! driving typed events through the pure `OPENAI_REALTIME_CONFIG` transforms.
+//! The host-facing entry point opens the WebSocket to OpenAI, then splices a
+//! client realtime stream to the upstream, driving typed events through the pure
+//! `OPENAI_REALTIME_CONFIG` transforms.
 //! Network, auth header, key resolution, and wire (de)serialization live here so
 //! the `transformation` module stays pure and typed.
 //!

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

@@ -3,15 +3,16 @@
 //! Two layers, split by feature so the Python `cdylib` can depend on the I/O
 //! without pulling in the HTTP server:
 //!
-//! - [`io`]: all network I/O (OCR HTTP call, realtime WebSocket splice, the
-//!   pre-warmed realtime pool). Always available — no feature required. The
-//!   Python bridge links this for `run_ocr`.
+//! - Call-type modules such as [`ocr`]: provider transforms, lifecycle hooks,
+//!   and provider I/O. Always available — no feature required.
+//! - [`io`]: compatibility exports and realtime WebSocket splice helpers.
 //! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling
 //!   `axum` are gated behind the `server` feature, which the `litellm-ai-gateway`
 //!   binary turns on. The `python-config` feature additionally pulls in [`python`]
 //!   for the load-time config reader.
 
 pub mod io;
+pub mod ocr;
 
 /// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and
 /// the `python-config` reader, so it is available without either feature.
@@ -27,9 +28,7 @@ pub mod state;
 // Realtime request logging. Only the server serves realtime, so these are
 // `server`-gated; `io::realtime` exposes the generic `observe` hook while the
 // collector and callback fan-out live here.
-#[cfg(feature = "server")]
 mod constants;
-#[cfg(feature = "server")]
 pub mod integrations;
 #[cfg(feature = "server")]
 mod realtime;

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

@@ -0,0 +1,14 @@
+use std::sync::OnceLock;
+use std::time::Duration;
+
+const OCR_TIMEOUT_SECS: u64 = 600;
+
+pub(super) fn http_client() -> &'static reqwest::Client {
+    static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
+    CLIENT.get_or_init(|| {
+        reqwest::Client::builder()
+            .timeout(Duration::from_secs(OCR_TIMEOUT_SECS))
+            .build()
+            .expect("failed to build reqwest client")
+    })
+}

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

@@ -18,7 +18,7 @@ use litellm_core::providers::vertex_ai::ocr::transformation::{
     VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
 };
 
-use super::http_client;
+use super::client::http_client;
 
 const ERROR_BODY_MAX_CHARS: usize = 256;
 const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
@@ -42,7 +42,6 @@ pub(super) fn ocr_provider_config(
         "azure_ai" if is_azure_document_intelligence_model(model) => {
             Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
         }
-        "azure_ai/doc-intelligence" => Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG),
         "azure_ai" => Some(&AZURE_AI_OCR_CONFIG),
         "vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG),
         "vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG),

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

@@ -0,0 +1,71 @@
+use litellm_core::error::CoreError;
+use litellm_core::ocr::transformation::OcrResponseHandling;
+use litellm_core::CoreResult;
+use serde_json::Value;
+
+use super::client::http_client;
+use super::common_utils::{poll_document_intelligence, truncate_error_body};
+use super::types::ProviderOcrRequest;
+
+pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult<Value> {
+    let mut request_builder = http_client().post(&request.url).json(&request.body);
+    for (key, value) in &request.upstream_headers {
+        request_builder = request_builder.header(key, value);
+    }
+    if let Some(duration) = request.timeout {
+        request_builder = request_builder.timeout(duration);
+    }
+
+    let response = request_builder
+        .send()
+        .await
+        .map_err(|err| CoreError::Network(err.to_string()))?;
+
+    let status = response.status();
+    if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
+        && status.as_u16() == 202
+    {
+        let operation_url = response
+            .headers()
+            .get("operation-location")
+            .and_then(|value| value.to_str().ok())
+            .map(str::to_string)
+            .ok_or_else(|| {
+                CoreError::InvalidResponse(
+                    "Azure Document Intelligence returned 202 but no Operation-Location header found"
+                        .to_string(),
+                )
+            })?;
+        let response_json = poll_document_intelligence(
+            &operation_url,
+            &request.url,
+            &request.upstream_headers,
+            request.timeout,
+        )
+        .await?;
+        return Ok(request
+            .config
+            .transform_ocr_response(&request.model, response_json)?
+            .into_json());
+    }
+
+    let text = response
+        .text()
+        .await
+        .map_err(|err| CoreError::Network(err.to_string()))?;
+
+    if !status.is_success() {
+        return Err(CoreError::Http {
+            status: status.as_u16(),
+            body: truncate_error_body(&text),
+        });
+    }
+
+    let response_json: Value = serde_json::from_str(&text)
+        .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
+
+    Ok(request
+        .config
+        .transform_ocr_response(&request.model, response_json)?
+        .into_json())
+}

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

@@ -0,0 +1,329 @@
+use std::future::Future;
+use std::pin::Pin;
+
+use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
+use litellm_core::error::CoreError;
+use litellm_core::ocr::transformation::OcrAuthStrategy;
+use litellm_core::CoreResult;
+use serde_json::{json, Map, Value};
+
+use super::common_utils::{
+    convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers,
+};
+use super::types::{PreparedOcrRequest, ProviderOcrRequest};
+use crate::integrations::custom_guardrail::{
+    CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
+};
+use crate::integrations::custom_logger::{
+    CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
+};
+use crate::integrations::types::{
+    RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload,
+};
+
+pub(crate) struct OcrLifecycleHooks {
+    logger_runner: CustomLoggerRunner,
+    guardrail_runner: CustomGuardrailRunner,
+    request_metadata: RequestMetadata,
+}
+
+type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
+type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
+
+impl OcrLifecycleHooks {
+    pub(crate) fn new(
+        logger_runner: CustomLoggerRunner,
+        guardrail_runner: CustomGuardrailRunner,
+        request_metadata: RequestMetadata,
+    ) -> Self {
+        Self {
+            logger_runner,
+            guardrail_runner,
+            request_metadata,
+        }
+    }
+
+    async fn run_pre_call_guardrails(
+        &self,
+        request: PreparedOcrRequest,
+    ) -> CoreResult<PreparedOcrRequest> {
+        if self.guardrail_runner.is_empty() {
+            return Ok(request);
+        }
+
+        let context = guardrail_context(&self.request_metadata);
+        let guardrail_request = GuardrailRequest::new(json!({
+            "model": request.model,
+            "custom_llm_provider": request.custom_llm_provider,
+            "document": request.document,
+            "optional_params": request.optional_params,
+        }));
+        let (guardrail_request, _) = self
+            .guardrail_runner
+            .run_pre_call(&context, guardrail_request)
+            .await
+            .map_err(guardrail_error_to_core_error)?;
+        let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?;
+        Ok(PreparedOcrRequest {
+            document,
+            optional_params,
+            ..request
+        })
+    }
+
+    async fn prepare_provider_request(
+        &self,
+        request: PreparedOcrRequest,
+    ) -> CoreResult<ProviderOcrRequest> {
+        let config = ocr_provider_config(&request.custom_llm_provider, &request.model)
+            .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?;
+        let env_lookup = |key: &str| std::env::var(key).ok();
+        let headers = string_headers(request.extra_headers)?;
+        let auth_strategy = config.auth_strategy();
+        let api_key = (!has_header(&headers, auth_strategy.header_name()))
+            .then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup))
+            .transpose()?;
+        let url = config.complete_url(
+            request.api_base.as_deref(),
+            &request.model,
+            &request.optional_params,
+            &env_lookup,
+        )?;
+        let filtered_params = config.map_ocr_params(&request.optional_params);
+        let model = request.model.clone();
+        let custom_llm_provider = request.custom_llm_provider.clone();
+        let document = if config.requires_data_uri_document() {
+            convert_document_url_to_data_uri(request.document).await?
+        } else {
+            request.document
+        };
+        let body = config
+            .transform_ocr_request(&request.model, document, filtered_params)?
+            .data;
+        let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref());
+        let body = self
+            .run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
+            .await?;
+        Ok(ProviderOcrRequest {
+            model,
+            config,
+            url,
+            body,
+            upstream_headers,
+            timeout: request.timeout,
+        })
+    }
+
+    async fn run_during_call_guardrails(
+        &self,
+        model: &str,
+        custom_llm_provider: &str,
+        url: &str,
+        body: Value,
+    ) -> CoreResult<Value> {
+        if self.guardrail_runner.is_empty() {
+            return Ok(body);
+        }
+
+        let context = guardrail_context(&self.request_metadata);
+        let guardrail_request = GuardrailRequest::new(json!({
+            "model": model,
+            "custom_llm_provider": custom_llm_provider,
+            "url": url,
+            "body": body,
+        }));
+        let (guardrail_request, _) = self
+            .guardrail_runner
+            .run_during_call(&context, guardrail_request)
+            .await
+            .map_err(guardrail_error_to_core_error)?;
+        parse_ocr_during_call_guardrail_request(guardrail_request)
+    }
+
+    fn standard_logging_payload(
+        &self,
+        context: &CallLifecycleContext,
+        timing: &CallLifecycleTiming,
+    ) -> StandardLoggingPayload {
+        StandardLoggingPayload {
+            id: context.litellm_call_id.clone(),
+            litellm_call_id: context.litellm_call_id.clone(),
+            call_type: context.call_type.clone(),
+            model: context.model.clone(),
+            custom_llm_provider: context.custom_llm_provider.clone(),
+            response_cost: 0.0,
+            prompt_tokens: 0,
+            completion_tokens: 0,
+            total_tokens: 0,
+            start_time: timing.start_time,
+            end_time: timing.end_time,
+            stream: false,
+            metadata: StandardLoggingMetadata {
+                user_api_key_hash: self.request_metadata.user_api_key_hash.clone(),
+                user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(),
+                user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(),
+                ..Default::default()
+            },
+            messages: None,
+        }
+    }
+}
+
+impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLifecycleHooks {
+    type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
+    type DuringCallFuture<'a> = OcrFuture<'a, ProviderOcrRequest>;
+    type SuccessFuture<'a> = OcrLogFuture<'a>;
+    type FailureFuture<'a> = OcrLogFuture<'a>;
+
+    fn async_pre_call_hook<'a>(
+        &'a self,
+        _context: &'a CallLifecycleContext,
+        request: PreparedOcrRequest,
+    ) -> Self::PreCallFuture<'a> {
+        Box::pin(async move { self.run_pre_call_guardrails(request).await })
+    }
+
+    fn async_during_call_hook<'a>(
+        &'a self,
+        _context: &'a CallLifecycleContext,
+        request: PreparedOcrRequest,
+    ) -> Self::DuringCallFuture<'a> {
+        Box::pin(async move { self.prepare_provider_request(request).await })
+    }
+
+    fn async_log_success_event<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        response: &'a Value,
+        timing: &'a CallLifecycleTiming,
+    ) -> Self::SuccessFuture<'a> {
+        Box::pin(async move {
+            if self.logger_runner.is_empty() {
+                return;
+            }
+            let response_obj = CallbackValue::new("ocr", response.clone());
+            self.logger_runner
+                .async_log_success_event(
+                    &ModelCallDetails::from_standard_logging_payload(
+                        self.standard_logging_payload(context, timing),
+                    ),
+                    &response_obj,
+                    CallbackTiming::new(timing.start_time, timing.end_time),
+                )
+                .await;
+        })
+    }
+
+    fn async_log_failure_event<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        error: &'a CoreError,
+        timing: &'a CallLifecycleTiming,
+    ) -> Self::FailureFuture<'a> {
+        Box::pin(async move {
+            if self.logger_runner.is_empty() {
+                return;
+            }
+            let logging_error = LoggingError {
+                message: error.to_string(),
+                kind: core_error_kind(error).to_string(),
+            };
+            let response_obj = CallbackValue::new(
+                "error",
+                json!({
+                    "message": logging_error.message,
+                    "kind": logging_error.kind,
+                }),
+            );
+            self.logger_runner
+                .async_log_failure_event(
+                    &ModelCallDetails::from_standard_logging_payload(
+                        self.standard_logging_payload(context, timing),
+                    )
+                    .with_failure_error(logging_error),
+                    Some(&response_obj),
+                    CallbackTiming::new(timing.start_time, timing.end_time),
+                )
+                .await;
+        })
+    }
+}
+
+fn upstream_headers(
+    headers: &[(String, String)],
+    auth_strategy: OcrAuthStrategy,
+    api_key: Option<&str>,
+) -> Vec<(String, String)> {
+    api_key
+        .map(|api_key| match auth_strategy {
+            OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")),
+            OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()),
+        })
+        .into_iter()
+        .chain(headers.iter().cloned())
+        .collect()
+}
+
+fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
+    GuardrailContext {
+        call_type: CallType::Ocr,
+        selected_guardrails: Vec::new(),
+        metadata: std::collections::HashMap::new(),
+        user_api_key_hash: metadata.user_api_key_hash.clone(),
+        user_api_key_user_id: metadata.user_api_key_user_id.clone(),
+        user_api_key_team_id: metadata.user_api_key_team_id.clone(),
+        trace_parent: None,
+    }
+}
+
+fn parse_ocr_pre_call_guardrail_request(
+    request: GuardrailRequest,
+) -> CoreResult<(Value, Map<String, Value>)> {
+    let Value::Object(mut data) = request.data else {
+        return Err(CoreError::InvalidRequest(
+            "OCR pre_call guardrail must return an object".to_string(),
+        ));
+    };
+    let document = data.remove("document").ok_or_else(|| {
+        CoreError::InvalidRequest("OCR pre_call guardrail removed document".to_string())
+    })?;
+    let optional_params = match data.remove("optional_params") {
+        Some(Value::Object(params)) => params,
+        Some(_) => {
+            return Err(CoreError::InvalidRequest(
+                "OCR pre_call guardrail optional_params must be an object".to_string(),
+            ))
+        }
+        None => Map::new(),
+    };
+    Ok((document, optional_params))
+}
+
+fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> CoreResult<Value> {
+    let Value::Object(mut data) = request.data else {
+        return Err(CoreError::InvalidRequest(
+            "OCR during_call guardrail must return an object".to_string(),
+        ));
+    };
+    data.remove("body").ok_or_else(|| {
+        CoreError::InvalidRequest("OCR during_call guardrail removed body".to_string())
+    })
+}
+
+fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError {
+    CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message))
+}
+
+fn core_error_kind(error: &CoreError) -> &'static str {
+    match error {
+        CoreError::Auth(_) => "AuthError",
+        CoreError::InvalidProvider(_) => "InvalidProvider",
+        CoreError::InvalidRequest(_) => "InvalidRequest",
+        CoreError::InvalidType { .. } => "InvalidType",
+        CoreError::MissingField(_) => "MissingField",
+        CoreError::Http { .. } => "HttpError",
+        CoreError::InvalidResponse(_) => "InvalidResponse",
+        CoreError::Network(_) => "NetworkError",
+        CoreError::Routing(_) => "RoutingError",
+    }
+}

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

@@ -0,0 +1,25 @@
+use litellm_core::call_lifecycle::CallLifecycle;
+use litellm_core::CoreResult;
+use serde_json::Value;
+
+mod client;
+mod common_utils;
+mod handler;
+mod hooks;
+mod prepare;
+mod types;
+
+pub use types::OcrRequest;
+
+use handler::execute_ocr_provider_call;
+use prepare::{prepare_ocr_call, PreparedOcrCall};
+
+pub async fn ocr(request: OcrRequest<'_>) -> CoreResult<Value> {
+    let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
+    CallLifecycle::default()
+        .run_request(request, &hooks, execute_ocr_provider_call)
+        .await
+}
+
+#[cfg(test)]
+mod tests;

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

@@ -0,0 +1,57 @@
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider};
+
+use super::hooks::OcrLifecycleHooks;
+use super::types::{OcrRequest, PreparedOcrRequest};
+use crate::integrations::custom_guardrail::CustomGuardrailRunner;
+use crate::integrations::custom_logger::CustomLoggerRunner;
+
+pub(crate) struct PreparedOcrCall {
+    pub(crate) request: PreparedOcrRequest,
+    pub(crate) hooks: OcrLifecycleHooks,
+}
+
+pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
+    let call_id = request
+        .litellm_call_id
+        .map(str::to_string)
+        .unwrap_or_else(new_ocr_call_id);
+    let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
+        .unwrap_or(CustomLlmProvider {
+            model: request.model,
+            custom_llm_provider: "mistral",
+        });
+    let model = provider_info.model.to_string();
+    let custom_llm_provider = provider_info.custom_llm_provider.to_string();
+
+    PreparedOcrCall {
+        request: PreparedOcrRequest {
+            model,
+            custom_llm_provider,
+            litellm_call_id: call_id,
+            document: request.document,
+            api_key: request.api_key.map(str::to_string),
+            api_base: request.api_base.map(str::to_string),
+            extra_headers: request.extra_headers,
+            optional_params: request.optional_params,
+            timeout: request.timeout,
+        },
+        hooks: OcrLifecycleHooks::new(
+            CustomLoggerRunner::new(request.callbacks),
+            CustomGuardrailRunner::new(request.guardrails),
+            request.request_metadata,
+        ),
+    }
+}
+
+fn new_ocr_call_id() -> String {
+    static COUNTER: AtomicU64 = AtomicU64::new(1);
+    let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
+    let timestamp = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|duration| duration.as_nanos())
+        .unwrap_or(0);
+    format!("ocr-{timestamp}-{sequence}")
+}

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

@@ -0,0 +1,610 @@
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+
+use litellm_core::error::CoreError;
+use litellm_core::ocr::transformation::OcrResponseHandling;
+use serde_json::{json, Map, Value};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::{TcpListener, TcpStream};
+
+use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body};
+use super::{ocr, OcrRequest};
+use crate::integrations::custom_guardrail::{
+    CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
+    GuardrailFuture, GuardrailRequest,
+};
+use crate::integrations::custom_logger::{
+    CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
+};
+use crate::integrations::types::RequestMetadata;
+
+async fn read_http_headers(socket: &mut TcpStream) -> String {
+    let mut request = Vec::new();
+    let mut buffer = [0_u8; 1024];
+    loop {
+        let n = socket.read(&mut buffer).await.expect("reads request");
+        if n == 0 {
+            break;
+        }
+        request.extend_from_slice(&buffer[..n]);
+        if request.windows(4).any(|window| window == b"\r\n\r\n") {
+            break;
+        }
+    }
+    String::from_utf8(request).expect("request is utf8")
+}
+
+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")
+}
+
+#[derive(Clone, Debug, PartialEq)]
+struct RecordedLogEvent {
+    hook: &'static str,
+    model: String,
+    call_type: String,
+    user_id: Option<String>,
+    response_object: Option<String>,
+    error_kind: Option<String>,
+}
+
+#[derive(Default)]
+struct RecordingOcrLogger {
+    events: Mutex<Vec<RecordedLogEvent>>,
+}
+
+impl RecordingOcrLogger {
+    fn events(&self) -> Vec<RecordedLogEvent> {
+        self.events.lock().unwrap().clone()
+    }
+}
+
+impl CustomLogger for RecordingOcrLogger {
+    fn async_log_success_event<'a>(
+        &'a self,
+        model_call_details: &'a ModelCallDetails,
+        response_obj: &'a CallbackValue,
+        _timing: CallbackTiming,
+    ) -> LogFuture<'a> {
+        Box::pin(async move {
+            self.events.lock().unwrap().push(RecordedLogEvent {
+                hook: "async_log_success_event",
+                model: model_call_details.model.clone(),
+                call_type: model_call_details.call_type.to_string(),
+                user_id: model_call_details.metadata.user_api_key_user_id.clone(),
+                response_object: Some(response_obj.object.clone()),
+                error_kind: None,
+            });
+            Ok(())
+        })
+    }
+
+    fn async_log_failure_event<'a>(
+        &'a self,
+        model_call_details: &'a ModelCallDetails,
+        response_obj: Option<&'a CallbackValue>,
+        _timing: CallbackTiming,
+    ) -> LogFuture<'a> {
+        Box::pin(async move {
+            self.events.lock().unwrap().push(RecordedLogEvent {
+                hook: "async_log_failure_event",
+                model: model_call_details.model.clone(),
+                call_type: model_call_details.call_type.to_string(),
+                user_id: model_call_details.metadata.user_api_key_user_id.clone(),
+                response_object: response_obj.map(|value| value.object.clone()),
+                error_kind: model_call_details
+                    .failure_error
+                    .as_ref()
+                    .map(|error| error.kind.clone()),
+            });
+            Ok(())
+        })
+    }
+}
+
+struct RecordingOcrGuardrail {
+    hooks: Vec<GuardrailEventHook>,
+    events: Mutex<Vec<&'static str>>,
+    block_pre_call: bool,
+}
+
+impl RecordingOcrGuardrail {
+    fn new(hooks: Vec<GuardrailEventHook>) -> Self {
+        Self {
+            hooks,
+            events: Mutex::new(Vec::new()),
+            block_pre_call: false,
+        }
+    }
+
+    fn blocking_pre_call() -> Self {
+        Self {
+            hooks: vec![GuardrailEventHook::PreCall],
+            events: Mutex::new(Vec::new()),
+            block_pre_call: true,
+        }
+    }
+
+    fn events(&self) -> Vec<&'static str> {
+        self.events.lock().unwrap().clone()
+    }
+}
+
+impl CustomGuardrail for RecordingOcrGuardrail {
+    fn guardrail_name(&self) -> &str {
+        "recording-ocr-guardrail"
+    }
+
+    fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
+        &self.hooks
+    }
+
+    fn async_pre_call_hook<'a>(
+        &'a self,
+        _context: &'a GuardrailContext,
+        mut request: GuardrailRequest,
+    ) -> GuardrailFuture<'a> {
+        Box::pin(async move {
+            self.events.lock().unwrap().push("async_pre_call_hook");
+            if self.block_pre_call {
+                return Ok(GuardrailDecision::Block(GuardrailError::blocked(
+                    "blocked before provider",
+                )));
+            }
+            request.data["document"]["guarded_pre"] = json!(true);
+            Ok(GuardrailDecision::Mask(request))
+        })
+    }
+
+    fn async_moderation_hook<'a>(
+        &'a self,
+        _context: &'a GuardrailContext,
+        mut request: GuardrailRequest,
+    ) -> GuardrailFuture<'a> {
+        Box::pin(async move {
+            self.events.lock().unwrap().push("async_moderation_hook");
+            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");
+}
+
+#[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 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"));
+}
+
+#[tokio::test]
+async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("test listener binds");
+    let addr = listener.local_addr().expect("listener has local addr");
+
+    let server = tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.expect("accepts one request");
+        let request = read_http_request(&mut socket).await;
+        let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
+        let response = format!(
+            "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+            response_body.len(),
+            response_body
+        );
+        socket
+            .write_all(response.as_bytes())
+            .await
+            .expect("writes response");
+        request
+    });
+
+    let logger = Arc::new(RecordingOcrLogger::default());
+    let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![
+        GuardrailEventHook::PreCall,
+        GuardrailEventHook::DuringCall,
+    ]));
+    let response = ocr(OcrRequest {
+        model: "mistral-ocr-latest",
+        document: json!({
+            "type": "document_url",
+            "document_url": "https://example.com/doc.pdf"
+        }),
+        api_key: Some("sk-test"),
+        api_base: Some(&format!("http://{addr}")),
+        custom_llm_provider: Some("mistral"),
+        extra_headers: None,
+        optional_params: Map::new(),
+        timeout: Some(Duration::from_secs(5)),
+        callbacks: vec![logger.clone()],
+        guardrails: vec![guardrail.clone()],
+        request_metadata: RequestMetadata {
+            user_api_key_user_id: Some("user-1".to_string()),
+            ..Default::default()
+        },
+        litellm_call_id: Some("ocr-call-1"),
+    })
+    .await
+    .expect("ocr request succeeds");
+
+    assert_eq!(response["pages"][0]["markdown"], "ok");
+    assert_eq!(
+        guardrail.events(),
+        vec!["async_pre_call_hook", "async_moderation_hook"]
+    );
+    assert_eq!(
+        logger.events(),
+        vec![RecordedLogEvent {
+            hook: "async_log_success_event",
+            model: "mistral-ocr-latest".to_string(),
+            call_type: "ocr".to_string(),
+            user_id: Some("user-1".to_string()),
+            response_object: Some("ocr".to_string()),
+            error_kind: None,
+        }]
+    );
+
+    let request = server.await.expect("server task completes");
+    assert!(request.contains(r#""guarded_pre":true"#), "{request}");
+    assert!(request.contains(r#""guarded_during":true"#), "{request}");
+}
+
+#[tokio::test]
+async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("test listener binds");
+    let addr = listener.local_addr().expect("listener has local addr");
+
+    let server = tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.expect("accepts one request");
+        let _request = read_http_request(&mut socket).await;
+        let response_body = "provider failed";
+        let response = format!(
+            "HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+            response_body.len(),
+            response_body
+        );
+        socket
+            .write_all(response.as_bytes())
+            .await
+            .expect("writes response");
+    });
+
+    let logger = Arc::new(RecordingOcrLogger::default());
+    let err = ocr(OcrRequest {
+        model: "mistral-ocr-latest",
+        document: json!({
+            "type": "document_url",
+            "document_url": "https://example.com/doc.pdf"
+        }),
+        api_key: Some("sk-test"),
+        api_base: Some(&format!("http://{addr}")),
+        custom_llm_provider: Some("mistral"),
+        extra_headers: None,
+        optional_params: Map::new(),
+        timeout: Some(Duration::from_secs(5)),
+        callbacks: vec![logger.clone()],
+        guardrails: Vec::new(),
+        request_metadata: RequestMetadata::default(),
+        litellm_call_id: Some("ocr-call-2"),
+    })
+    .await
+    .expect_err("provider error propagates");
+
+    assert!(matches!(err, CoreError::Http { status: 500, .. }));
+    server.await.expect("server task completes");
+    assert_eq!(
+        logger.events(),
+        vec![RecordedLogEvent {
+            hook: "async_log_failure_event",
+            model: "mistral-ocr-latest".to_string(),
+            call_type: "ocr".to_string(),
+            user_id: None,
+            response_object: Some("error".to_string()),
+            error_kind: Some("HttpError".to_string()),
+        }]
+    );
+}
+
+#[tokio::test]
+async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("test listener binds");
+    let addr = listener.local_addr().expect("listener has local addr");
+    let logger = Arc::new(RecordingOcrLogger::default());
+    let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call());
+
+    let err = ocr(OcrRequest {
+        model: "mistral-ocr-latest",
+        document: json!({
+            "type": "document_url",
+            "document_url": "https://example.com/doc.pdf"
+        }),
+        api_key: Some("sk-test"),
+        api_base: Some(&format!("http://{addr}")),
+        custom_llm_provider: Some("mistral"),
+        extra_headers: None,
+        optional_params: Map::new(),
+        timeout: Some(Duration::from_millis(100)),
+        callbacks: vec![logger.clone()],
+        guardrails: vec![guardrail.clone()],
+        request_metadata: RequestMetadata::default(),
+        litellm_call_id: Some("ocr-call-3"),
+    })
+    .await
+    .expect_err("guardrail blocks request");
+
+    assert!(matches!(err, CoreError::InvalidRequest(_)));
+    assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]);
+    assert_eq!(
+        logger.events(),
+        vec![RecordedLogEvent {
+            hook: "async_log_failure_event",
+            model: "mistral-ocr-latest".to_string(),
+            call_type: "ocr".to_string(),
+            user_id: None,
+            response_object: Some("error".to_string()),
+            error_kind: Some("InvalidRequest".to_string()),
+        }]
+    );
+    let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
+    assert!(accepted.is_err(), "provider socket should not be touched");
+}
+
+#[tokio::test]
+async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("test listener binds");
+    let addr = listener.local_addr().expect("listener has local addr");
+
+    let server = tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.expect("accepts one request");
+        let request = read_http_headers(&mut socket).await;
+        let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
+        let response = format!(
+            "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+            response_body.len(),
+            response_body
+        );
+        socket
+            .write_all(response.as_bytes())
+            .await
+            .expect("writes response");
+        request
+    });
+
+    let mut headers = Map::new();
+    headers.insert(
+        "Authorization".to_string(),
+        Value::String("Bearer sk-from-python".to_string()),
+    );
+    headers.insert(
+        "x-trace-id".to_string(),
+        Value::String("trace-1".to_string()),
+    );
+
+    let response = ocr(OcrRequest {
+        model: "mistral-ocr-latest",
+        document: json!({
+            "type": "document_url",
+            "document_url": "https://example.com/doc.pdf"
+        }),
+        api_key: Some("sk-for-rust-fallback"),
+        api_base: Some(&format!("http://{addr}")),
+        custom_llm_provider: Some("mistral"),
+        extra_headers: Some(headers),
+        optional_params: Map::new(),
+        timeout: Some(Duration::from_secs(5)),
+        callbacks: Vec::new(),
+        guardrails: Vec::new(),
+        request_metadata: RequestMetadata::default(),
+        litellm_call_id: None,
+    })
+    .await
+    .expect("ocr request succeeds");
+
+    assert_eq!(response["pages"][0]["markdown"], "ok");
+
+    let request = server.await.expect("server task completes");
+    let authorization_count = request
+        .lines()
+        .filter(|line| line.to_ascii_lowercase().starts_with("authorization:"))
+        .count();
+    assert_eq!(authorization_count, 1, "{request}");
+    assert!(
+        request.contains("authorization: Bearer sk-from-python")
+            || request.contains("Authorization: Bearer sk-from-python"),
+        "{request}"
+    );
+}
+
+#[tokio::test]
+async fn document_intelligence_poll_uses_resolved_subscription_key() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("test listener binds");
+    let addr = listener.local_addr().expect("listener has local addr");
+    let operation_url = format!("http://{addr}/operations/1");
+
+    let server = tokio::spawn(async move {
+        let (mut post_socket, _) = listener.accept().await.expect("accepts post request");
+        let post_request = read_http_headers(&mut post_socket).await;
+        let post_response = format!(
+            "HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n"
+        );
+        post_socket
+            .write_all(post_response.as_bytes())
+            .await
+            .expect("writes post response");
+
+        let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request");
+        let poll_request = read_http_headers(&mut poll_socket).await;
+        let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#;
+        let poll_response = format!(
+            "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+            response_body.len(),
+            response_body
+        );
+        poll_socket
+            .write_all(poll_response.as_bytes())
+            .await
+            .expect("writes poll response");
+        (post_request, poll_request)
+    });
+
+    let response = ocr(OcrRequest {
+        model: "doc-intelligence/prebuilt-read",
+        document: json!({
+            "type": "document_url",
+            "document_url": "https://example.com/doc.pdf"
+        }),
+        api_key: Some("di-key"),
+        api_base: Some(&format!("http://{addr}")),
+        custom_llm_provider: Some("azure_ai"),
+        extra_headers: None,
+        optional_params: Map::new(),
+        timeout: Some(Duration::from_secs(5)),
+        callbacks: Vec::new(),
+        guardrails: Vec::new(),
+        request_metadata: RequestMetadata::default(),
+        litellm_call_id: None,
+    })
+    .await
+    .expect("document intelligence request succeeds");
+
+    assert_eq!(response["pages"][0]["markdown"], "ok");
+
+    let (post_request, poll_request) = server.await.expect("server task completes");
+    assert!(
+        post_request
+            .to_ascii_lowercase()
+            .contains("ocp-apim-subscription-key: di-key"),
+        "{post_request}"
+    );
+    assert!(
+        poll_request
+            .to_ascii_lowercase()
+            .contains("ocp-apim-subscription-key: di-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,
+        CoreError::InvalidRequest(
+            "OCR extra_headers.x-retry-count must be a string, got number".to_string()
+        )
+    );
+}

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

@@ -0,0 +1,57 @@
+use std::sync::Arc;
+use std::time::Duration;
+
+use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
+use litellm_core::ocr::transformation::OcrProviderConfig;
+use serde_json::{Map, Value};
+
+use crate::integrations::custom_guardrail::CustomGuardrail;
+use crate::integrations::custom_logger::CustomLogger;
+use crate::integrations::types::RequestMetadata;
+
+pub struct OcrRequest<'a> {
+    pub model: &'a str,
+    pub document: Value,
+    pub api_key: Option<&'a str>,
+    pub api_base: Option<&'a str>,
+    pub custom_llm_provider: Option<&'a str>,
+    pub extra_headers: Option<Map<String, Value>>,
+    pub optional_params: Map<String, Value>,
+    pub timeout: Option<Duration>,
+    pub callbacks: Vec<Arc<dyn CustomLogger>>,
+    pub guardrails: Vec<Arc<dyn CustomGuardrail>>,
+    pub request_metadata: RequestMetadata,
+    pub litellm_call_id: Option<&'a str>,
+}
+
+pub(crate) struct PreparedOcrRequest {
+    pub(crate) model: String,
+    pub(crate) custom_llm_provider: String,
+    pub(crate) litellm_call_id: String,
+    pub(crate) document: Value,
+    pub(crate) api_key: Option<String>,
+    pub(crate) api_base: Option<String>,
+    pub(crate) extra_headers: Option<Map<String, Value>>,
+    pub(crate) optional_params: Map<String, Value>,
+    pub(crate) timeout: Option<Duration>,
+}
+
+impl CallLifecycleRequest for PreparedOcrRequest {
+    fn lifecycle_context(&self) -> CallLifecycleContext {
+        CallLifecycleContext::new(
+            "ocr",
+            self.model.clone(),
+            self.custom_llm_provider.clone(),
+            self.litellm_call_id.clone(),
+        )
+    }
+}
+
+pub(crate) struct ProviderOcrRequest {
+    pub(crate) model: String,
+    pub(crate) config: &'static dyn OcrProviderConfig,
+    pub(crate) url: String,
+    pub(crate) body: Value,
+    pub(crate) upstream_headers: Vec<(String, String)>,
+    pub(crate) timeout: Option<Duration>,
+}

litellm-rust/crates/ai-gateway/src/realtime/streaming.rs

@@ -13,7 +13,9 @@ use litellm_core::realtime::types::RealtimeEvent;
 use serde_json::Value;
 
 use crate::constants::DEFAULT_PROVIDER;
-use crate::integrations::custom_logger::CustomLogger;
+use crate::integrations::custom_logger::{
+    CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails,
+};
 use crate::integrations::types::{
     RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, Usage,
 };
@@ -183,30 +185,45 @@ impl RealTimeStreaming {
     /// Finish the session: stamp the end time and fan the payload out to every
     /// callback. On a logger enqueue error we bump a non-fatal counter (the
     /// realtime session has already ended; a dropped log must never propagate).
-    pub fn log_messages(&mut self, status: SessionStatus) {
+    pub async fn log_messages(&mut self, status: SessionStatus) {
         self.end_time = epoch_seconds();
         let payload = self.build_payload();
+        let timing = CallbackTiming::new(payload.start_time, payload.end_time);
+        let runner = CustomLoggerRunner::new(self.callbacks.clone());
 
         match status {
             SessionStatus::Success => {
-                for callback in &self.callbacks {
-                    if let Err(err) = callback.log_success_event(&payload) {
-                        self.dropped += 1;
-                        eprintln!("litellm-ai-gateway: log_success_event dropped: {err}");
-                    }
-                }
+                let response = CallbackValue::new("realtime", serde_json::Value::Null);
+                let report = runner
+                    .async_log_success_event(
+                        &ModelCallDetails::from_standard_logging_payload(payload),
+                        &response,
+                        timing,
+                    )
+                    .await;
+                self.dropped += report.dropped as u64;
             }
             SessionStatus::Failure => {
-                let error = crate::integrations::types::LoggingError {
+                let error = LoggingError {
                     message: "realtime session ended in failure".to_string(),
                     kind: "RealtimeSessionError".to_string(),
                 };
-                for callback in &self.callbacks {
-                    if let Err(err) = callback.log_failure_event(&payload, &error) {
-                        self.dropped += 1;
-                        eprintln!("litellm-ai-gateway: log_failure_event dropped: {err}");
-                    }
-                }
+                let response = CallbackValue::new(
+                    "error",
+                    serde_json::json!({
+                        "message": error.message,
+                        "kind": error.kind,
+                    }),
+                );
+                let report = runner
+                    .async_log_failure_event(
+                        &ModelCallDetails::from_standard_logging_payload(payload)
+                            .with_failure_error(error),
+                        Some(&response),
+                        timing,
+                    )
+                    .await;
+                self.dropped += report.dropped as u64;
             }
         }
     }
@@ -215,7 +232,8 @@ impl RealTimeStreaming {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::integrations::types::{LogError, LoggingError};
+    use crate::integrations::custom_logger::LogError;
+    use crate::integrations::custom_logger::LogFuture;
     use std::sync::atomic::{AtomicU64, Ordering};
 
     fn event(raw: &str) -> RealtimeEvent {
@@ -231,17 +249,28 @@ mod tests {
     }
 
     impl CustomLogger for CapturingLogger {
-        fn log_success_event(&self, payload: &StandardLoggingPayload) -> Result<(), LogError> {
-            self.calls.fetch_add(1, Ordering::SeqCst);
-            *self.last_model.lock().unwrap() = Some(payload.model.clone());
-            self.last_total_tokens
-                .store(payload.total_tokens, Ordering::SeqCst);
-            Ok(())
+        fn async_log_success_event<'a>(
+            &'a self,
+            model_call_details: &'a ModelCallDetails,
+            _response_obj: &'a CallbackValue,
+            _timing: CallbackTiming,
+        ) -> LogFuture<'a> {
+            Box::pin(async move {
+                let payload = model_call_details
+                    .standard_logging_payload
+                    .as_ref()
+                    .expect("standard logging payload");
+                self.calls.fetch_add(1, Ordering::SeqCst);
+                *self.last_model.lock().unwrap() = Some(payload.model.clone());
+                self.last_total_tokens
+                    .store(payload.total_tokens, Ordering::SeqCst);
+                Ok(())
+            })
         }
     }
 
-    #[test]
-    fn observe_accumulates_model_and_tokens_then_logs() {
+    #[tokio::test]
+    async fn observe_accumulates_model_and_tokens_then_logs() {
         let logger = Arc::new(CapturingLogger::default());
         let callbacks: Vec<Arc<dyn CustomLogger>> = vec![logger.clone()];
         let mut streaming = RealTimeStreaming::new(
@@ -284,7 +313,7 @@ mod tests {
             Some("hash123")
         );
 
-        streaming.log_messages(SessionStatus::Success);
+        streaming.log_messages(SessionStatus::Success).await;
         assert_eq!(logger.calls.load(Ordering::SeqCst), 1);
         assert_eq!(
             logger.last_model.lock().unwrap().as_deref(),
@@ -324,19 +353,26 @@ mod tests {
 
     /// A logger whose enqueue always fails should bump the dropped counter, not
     /// panic or propagate.
-    #[test]
-    fn failing_logger_bumps_dropped_counter() {
+    #[tokio::test]
+    async fn failing_logger_bumps_dropped_counter() {
         struct FailingLogger;
         impl CustomLogger for FailingLogger {
-            fn log_success_event(&self, _p: &StandardLoggingPayload) -> Result<(), LogError> {
-                Err(LogError::channel_full())
+            fn async_log_success_event<'a>(
+                &'a self,
+                _model_call_details: &'a ModelCallDetails,
+                _response_obj: &'a CallbackValue,
+                _timing: CallbackTiming,
+            ) -> LogFuture<'a> {
+                Box::pin(async { Err(LogError::channel_full()) })
             }
-            fn log_failure_event(
-                &self,
-                _p: &StandardLoggingPayload,
-                _e: &LoggingError,
-            ) -> Result<(), LogError> {
-                Err(LogError::channel_closed())
+
+            fn async_log_failure_event<'a>(
+                &'a self,
+                _model_call_details: &'a ModelCallDetails,
+                _response_obj: Option<&'a CallbackValue>,
+                _timing: CallbackTiming,
+            ) -> LogFuture<'a> {
+                Box::pin(async { Err(LogError::channel_closed()) })
             }
         }
         let callbacks: Vec<Arc<dyn CustomLogger>> = vec![Arc::new(FailingLogger)];
@@ -346,7 +382,7 @@ mod tests {
             "gpt-realtime".to_string(),
             RequestMetadata::default(),
         );
-        streaming.log_messages(SessionStatus::Success);
+        streaming.log_messages(SessionStatus::Success).await;
         assert_eq!(streaming.dropped(), 1);
     }
 }

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

@@ -162,5 +162,5 @@ async fn bridge(
     } else {
         SessionStatus::Failure
     };
-    collector.log_messages(status);
+    collector.log_messages(status).await;
 }

litellm-rust/crates/core/Cargo.toml

@@ -10,3 +10,6 @@ rand.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true
+
+[dev-dependencies]
+tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

litellm-rust/crates/core/src/call_lifecycle/README.md

@@ -0,0 +1,167 @@
+# Call lifecycle
+
+`litellm_core::call_lifecycle` is the shared execution wrapper for LiteLLM call
+types migrated to Rust. It owns lifecycle ordering, phase timing, and trace
+observer calls. It must not know about OCR, chat, messages, responses,
+completions, provider auth, request transforms, or response normalization.
+
+Call-type modules own their domain behavior. For example, OCR owns document
+payloads, OCR provider transforms, safe document fetch, guardrail payload shape,
+callback payload shape, and provider HTTP execution.
+
+## Runtime order
+
+Every wrapped call runs in this order:
+
+1. `async_pre_call_hook`
+2. `async_during_call_hook`
+3. provider call
+4. `async_log_success_event` or `async_log_failure_event`
+
+`async_pre_call_hook` receives the initial LiteLLM request shape. It is where
+pre-call custom guardrails run.
+
+`async_during_call_hook` converts the initial request into the provider-ready
+request. It is where provider config selection, parameter mapping, auth/header
+resolution, request transforms, and during-call guardrails belong.
+
+The provider call receives only the provider-ready request. It should execute
+I/O and call the provider response transform.
+
+Success and failure callbacks receive `CallLifecycleTiming`. Callback failures
+must not replace the original provider or guardrail result.
+
+## Trace contract
+
+The lifecycle runner records:
+
+- full call start and end time
+- `pre_call` phase timing
+- `during_call` phase timing
+- `provider_call` phase timing
+- `success_callback` phase timing
+- `failure_callback` phase timing
+
+`CallLifecycleObserver` receives phase start and end events. The default
+observer is a no-op. Future OTEL support should implement this observer instead
+of editing OCR, chat, messages, responses, completions, or provider modules.
+
+## Required shape
+
+Each migrated call type should use this folder shape:
+
+```text
+litellm-rust/crates/ai-gateway/src/<call_type>/
+  mod.rs          # thin public entrypoint
+  types.rs        # public request, prepared request, provider request, response types
+  prepare.rs      # model/provider/callback/guardrail setup
+  hooks.rs        # CallLifecycleHooks implementation
+  handler.rs      # provider I/O and response normalization
+  tests.rs        # call-type lifecycle and handler tests
+```
+
+Provider transforms can live in `litellm-rust/crates/core/src/providers/...`.
+Shared call-type helpers can live beside the call type, but generic lifecycle
+code stays in this folder.
+
+## Core API
+
+The prepared request implements `CallLifecycleRequest`:
+
+```rust
+impl CallLifecycleRequest for PreparedMessagesRequest {
+    fn lifecycle_context(&self) -> CallLifecycleContext {
+        CallLifecycleContext::new(
+            "messages",
+            self.model.clone(),
+            self.custom_llm_provider.clone(),
+            self.litellm_call_id.clone(),
+        )
+    }
+}
+```
+
+The call-type hooks implement `CallLifecycleHooks`:
+
+```rust
+impl CallLifecycleHooks<
+    PreparedMessagesRequest,
+    ProviderMessagesRequest,
+    MessagesResponse,
+> for MessagesLifecycleHooks {
+    fn async_pre_call_hook(...) {
+        // run pre-call custom guardrails against the LiteLLM request shape
+    }
+
+    fn async_during_call_hook(...) {
+        // map params, validate env, transform request, run during-call guardrails
+    }
+
+    fn async_log_success_event(...) {
+        // call async_log_success_event on configured custom loggers
+    }
+
+    fn async_log_failure_event(...) {
+        // call async_log_failure_event without swallowing the original error
+    }
+}
+```
+
+The public entrypoint stays thin:
+
+```rust
+pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<MessagesResponse> {
+    let PreparedMessagesCall { request, hooks } = prepare_messages_call(request)?;
+
+    CallLifecycle::default()
+        .run_request(request, &hooks, execute_messages_provider_call)
+        .await
+}
+```
+
+Use `run_request` for new call types. Keep `run` available only for specialized
+tests or existing code that already has a `CallLifecycleContext`.
+
+## Adding a new call type
+
+1. Add `<call_type>/types.rs`
+
+Define the public request accepted by the bridge, the prepared request used by
+the lifecycle runner, and the provider request consumed by the handler.
+
+2. Implement `CallLifecycleRequest`
+
+Return `call_type`, `model`, `custom_llm_provider`, and `litellm_call_id`.
+Do not put provider-specific logic here.
+
+3. Add `<call_type>/prepare.rs`
+
+Resolve model/provider once, generate or preserve `litellm_call_id`, construct
+callback and guardrail runners, and return `Prepared<CallType>Call`.
+
+4. Add `<call_type>/hooks.rs`
+
+Implement `CallLifecycleHooks`. Put pre-call guardrail payload construction,
+provider config selection, param mapping, request transform, during-call
+guardrail payload construction, and callback payload construction here.
+
+5. Add `<call_type>/handler.rs`
+
+Execute the provider request and normalize the provider response. Do not repeat
+provider-specific transforms here; call the provider config.
+
+6. Add tests
+
+Cover hook order, success callback payload, failure callback payload, pre-call
+guardrail blocking before provider I/O, during-call body mutation, and provider
+error mapping.
+
+## Review checklist
+
+- Core lifecycle has no call-type or provider-specific branches
+- Public call-type entrypoint only prepares and calls `run_request`
+- Provider behavior lives behind provider config/transformation code
+- Hook method names map to the Python custom logger and guardrail concepts
+- Phase timing is recorded once in lifecycle, not separately per call type
+- Callback failures never hide the original provider or guardrail error
+- Tests prove the provider socket is not touched when pre-call guardrails block

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

@@ -0,0 +1,414 @@
+use std::future::Future;
+use std::time::{Instant, SystemTime, UNIX_EPOCH};
+
+use crate::{CoreError, CoreResult};
+
+pub mod types;
+
+pub use types::{
+    CallLifecycleContext, CallLifecyclePhase, CallLifecyclePhaseTiming, CallLifecycleRequest,
+    CallLifecycleTiming,
+};
+
+pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
+    type PreCallFuture<'a>: Future<Output = CoreResult<InitialReq>> + Send + 'a
+    where
+        Self: 'a,
+        InitialReq: 'a,
+        ProviderReq: 'a,
+        Resp: 'a;
+
+    type DuringCallFuture<'a>: Future<Output = CoreResult<ProviderReq>> + Send + 'a
+    where
+        Self: 'a,
+        InitialReq: 'a,
+        ProviderReq: 'a,
+        Resp: 'a;
+
+    type SuccessFuture<'a>: Future<Output = ()> + Send + 'a
+    where
+        Self: 'a,
+        Resp: 'a;
+
+    type FailureFuture<'a>: Future<Output = ()> + Send + 'a
+    where
+        Self: 'a;
+
+    fn async_pre_call_hook<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        request: InitialReq,
+    ) -> Self::PreCallFuture<'a>;
+
+    fn async_during_call_hook<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        request: InitialReq,
+    ) -> Self::DuringCallFuture<'a>;
+
+    fn async_log_success_event<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        response: &'a Resp,
+        timing: &'a CallLifecycleTiming,
+    ) -> Self::SuccessFuture<'a>;
+
+    fn async_log_failure_event<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        error: &'a CoreError,
+        timing: &'a CallLifecycleTiming,
+    ) -> Self::FailureFuture<'a>;
+}
+
+pub trait CallLifecycleObserver: Send + Sync {
+    fn on_phase_start(&self, _context: &CallLifecycleContext, _phase: CallLifecyclePhase) {}
+
+    fn on_phase_end(&self, _context: &CallLifecycleContext, _timing: &CallLifecyclePhaseTiming) {}
+}
+
+#[derive(Default)]
+pub struct NoopCallLifecycleObserver;
+
+impl CallLifecycleObserver for NoopCallLifecycleObserver {}
+
+pub struct CallLifecycle<'a> {
+    observer: &'a dyn CallLifecycleObserver,
+}
+
+impl<'a> CallLifecycle<'a> {
+    pub fn new(observer: &'a dyn CallLifecycleObserver) -> Self {
+        Self { observer }
+    }
+
+    pub async fn run_request<InitialReq, ProviderReq, Resp, Hooks, ProviderCall, ProviderFuture>(
+        &self,
+        request: InitialReq,
+        hooks: &Hooks,
+        provider_call: ProviderCall,
+    ) -> CoreResult<Resp>
+    where
+        InitialReq: CallLifecycleRequest,
+        Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
+        ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
+        ProviderFuture: Future<Output = CoreResult<Resp>>,
+    {
+        let context = request.lifecycle_context();
+        self.run(context, request, hooks, provider_call).await
+    }
+
+    pub async fn run<InitialReq, ProviderReq, Resp, Hooks, ProviderCall, ProviderFuture>(
+        &self,
+        context: CallLifecycleContext,
+        request: InitialReq,
+        hooks: &Hooks,
+        provider_call: ProviderCall,
+    ) -> CoreResult<Resp>
+    where
+        Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
+        ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
+        ProviderFuture: Future<Output = CoreResult<Resp>>,
+    {
+        let call_start = epoch_seconds();
+        let mut phases = Vec::new();
+
+        let pre_call = self.start_phase(&context, CallLifecyclePhase::PreCall);
+        let request = match hooks.async_pre_call_hook(&context, request).await {
+            Ok(request) => {
+                phases.push(self.finish_phase(&context, pre_call));
+                request
+            }
+            Err(error) => {
+                phases.push(self.finish_phase(&context, pre_call));
+                self.log_failure(&context, hooks, &error, call_start, &mut phases)
+                    .await;
+                return Err(error);
+            }
+        };
+
+        let during_call = self.start_phase(&context, CallLifecyclePhase::DuringCall);
+        let provider_request = match hooks.async_during_call_hook(&context, request).await {
+            Ok(request) => {
+                phases.push(self.finish_phase(&context, during_call));
+                request
+            }
+            Err(error) => {
+                phases.push(self.finish_phase(&context, during_call));
+                self.log_failure(&context, hooks, &error, call_start, &mut phases)
+                    .await;
+                return Err(error);
+            }
+        };
+
+        let provider_phase = self.start_phase(&context, CallLifecyclePhase::ProviderCall);
+        let result = provider_call(provider_request).await;
+        phases.push(self.finish_phase(&context, provider_phase));
+
+        match &result {
+            Ok(response) => {
+                let success_phase = self.start_phase(&context, CallLifecyclePhase::SuccessCallback);
+                let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone());
+                hooks
+                    .async_log_success_event(&context, response, &timing)
+                    .await;
+                phases.push(self.finish_phase(&context, success_phase));
+            }
+            Err(error) => {
+                self.log_failure(&context, hooks, error, call_start, &mut phases)
+                    .await;
+            }
+        }
+
+        result
+    }
+
+    async fn log_failure<InitialReq, ProviderReq, Resp, Hooks>(
+        &self,
+        context: &CallLifecycleContext,
+        hooks: &Hooks,
+        error: &CoreError,
+        call_start: f64,
+        phases: &mut Vec<CallLifecyclePhaseTiming>,
+    ) where
+        Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
+    {
+        let failure_phase = self.start_phase(context, CallLifecyclePhase::FailureCallback);
+        let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone());
+        hooks.async_log_failure_event(context, error, &timing).await;
+        phases.push(self.finish_phase(context, failure_phase));
+    }
+
+    fn start_phase(&self, context: &CallLifecycleContext, phase: CallLifecyclePhase) -> PhaseStart {
+        self.observer.on_phase_start(context, phase);
+        PhaseStart {
+            phase,
+            start_time: epoch_seconds(),
+            started_at: Instant::now(),
+        }
+    }
+
+    fn finish_phase(
+        &self,
+        context: &CallLifecycleContext,
+        phase_start: PhaseStart,
+    ) -> CallLifecyclePhaseTiming {
+        let timing = CallLifecyclePhaseTiming {
+            phase: phase_start.phase,
+            start_time: phase_start.start_time,
+            end_time: epoch_seconds(),
+            duration: phase_start.started_at.elapsed(),
+        };
+        self.observer.on_phase_end(context, &timing);
+        timing
+    }
+}
+
+impl Default for CallLifecycle<'static> {
+    fn default() -> Self {
+        static OBSERVER: NoopCallLifecycleObserver = NoopCallLifecycleObserver;
+        Self::new(&OBSERVER)
+    }
+}
+
+struct PhaseStart {
+    phase: CallLifecyclePhase,
+    start_time: f64,
+    started_at: Instant,
+}
+
+fn epoch_seconds() -> f64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|duration| duration.as_secs_f64())
+        .unwrap_or(0.0)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::pin::Pin;
+    use std::sync::Mutex;
+
+    type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
+
+    #[derive(Default)]
+    struct RecordingHooks {
+        events: Mutex<Vec<&'static str>>,
+    }
+
+    struct RecordingRequest(String);
+
+    impl CallLifecycleRequest for RecordingRequest {
+        fn lifecycle_context(&self) -> CallLifecycleContext {
+            CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1")
+        }
+    }
+
+    impl RecordingHooks {
+        fn events(&self) -> Vec<&'static str> {
+            self.events.lock().unwrap().clone()
+        }
+    }
+
+    impl CallLifecycleHooks<String, String, String> for RecordingHooks {
+        type PreCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
+        type DuringCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
+        type SuccessFuture<'a> = BoxFuture<'a, ()>;
+        type FailureFuture<'a> = BoxFuture<'a, ()>;
+
+        fn async_pre_call_hook<'a>(
+            &'a self,
+            _context: &'a CallLifecycleContext,
+            request: String,
+        ) -> Self::PreCallFuture<'a> {
+            Box::pin(async move {
+                self.events.lock().unwrap().push("pre_call");
+                Ok(format!("{request}:pre"))
+            })
+        }
+
+        fn async_during_call_hook<'a>(
+            &'a self,
+            _context: &'a CallLifecycleContext,
+            request: String,
+        ) -> Self::DuringCallFuture<'a> {
+            Box::pin(async move {
+                self.events.lock().unwrap().push("during_call");
+                Ok(format!("{request}:during"))
+            })
+        }
+
+        fn async_log_success_event<'a>(
+            &'a self,
+            _context: &'a CallLifecycleContext,
+            _response: &'a String,
+            timing: &'a CallLifecycleTiming,
+        ) -> Self::SuccessFuture<'a> {
+            Box::pin(async move {
+                assert!(timing.end_time >= timing.start_time);
+                assert_eq!(timing.phases.len(), 3);
+                self.events.lock().unwrap().push("success");
+            })
+        }
+
+        fn async_log_failure_event<'a>(
+            &'a self,
+            _context: &'a CallLifecycleContext,
+            _error: &'a CoreError,
+            _timing: &'a CallLifecycleTiming,
+        ) -> Self::FailureFuture<'a> {
+            Box::pin(async move {
+                self.events.lock().unwrap().push("failure");
+            })
+        }
+    }
+
+    impl CallLifecycleHooks<RecordingRequest, String, String> for RecordingHooks {
+        type PreCallFuture<'a> = BoxFuture<'a, CoreResult<RecordingRequest>>;
+        type DuringCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
+        type SuccessFuture<'a> = BoxFuture<'a, ()>;
+        type FailureFuture<'a> = BoxFuture<'a, ()>;
+
+        fn async_pre_call_hook<'a>(
+            &'a self,
+            _context: &'a CallLifecycleContext,
+            request: RecordingRequest,
+        ) -> Self::PreCallFuture<'a> {
+            Box::pin(async move {
+                self.events.lock().unwrap().push("pre_call");
+                Ok(RecordingRequest(format!("{}:pre", request.0)))
+            })
+        }
+
+        fn async_during_call_hook<'a>(
+            &'a self,
+            _context: &'a CallLifecycleContext,
+            request: RecordingRequest,
+        ) -> Self::DuringCallFuture<'a> {
+            Box::pin(async move {
+                self.events.lock().unwrap().push("during_call");
+                Ok(format!("{}:during", request.0))
+            })
+        }
+
+        fn async_log_success_event<'a>(
+            &'a self,
+            _context: &'a CallLifecycleContext,
+            _response: &'a String,
+            _timing: &'a CallLifecycleTiming,
+        ) -> Self::SuccessFuture<'a> {
+            Box::pin(async move {
+                self.events.lock().unwrap().push("success");
+            })
+        }
+
+        fn async_log_failure_event<'a>(
+            &'a self,
+            _context: &'a CallLifecycleContext,
+            _error: &'a CoreError,
+            _timing: &'a CallLifecycleTiming,
+        ) -> Self::FailureFuture<'a> {
+            Box::pin(async move {
+                self.events.lock().unwrap().push("failure");
+            })
+        }
+    }
+
+    #[tokio::test]
+    async fn lifecycle_runs_hooks_around_provider_call() {
+        let hooks = RecordingHooks::default();
+        let response = CallLifecycle::default()
+            .run(
+                CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"),
+                "request".to_string(),
+                &hooks,
+                |request| async move {
+                    assert_eq!(request, "request:pre:during");
+                    Ok("response".to_string())
+                },
+            )
+            .await
+            .expect("call succeeds");
+
+        assert_eq!(response, "response");
+        assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]);
+    }
+
+    #[tokio::test]
+    async fn lifecycle_logs_failure_when_provider_fails() {
+        let hooks = RecordingHooks::default();
+        let error = CallLifecycle::default()
+            .run(
+                CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"),
+                "request".to_string(),
+                &hooks,
+                |_request| async move {
+                    Err::<String, CoreError>(CoreError::Network("provider down".to_string()))
+                },
+            )
+            .await
+            .expect_err("call fails");
+
+        assert_eq!(error, CoreError::Network("provider down".to_string()));
+        assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]);
+    }
+
+    #[tokio::test]
+    async fn lifecycle_can_run_any_request_with_embedded_context() {
+        let hooks = RecordingHooks::default();
+        let response = CallLifecycle::default()
+            .run_request(
+                RecordingRequest("request".to_string()),
+                &hooks,
+                |request| async move {
+                    assert_eq!(request, "request:pre:during");
+                    Ok("response".to_string())
+                },
+            )
+            .await
+            .expect("call succeeds");
+
+        assert_eq!(response, "response");
+        assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]);
+    }
+}

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

@@ -0,0 +1,75 @@
+use std::time::Duration;
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct CallLifecycleContext {
+    pub call_type: String,
+    pub model: String,
+    pub custom_llm_provider: String,
+    pub litellm_call_id: String,
+}
+
+impl CallLifecycleContext {
+    pub fn new(
+        call_type: impl Into<String>,
+        model: impl Into<String>,
+        custom_llm_provider: impl Into<String>,
+        litellm_call_id: impl Into<String>,
+    ) -> Self {
+        Self {
+            call_type: call_type.into(),
+            model: model.into(),
+            custom_llm_provider: custom_llm_provider.into(),
+            litellm_call_id: litellm_call_id.into(),
+        }
+    }
+}
+
+pub trait CallLifecycleRequest {
+    fn lifecycle_context(&self) -> CallLifecycleContext;
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum CallLifecyclePhase {
+    PreCall,
+    DuringCall,
+    ProviderCall,
+    SuccessCallback,
+    FailureCallback,
+}
+
+impl CallLifecyclePhase {
+    pub fn as_str(self) -> &'static str {
+        match self {
+            Self::PreCall => "pre_call",
+            Self::DuringCall => "during_call",
+            Self::ProviderCall => "provider_call",
+            Self::SuccessCallback => "success_callback",
+            Self::FailureCallback => "failure_callback",
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct CallLifecyclePhaseTiming {
+    pub phase: CallLifecyclePhase,
+    pub start_time: f64,
+    pub end_time: f64,
+    pub duration: Duration,
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct CallLifecycleTiming {
+    pub start_time: f64,
+    pub end_time: f64,
+    pub phases: Vec<CallLifecyclePhaseTiming>,
+}
+
+impl CallLifecycleTiming {
+    pub fn new(start_time: f64, end_time: f64, phases: Vec<CallLifecyclePhaseTiming>) -> Self {
+        Self {
+            start_time,
+            end_time,
+            phases,
+        }
+    }
+}

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

@@ -1,7 +1,9 @@
+pub mod call_lifecycle;
 pub mod error;
 pub mod ocr;
 pub mod providers;
 pub mod realtime;
 pub mod router;
+pub mod routing_utils;
 
 pub use error::{CoreError, CoreResult};

litellm-rust/crates/core/src/routing_utils/README.md

@@ -0,0 +1,7 @@
+# Routing Utils
+
+Shared helpers for deciding how a LiteLLM model routes to an LLM provider.
+Keep provider-name parsing, explicit `custom_llm_provider` handling, and model-prefix normalization here.
+Do not put deployment selection or load-balancing logic here; that belongs in `router`.
+Do not put provider HTTP transformation logic here; that belongs in `providers`.
+Helpers in this folder should be deterministic and easy to unit test without network calls.

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

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

litellm-rust/crates/core/src/routing_utils/provider.rs

@@ -0,0 +1,77 @@
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct CustomLlmProvider<'a> {
+    pub model: &'a str,
+    pub custom_llm_provider: &'a str,
+}
+
+pub fn get_custom_llm_provider<'a>(
+    model: &'a str,
+    custom_llm_provider: Option<&'a str>,
+) -> Option<CustomLlmProvider<'a>> {
+    if let Some(custom_llm_provider) = custom_llm_provider.filter(|provider| !provider.is_empty()) {
+        return Some(CustomLlmProvider {
+            model: strip_custom_llm_provider_prefix(model, custom_llm_provider),
+            custom_llm_provider,
+        });
+    }
+
+    let (custom_llm_provider, model) = model.split_once('/')?;
+    if custom_llm_provider.is_empty() || model.is_empty() {
+        return None;
+    }
+    Some(CustomLlmProvider {
+        model,
+        custom_llm_provider,
+    })
+}
+
+fn strip_custom_llm_provider_prefix<'a>(model: &'a str, custom_llm_provider: &str) -> &'a str {
+    model
+        .strip_prefix(custom_llm_provider)
+        .and_then(|model| model.strip_prefix('/'))
+        .unwrap_or(model)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn gets_custom_llm_provider_from_model_prefix() {
+        assert_eq!(
+            get_custom_llm_provider("mistral/mistral-ocr-latest", None),
+            Some(CustomLlmProvider {
+                model: "mistral-ocr-latest",
+                custom_llm_provider: "mistral",
+            })
+        );
+        assert_eq!(
+            get_custom_llm_provider("azure_ai/doc-intelligence/prebuilt-layout", None),
+            Some(CustomLlmProvider {
+                model: "doc-intelligence/prebuilt-layout",
+                custom_llm_provider: "azure_ai",
+            })
+        );
+        assert_eq!(get_custom_llm_provider("mistral-ocr-latest", None), None);
+        assert_eq!(get_custom_llm_provider("/model", None), None);
+        assert_eq!(get_custom_llm_provider("provider/", None), None);
+    }
+
+    #[test]
+    fn explicit_custom_llm_provider_strips_matching_model_prefix() {
+        assert_eq!(
+            get_custom_llm_provider("mistral/mistral-ocr-latest", Some("mistral")),
+            Some(CustomLlmProvider {
+                model: "mistral-ocr-latest",
+                custom_llm_provider: "mistral",
+            })
+        );
+        assert_eq!(
+            get_custom_llm_provider("mistral/mistral-ocr-latest", Some("vertex_ai")),
+            Some(CustomLlmProvider {
+                model: "mistral/mistral-ocr-latest",
+                custom_llm_provider: "vertex_ai",
+            })
+        );
+    }
+}

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

@@ -96,7 +96,6 @@ fn ocr(
     optional_params: Option<Py<PyAny>>,
     timeout_seconds: Option<f64>,
 ) -> PyResult<Py<PyAny>> {
-    let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string());
     let (document, extra_headers, optional_params, timeout) = marshal_inputs(
         py,
         document,
@@ -111,10 +110,14 @@ fn ocr(
             document,
             api_key: api_key.as_deref(),
             api_base: api_base.as_deref(),
-            custom_llm_provider: &custom_llm_provider,
+            custom_llm_provider: custom_llm_provider.as_deref(),
             extra_headers,
             optional_params,
             timeout,
+            callbacks: Vec::new(),
+            guardrails: Vec::new(),
+            request_metadata: Default::default(),
+            litellm_call_id: None,
         }))
     });
 
@@ -138,7 +141,6 @@ fn aocr(
     optional_params: Option<Py<PyAny>>,
     timeout_seconds: Option<f64>,
 ) -> PyResult<Bound<'_, PyAny>> {
-    let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string());
     let (document, extra_headers, optional_params, timeout) = marshal_inputs(
         py,
         document,
@@ -153,10 +155,14 @@ fn aocr(
             document,
             api_key: api_key.as_deref(),
             api_base: api_base.as_deref(),
-            custom_llm_provider: &custom_llm_provider,
+            custom_llm_provider: custom_llm_provider.as_deref(),
             extra_headers,
             optional_params,
             timeout,
+            callbacks: Vec::new(),
+            guardrails: Vec::new(),
+            request_metadata: Default::default(),
+            litellm_call_id: None,
         })
         .await
         .map_err(core_error_to_pyerr)?;

litellm/__init__.py

@@ -1394,7 +1394,7 @@ from .skills.main import (
 )
 from .containers.main import *
 from .ocr.main import *
-from .ocr.rust_bridge import use_litellm_rust
+from .rust_bridge.ocr import use_litellm_rust
 from .rag.main import *
 from .sandbox.main import *
 from .search.main import *

litellm/caching/_embedding_router.py

@@ -0,0 +1,45 @@
+"""Shared selection of the embedding path for semantic caches.
+
+Both the Redis and qdrant semantic caches need the same decision: when the
+configured embedding model is a proxy Router deployment, embeddings must run
+through the Router so per-deployment auth (e.g. Bedrock aws_role_name) is
+applied. Otherwise fall back to a direct litellm embedding call.
+
+This module is dependency-injected: callers pass the proxy ``llm_router`` and
+``llm_model_list`` in, so the decision logic is unit-testable without importing
+``litellm.proxy.proxy_server``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from litellm.router import Router
+
+
+def resolve_embedding_router(
+    embedding_model: str,
+    llm_router: Router | None,
+    llm_model_list: list[dict[str, Any]] | None,
+) -> Router | None:
+    """Return ``llm_router`` iff it serves ``embedding_model`` as a deployment."""
+    if llm_router is None:
+        return None
+    router_model_names: list[str] = (
+        [m["model_name"] for m in llm_model_list if "model_name" in m]
+        if llm_model_list is not None
+        else []
+    )
+    if embedding_model in router_model_names:
+        return llm_router
+    return None
+
+
+def build_router_embedding_metadata(
+    request_metadata: dict[str, Any] | None,
+) -> dict[str, Any]:
+    """Forward the caller's full metadata, flagged as a semantic-cache embedding."""
+    metadata: dict[str, Any] = dict(request_metadata or {})
+    metadata["semantic-cache-embedding"] = True
+    return metadata

litellm/caching/caching.py

@@ -574,8 +574,9 @@ class Cache:
             if prompt_kwarg in kwargs:
                 cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg]
 
-        if isinstance(kwargs.get("metadata"), dict):
-            cache_lookup_kwargs["metadata"] = {}
+        metadata = kwargs.get("metadata")
+        if isinstance(metadata, dict):
+            cache_lookup_kwargs["metadata"] = dict(metadata)
 
         return cache_lookup_kwargs
 

litellm/caching/qdrant_semantic_cache.py

@@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
 )
 from litellm.types.utils import EmbeddingResponse
 
+from ._embedding_router import build_router_embedding_metadata, resolve_embedding_router
 from .base_cache import BaseCache
 
 
@@ -219,37 +220,50 @@ class QdrantSemanticCache(BaseCache):
         cached_key = payload.get(self.CACHE_KEY_FIELD_NAME)
         return cached_key is not None and str(cached_key) == str(key)
 
-    async def _get_async_embedding(self, prompt: str, **kwargs) -> Any:
-        llm_model_list = None
-        llm_router = None
-
+    def _get_embedding(
+        self, prompt: str, metadata: Dict[str, Any] | None = None
+    ) -> EmbeddingResponse:
+        """Embed via the proxy Router when it serves the model, else direct."""
         try:
-            from litellm.proxy.proxy_server import (
-                llm_model_list as proxy_llm_model_list,
-                llm_router as proxy_llm_router,
-            )
-
-            llm_model_list = proxy_llm_model_list
-            llm_router = proxy_llm_router
+            from litellm.proxy.proxy_server import llm_model_list, llm_router
         except ImportError:
-            pass
+            llm_model_list = None
+            llm_router = None
 
-        router_model_names = (
-            [m["model_name"] for m in llm_model_list]
-            if llm_model_list is not None
-            else []
+        router = resolve_embedding_router(
+            self.embedding_model, llm_router, llm_model_list
         )
-        if llm_router is not None and self.embedding_model in router_model_names:
-            user_api_key = kwargs.get("metadata", {}).get("user_api_key", "")
-            return await llm_router.aembedding(
+        if router is not None:
+            return router.embedding(
                 model=self.embedding_model,
                 input=prompt,
                 cache={"no-store": True, "no-cache": True},
-                metadata={
-                    "user_api_key": user_api_key,
-                    "semantic-cache-embedding": True,
-                    "trace_id": kwargs.get("metadata", {}).get("trace_id", None),
-                },
+                metadata=build_router_embedding_metadata(metadata),
+            )
+        return litellm.embedding(
+            model=self.embedding_model,
+            input=prompt,
+            cache={"no-store": True, "no-cache": True},
+        )
+
+    async def _get_async_embedding(
+        self, prompt: str, metadata: Dict[str, Any] | None = None
+    ) -> EmbeddingResponse:
+        try:
+            from litellm.proxy.proxy_server import llm_model_list, llm_router
+        except ImportError:
+            llm_model_list = None
+            llm_router = None
+
+        router = resolve_embedding_router(
+            self.embedding_model, llm_router, llm_model_list
+        )
+        if router is not None:
+            return await router.aembedding(
+                model=self.embedding_model,
+                input=prompt,
+                cache={"no-store": True, "no-cache": True},
+                metadata=build_router_embedding_metadata(metadata),
             )
 
         return await litellm.aembedding(
@@ -269,11 +283,7 @@ class QdrantSemanticCache(BaseCache):
         # create an embedding for prompt
         embedding_response = cast(
             EmbeddingResponse,
-            litellm.embedding(
-                model=self.embedding_model,
-                input=prompt,
-                cache={"no-store": True, "no-cache": True},
-            ),
+            self._get_embedding(prompt, metadata=kwargs.get("metadata")),
         )
 
         # get the embedding
@@ -312,11 +322,7 @@ class QdrantSemanticCache(BaseCache):
         # convert to embedding
         embedding_response = cast(
             EmbeddingResponse,
-            litellm.embedding(
-                model=self.embedding_model,
-                input=prompt,
-                cache={"no-store": True, "no-cache": True},
-            ),
+            self._get_embedding(prompt, metadata=kwargs.get("metadata")),
         )
 
         # get the embedding
@@ -388,7 +394,9 @@ class QdrantSemanticCache(BaseCache):
         # get the prompt
         messages = kwargs["messages"]
         prompt = get_str_from_messages(messages)
-        embedding_response = await self._get_async_embedding(prompt, **kwargs)
+        embedding_response = await self._get_async_embedding(
+            prompt, metadata=kwargs.get("metadata")
+        )
 
         # get the embedding
         embedding = embedding_response["data"][0]["embedding"]
@@ -424,7 +432,9 @@ class QdrantSemanticCache(BaseCache):
         messages = kwargs["messages"]
         prompt = get_str_from_messages(messages)
 
-        embedding_response = await self._get_async_embedding(prompt, **kwargs)
+        embedding_response = await self._get_async_embedding(
+            prompt, metadata=kwargs.get("metadata")
+        )
 
         # get the embedding
         embedding = embedding_response["data"][0]["embedding"]

litellm/caching/redis_semantic_cache.py

@@ -16,12 +16,13 @@ import os
 from typing import Any, Dict, List, Optional, Tuple, cast
 
 import litellm
-from litellm._logging import print_verbose
+from litellm._logging import print_verbose, verbose_logger
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     get_str_from_messages,
 )
 from litellm.types.utils import EmbeddingResponse
 
+from ._embedding_router import build_router_embedding_metadata, resolve_embedding_router
 from .base_cache import BaseCache
 
 
@@ -67,9 +68,6 @@ class RedisSemanticCache(BaseCache):
             Exception: If similarity_threshold is not provided or required Redis
                 connection information is missing
         """
-        from redisvl.extensions.llmcache import SemanticCache  # type: ignore[import-not-found, import-untyped]
-        from redisvl.utils.vectorize import CustomTextVectorizer  # type: ignore[import-not-found, import-untyped]
-
         if index_name is None:
             index_name = self.DEFAULT_REDIS_INDEX_NAME
 
@@ -107,15 +105,42 @@ class RedisSemanticCache(BaseCache):
 
         print_verbose(f"Redis semantic-cache redis_url: {redis_url}")
 
-        # Initialize the Redis vectorizer and cache
-        cache_vectorizer = CustomTextVectorizer(self._get_embedding)
+        # Defer redisvl index construction until first use. redisvl's
+        # CustomTextVectorizer eagerly embeds a probe string at construction;
+        # building lazily ensures that probe runs after llm_router is wired so
+        # per-deployment auth (e.g. Bedrock aws_role_name) is applied.
+        self._index_name = index_name
+        self._redis_url = redis_url
+        self._llmcache = None
 
-        self.llmcache = self._init_semantic_cache(
-            semantic_cache_cls=SemanticCache,
-            index_name=index_name,
-            redis_url=redis_url,
-            cache_vectorizer=cache_vectorizer,
-        )
+    @property
+    def llmcache(self) -> object:
+        if getattr(self, "_llmcache", None) is None:
+            self._llmcache = self._build_llmcache()
+        return self._llmcache
+
+    @llmcache.setter
+    def llmcache(self, value: object) -> None:
+        self._llmcache = value
+
+    def _build_llmcache(self) -> object:
+        # CustomTextVectorizer probes its embedding dimension at construction by
+        # embedding "dimension test", so the first cache request issues one extra
+        # billable embedding on top of the request's own.
+        from redisvl.extensions.llmcache import SemanticCache  # type: ignore[import-not-found, import-untyped]
+        from redisvl.utils.vectorize import CustomTextVectorizer  # type: ignore[import-not-found, import-untyped]
+
+        try:
+            cache_vectorizer = CustomTextVectorizer(self._get_embedding)
+            return self._init_semantic_cache(
+                semantic_cache_cls=SemanticCache,
+                index_name=self._index_name,
+                redis_url=self._redis_url,
+                cache_vectorizer=cache_vectorizer,
+            )
+        except Exception as e:
+            verbose_logger.error(f"Redis semantic-cache index build failed: {e}")
+            raise
 
     @classmethod
     def _cache_key_filterable_field(cls) -> Dict[str, str]:
@@ -285,27 +310,43 @@ class RedisSemanticCache(BaseCache):
             return dict_method()
         return value
 
-    def _get_embedding(self, prompt: str) -> List[float]:
+    def _get_embedding(
+        self, prompt: str, metadata: Dict[str, Any] | None = None
+    ) -> List[float]:
         """
-        Generate an embedding vector for the given prompt using the configured embedding model.
-
-        Args:
-            prompt: The text to generate an embedding for
-
-        Returns:
-            List[float]: The embedding vector
+        Routes through the proxy Router when the embedding model is a Router
+        deployment so per-deployment auth (e.g. Bedrock aws_role_name) applies,
+        mirroring ``_get_async_embedding``; otherwise embeds directly.
         """
-        # Create an embedding from prompt
-        embedding_response = cast(
-            EmbeddingResponse,
-            litellm.embedding(
-                model=self.embedding_model,
-                input=prompt,
-                cache={"no-store": True, "no-cache": True},
-            ),
+        try:
+            from litellm.proxy.proxy_server import llm_model_list, llm_router
+        except ImportError:
+            llm_model_list = None
+            llm_router = None
+
+        router = resolve_embedding_router(
+            self.embedding_model, llm_router, llm_model_list
         )
-        embedding = embedding_response["data"][0]["embedding"]
-        return embedding
+        if router is not None:
+            embedding_response = cast(
+                EmbeddingResponse,
+                router.embedding(
+                    model=self.embedding_model,
+                    input=prompt,
+                    cache={"no-store": True, "no-cache": True},
+                    metadata=build_router_embedding_metadata(metadata),
+                ),
+            )
+        else:
+            embedding_response = cast(
+                EmbeddingResponse,
+                litellm.embedding(
+                    model=self.embedding_model,
+                    input=prompt,
+                    cache={"no-store": True, "no-cache": True},
+                ),
+            )
+        return embedding_response["data"][0]["embedding"]
 
     def _get_cache_logic(self, cached_response: Any) -> Any:
         """
@@ -357,7 +398,12 @@ class RedisSemanticCache(BaseCache):
 
             value_str = str(value)
 
-            store_kwargs: Dict[str, Any] = {
+            prompt_embedding = self._get_embedding(
+                prompt, metadata=kwargs.get("metadata")
+            )
+
+            store_kwargs: dict[str, Any] = {
+                "vector": prompt_embedding,
                 "filters": self._get_cache_filters(key),
             }
 
@@ -393,8 +439,12 @@ class RedisSemanticCache(BaseCache):
 
             # Check the cache for semantically similar prompts in this exact
             # LiteLLM cache-key scope.
-            check_kwargs: Dict[str, Any] = {
+            prompt_embedding = self._get_embedding(
+                prompt, metadata=kwargs.get("metadata")
+            )
+            check_kwargs: dict[str, Any] = {
                 "prompt": prompt,
+                "vector": prompt_embedding,
                 "filter_expression": self._get_cache_key_filter_expression(key),
             }
             results = self.llmcache.check(**check_kwargs)
@@ -435,49 +485,42 @@ class RedisSemanticCache(BaseCache):
             print_verbose(f"Error retrieving from Redis semantic cache: {str(e)}")
             kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
 
-    async def _get_async_embedding(self, prompt: str, **kwargs) -> List[float]:
+    async def _get_async_embedding(
+        self, prompt: str, metadata: Dict[str, Any] | None = None
+    ) -> List[float]:
         """
         Asynchronously generate an embedding for the given prompt.
 
         Args:
             prompt: The text to generate an embedding for
-            **kwargs: Additional arguments that may contain metadata
+            metadata: Request metadata forwarded to the Router embedding call
 
         Returns:
             List[float]: The embedding vector
         """
-        from litellm.proxy.proxy_server import llm_model_list, llm_router
-
-        # Route the embedding request through the proxy if appropriate
-        router_model_names = (
-            [m["model_name"] for m in llm_model_list]
-            if llm_model_list is not None
-            else []
-        )
-
         try:
-            if llm_router is not None and self.embedding_model in router_model_names:
-                # Use the router for embedding generation
-                user_api_key = kwargs.get("metadata", {}).get("user_api_key", "")
-                embedding_response = await llm_router.aembedding(
+            from litellm.proxy.proxy_server import llm_model_list, llm_router
+        except ImportError:
+            llm_model_list = None
+            llm_router = None
+
+        router = resolve_embedding_router(
+            self.embedding_model, llm_router, llm_model_list
+        )
+        try:
+            if router is not None:
+                embedding_response = await router.aembedding(
                     model=self.embedding_model,
                     input=prompt,
                     cache={"no-store": True, "no-cache": True},
-                    metadata={
-                        "user_api_key": user_api_key,
-                        "semantic-cache-embedding": True,
-                        "trace_id": kwargs.get("metadata", {}).get("trace_id", None),
-                    },
+                    metadata=build_router_embedding_metadata(metadata),
                 )
             else:
-                # Generate embedding directly
                 embedding_response = await litellm.aembedding(
                     model=self.embedding_model,
                     input=prompt,
                     cache={"no-store": True, "no-cache": True},
                 )
-
-            # Extract and return the embedding vector
             return embedding_response["data"][0]["embedding"]
         except Exception as e:
             print_verbose(f"Error generating async embedding: {str(e)}")
@@ -504,9 +547,11 @@ class RedisSemanticCache(BaseCache):
             value_str = str(value)
 
             # Generate embedding for the value (response) to cache
-            prompt_embedding = await self._get_async_embedding(prompt, **kwargs)
+            prompt_embedding = await self._get_async_embedding(
+                prompt, metadata=kwargs.get("metadata")
+            )
 
-            store_kwargs: Dict[str, Any] = {
+            store_kwargs: dict[str, Any] = {
                 "vector": prompt_embedding,
                 "filters": self._get_cache_filters(key),
             }
@@ -544,11 +589,13 @@ class RedisSemanticCache(BaseCache):
                 return None
 
             # Generate embedding for the prompt
-            prompt_embedding = await self._get_async_embedding(prompt, **kwargs)
+            prompt_embedding = await self._get_async_embedding(
+                prompt, metadata=kwargs.get("metadata")
+            )
 
             # Check the cache for semantically similar prompts in this exact
             # LiteLLM cache-key scope.
-            check_kwargs: Dict[str, Any] = {
+            check_kwargs: dict[str, Any] = {
                 "prompt": prompt,
                 "vector": prompt_embedding,
                 "filter_expression": self._get_cache_key_filter_expression(key),

litellm/google_genai/main.py

@@ -49,6 +49,7 @@ class GenerateContentSetupResult(BaseModel):
     custom_llm_provider: str
     generate_content_provider_config: Optional[BaseGoogleGenAIGenerateContentConfig]
     generate_content_config_dict: Dict[str, Any]
+    native_request_fields: dict[str, object]
     litellm_params: GenericLiteLLMParams
     litellm_logging_obj: LiteLLMLoggingObj
     litellm_call_id: Optional[str]
@@ -152,6 +153,7 @@ class GenerateContentHelper:
                 request_body={},  # Will be handled by adapter
                 generate_content_provider_config=None,  # type: ignore
                 generate_content_config_dict=dict(config or {}),
+                native_request_fields={},
                 litellm_params=litellm_params,
                 litellm_logging_obj=litellm_logging_obj,
                 litellm_call_id=litellm_call_id,
@@ -171,6 +173,12 @@ class GenerateContentHelper:
         system_instruction = kwargs.get("systemInstruction") or kwargs.get(
             "system_instruction"
         )
+        # Native top-level REST fields arrive as loose kwargs and are otherwise dropped.
+        native_request_fields: dict[str, object] = {
+            field: kwargs[field]
+            for field in generate_content_provider_config.get_generate_content_request_top_level_fields()
+            if field in kwargs
+        }
         request_body = (
             generate_content_provider_config.transform_generate_content_request(
                 model=model,
@@ -201,12 +209,29 @@ class GenerateContentHelper:
             request_body=request_body,
             generate_content_provider_config=generate_content_provider_config,
             generate_content_config_dict=generate_content_config_dict,
+            native_request_fields=native_request_fields,
             litellm_params=litellm_params,
             litellm_logging_obj=litellm_logging_obj,
             litellm_call_id=litellm_call_id,
         )
 
 
+def _merge_native_request_fields(
+    native_request_fields: dict[str, object],
+    extra_body: dict[str, object] | None,
+) -> dict[str, object] | None:
+    """
+    Merge native top-level request fields into ``extra_body`` so the HTTP handler
+    forwards them verbatim onto the outgoing request body. An explicit ``extra_body``
+    value wins on conflict. Returns ``None`` only when there is genuinely nothing to
+    forward (no native fields and no caller-supplied ``extra_body``), preserving the
+    prior behavior without discarding an explicit ``extra_body={}``.
+    """
+    if not native_request_fields and extra_body is None:
+        return None
+    return {**native_request_fields, **(extra_body or {})}
+
+
 @client
 async def agenerate_content(
     model: str,
@@ -350,7 +375,9 @@ def generate_content(
             litellm_params=setup_result.litellm_params,
             logging_obj=setup_result.litellm_logging_obj,
             extra_headers=extra_headers,
-            extra_body=extra_body,
+            extra_body=_merge_native_request_fields(
+                setup_result.native_request_fields, extra_body
+            ),
             timeout=timeout or request_timeout,
             _is_async=_is_async,
             client=kwargs.get("client"),
@@ -447,7 +474,9 @@ async def agenerate_content_stream(
             litellm_params=setup_result.litellm_params,
             logging_obj=setup_result.litellm_logging_obj,
             extra_headers=extra_headers,
-            extra_body=extra_body,
+            extra_body=_merge_native_request_fields(
+                setup_result.native_request_fields, extra_body
+            ),
             timeout=timeout or request_timeout,
             _is_async=True,
             client=kwargs.get("client"),
@@ -503,6 +532,11 @@ def generate_content_stream(
             **kwargs,
         )
 
+        # Extract systemInstruction from kwargs to pass to handler
+        system_instruction = kwargs.get("systemInstruction") or kwargs.get(
+            "system_instruction"
+        )
+
         # Check if we should use the adapter (when provider config is None)
         if setup_result.generate_content_provider_config is None:
             if "stream" in kwargs:
@@ -531,12 +565,15 @@ def generate_content_stream(
             litellm_params=setup_result.litellm_params,
             logging_obj=setup_result.litellm_logging_obj,
             extra_headers=extra_headers,
-            extra_body=extra_body,
+            extra_body=_merge_native_request_fields(
+                setup_result.native_request_fields, extra_body
+            ),
             timeout=timeout or request_timeout,
             _is_async=_is_async,
             client=kwargs.get("client"),
             stream=True,
             litellm_metadata=kwargs.get("litellm_metadata", {}),
+            system_instruction=system_instruction,
         )
 
     except Exception as e:

litellm/litellm_core_utils/llm_cost_calc/utils.py

@@ -485,6 +485,7 @@ class PromptTokensDetailsResult(TypedDict):
     character_count: int
     image_count: int
     video_length_seconds: float
+    audio_length_seconds: float
 
 
 def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
@@ -535,6 +536,13 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
         )
         or 0.0
     )
+    audio_length_seconds = (
+        cast(
+            Optional[float],
+            getattr(usage.prompt_tokens_details, "audio_length_seconds", 0),
+        )
+        or 0.0
+    )
 
     return PromptTokensDetailsResult(
         cache_hit_tokens=cache_hit_tokens,
@@ -546,6 +554,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
         character_count=character_count,
         image_count=image_count,
         video_length_seconds=float(video_length_seconds),
+        audio_length_seconds=float(audio_length_seconds),
     )
 
 
@@ -667,6 +676,14 @@ def _calculate_input_cost(
             prompt_tokens_details["video_length_seconds"],
         )
 
+    ### AUDIO LENGTH COST
+    if prompt_tokens_details["audio_length_seconds"]:
+        prompt_cost += calculate_cost_component(
+            model_info,
+            "input_cost_per_audio_per_second",
+            prompt_tokens_details["audio_length_seconds"],
+        )
+
     return prompt_cost
 
 
@@ -743,6 +760,7 @@ def generic_cost_per_token(
         character_count=0,
         image_count=0,
         video_length_seconds=0.0,
+        audio_length_seconds=0.0,
     )
     if usage.prompt_tokens_details:
         prompt_tokens_details = _parse_prompt_tokens_details(usage)

litellm/litellm_core_utils/prompt_templates/common_utils.py

@@ -466,7 +466,7 @@ def get_format_from_file_id(file_id: Optional[str]) -> Optional[str]:
 
 def update_messages_with_model_file_ids(
     messages: List[AllMessageValues],
-    model_id: str,
+    model_id: str | None,
     model_file_id_mapping: Dict[str, Dict[str, str]],
 ) -> List[AllMessageValues]:
     """
@@ -519,7 +519,7 @@ def update_messages_with_model_file_ids(
                         if file_id:
                             provider_file_id = (
                                 model_file_id_mapping.get(file_id, {}).get(model_id)
-                                if model_file_id_mapping
+                                if model_file_id_mapping and model_id is not None
                                 else None
                             )
                             if (

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -585,6 +585,17 @@ class ChunkProcessor:
         # # Update usage information if needed
         prompt_tokens = 0
         completion_tokens = 0
+        # Anthropic's `message_start` SSE event carries usage.output_tokens=1 as a
+        # cursor/placeholder; the real value only arrives in `message_delta`.
+        # If a stream is cancelled before `message_delta` lands, the last-wins
+        # accumulator below leaves completion_tokens stuck at 1 — which then
+        # bypasses the `completion_tokens or token_counter(...)` fallback in
+        # calculate_usage() because 1 is truthy. Count the completion-bearing
+        # usage events so `_reset_anthropic_cursor_completion_tokens` can tell a
+        # legitimate single-token reply (Anthropic emits 1 in BOTH message_start
+        # AND message_delta, so >=2 events is positive evidence message_delta
+        # arrived) from a stale lone cursor.
+        completion_usage_updates = 0
         ## anthropic prompt caching information ##
         cache_creation_input_tokens: Optional[int] = None
         cache_read_input_tokens: Optional[int] = None
@@ -617,6 +628,7 @@ class ChunkProcessor:
                     and usage_chunk_dict["completion_tokens"] > 0
                 ):
                     completion_tokens = usage_chunk_dict["completion_tokens"]
+                    completion_usage_updates += 1
                 if usage_chunk_dict["cache_creation_input_tokens"] is not None and (
                     usage_chunk_dict["cache_creation_input_tokens"] > 0
                     or cache_creation_input_tokens is None
@@ -667,6 +679,12 @@ class ChunkProcessor:
 
                 prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"]
 
+        completion_tokens = self._reset_anthropic_cursor_completion_tokens(
+            chunks=chunks,
+            completion_tokens=completion_tokens,
+            completion_usage_updates=completion_usage_updates,
+        )
+
         return UsagePerChunk(
             prompt_tokens=prompt_tokens,
             completion_tokens=completion_tokens,
@@ -678,6 +696,47 @@ class ChunkProcessor:
             prompt_tokens_details=prompt_tokens_details,
         )
 
+    @staticmethod
+    def _reset_anthropic_cursor_completion_tokens(
+        chunks: list[dict[str, Any] | ModelResponse],
+        completion_tokens: int,
+        completion_usage_updates: int,
+    ) -> int:
+        """Reset a stale Anthropic ``message_start`` cursor placeholder to 0.
+
+        See the ``completion_usage_updates`` comment in
+        ``_calculate_usage_per_chunk``. The accumulated value is NOT a stale
+        cursor when either it is > 1 (definitely not a placeholder) or we saw
+        >= 2 completion-bearing usage events (positive evidence ``message_delta``
+        arrived). Otherwise — the only completion update we ever saw was the
+        Anthropic ``message_start`` cursor (=1) — reset to 0 so
+        ``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates
+        from the actually-received completion text instead of trusting the
+        placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the
+        heuristic (which encodes Anthropic's specific message_start SSE shape)
+        does not silently affect other providers that may legitimately report
+        ``completion_tokens=1`` from a single usage event.
+        """
+        saw_non_cursor_completion = (
+            completion_tokens > 1 or completion_usage_updates >= 2
+        )
+        if saw_non_cursor_completion:
+            return completion_tokens
+
+        custom_llm_provider: Optional[str] = None
+        if chunks:
+            first_chunk = chunks[0]
+            if isinstance(first_chunk, dict):
+                hp = first_chunk.get("_hidden_params")
+            else:
+                hp = getattr(first_chunk, "_hidden_params", None)
+            if isinstance(hp, dict):
+                custom_llm_provider = hp.get("custom_llm_provider")
+
+        if custom_llm_provider == "anthropic" and completion_tokens == 1:
+            return 0
+        return completion_tokens
+
     def calculate_usage(
         self,
         chunks: List[Union[Dict[str, Any], ModelResponse]],

litellm/llms/anthropic/batches/transformation.py

@@ -43,7 +43,9 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
         api_base: Optional[str] = None,
     ) -> dict:
         """Validate and prepare environment-specific headers and parameters."""
-        auth_header = self.anthropic_model_info.get_auth_header(api_key)
+        if api_base is None and isinstance(litellm_params, dict):
+            api_base = litellm_params.get("api_base")
+        auth_header = self.anthropic_model_info.get_auth_header(api_key, api_base)
         if auth_header is None:
             raise ValueError(
                 "Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params"

litellm/llms/anthropic/chat/handler.py

@@ -364,6 +364,7 @@ class AnthropicChatCompletion(BaseLLM):
             messages=messages,
             optional_params={**optional_params, "is_vertex_request": is_vertex_request},
             litellm_params=litellm_params,
+            api_base=api_base,
         )
 
         config = ProviderConfigManager.get_provider_chat_config(

litellm/llms/anthropic/common_utils.py

@@ -548,6 +548,19 @@ class AnthropicModelInfo(BaseLLMModelInfo):
 
         return list(set(betas))
 
+    @staticmethod
+    def _make_api_key_auth_header(
+        api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False
+    ) -> dict:
+        if use_bearer_for_custom_base and (
+            api_base
+            and "api.anthropic.com" not in api_base
+            and not api_key.startswith("sk-ant-")
+        ):
+            value = api_key if api_key.startswith("Bearer ") else f"Bearer {api_key}"
+            return {"authorization": value}
+        return {"x-api-key": api_key}
+
     def get_anthropic_headers(
         self,
         api_key: Optional[str] = None,
@@ -567,6 +580,8 @@ class AnthropicModelInfo(BaseLLMModelInfo):
         user_anthropic_beta_headers: Optional[List[str]] = None,
         code_execution_tool_used: bool = False,
         container_with_skills_used: bool = False,
+        api_base: str | None = None,
+        use_bearer_for_custom_base: bool = False,
     ) -> dict:
         betas = set()
         # Anthropic no longer requires the prompt-caching beta header
@@ -615,7 +630,11 @@ class AnthropicModelInfo(BaseLLMModelInfo):
         elif auth_token and not api_key:
             headers["authorization"] = f"Bearer {auth_token}"
         elif api_key:
-            headers["x-api-key"] = api_key
+            headers.update(
+                self._make_api_key_auth_header(
+                    api_key, api_base, use_bearer_for_custom_base
+                )
+            )
 
         if user_anthropic_beta_headers is not None:
             betas.update(user_anthropic_beta_headers)
@@ -644,6 +663,12 @@ class AnthropicModelInfo(BaseLLMModelInfo):
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
     ) -> Dict:
+        if api_base is None and isinstance(litellm_params, dict):
+            api_base = litellm_params.get("api_base")
+        use_bearer_for_custom_base: bool = bool(
+            isinstance(litellm_params, dict)
+            and litellm_params.get("use_bearer_for_custom_base", False)
+        )
         # Check for Anthropic OAuth token in headers
         headers, api_key = optionally_handle_anthropic_oauth(
             headers=headers, api_key=api_key
@@ -699,6 +724,8 @@ class AnthropicModelInfo(BaseLLMModelInfo):
             effort_used=effort_used,
             code_execution_tool_used=code_execution_tool_used,
             container_with_skills_used=container_with_skills_used,
+            api_base=api_base,
+            use_bearer_for_custom_base=use_bearer_for_custom_base,
         )
 
         headers = {**headers, **anthropic_headers}
@@ -734,18 +761,24 @@ class AnthropicModelInfo(BaseLLMModelInfo):
         return auth_token or get_secret_str("ANTHROPIC_AUTH_TOKEN")
 
     @staticmethod
-    def get_auth_header(api_key: Optional[str] = None) -> Optional[dict]:
+    def get_auth_header(
+        api_key: str | None = None,
+        api_base: str | None = None,
+        use_bearer_for_custom_base: bool = False,
+    ) -> dict | None:
         """Resolve Anthropic credentials and return the appropriate auth header dict.
 
-        Checks ANTHROPIC_API_KEY first (-> x-api-key), then
-        ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer).
+        Checks ANTHROPIC_API_KEY first (-> x-api-key or Bearer depending on
+        use_bearer_for_custom_base), then ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer).
         Returns None if neither is available.
         """
         resolved_key = AnthropicModelInfo.get_api_key(api_key)
         if resolved_key is not None:
             if is_anthropic_oauth_key(resolved_key):
                 return {"authorization": f"Bearer {resolved_key}"}
-            return {"x-api-key": resolved_key}
+            return AnthropicModelInfo._make_api_key_auth_header(
+                resolved_key, api_base, use_bearer_for_custom_base
+            )
         auth_token = AnthropicModelInfo.get_auth_token()
         if auth_token is not None:
             return {"authorization": f"Bearer {auth_token}"}
@@ -759,7 +792,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
         self, api_key: Optional[str] = None, api_base: Optional[str] = None
     ) -> List[str]:
         api_base = AnthropicModelInfo.get_api_base(api_base)
-        auth_header = AnthropicModelInfo.get_auth_header(api_key)
+        auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base)
         if api_base is None or auth_header is None:
             raise ValueError(
                 "ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN is not set. Please set the environment variable, to query Anthropic's `/models` endpoint."

litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py

@@ -205,32 +205,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
         if "delta" not in merged_chunk:
             merged_chunk["delta"] = {}
 
-        uncached_input_tokens = chunk.usage.prompt_tokens or 0
-        if (
-            hasattr(chunk.usage, "prompt_tokens_details")
-            and chunk.usage.prompt_tokens_details
-        ):
-            cached_tokens = (
-                getattr(chunk.usage.prompt_tokens_details, "cached_tokens", 0) or 0
-            )
-            uncached_input_tokens -= cached_tokens
+        from .transformation import LiteLLMAnthropicMessagesAdapter
 
-        usage_dict: UsageDelta = {
-            "input_tokens": uncached_input_tokens,
-            "output_tokens": chunk.usage.completion_tokens or 0,
-        }
-        if (
-            hasattr(chunk.usage, "_cache_creation_input_tokens")
-            and chunk.usage._cache_creation_input_tokens > 0
-        ):
-            usage_dict["cache_creation_input_tokens"] = (
-                chunk.usage._cache_creation_input_tokens
-            )
-        if (
-            hasattr(chunk.usage, "_cache_read_input_tokens")
-            and chunk.usage._cache_read_input_tokens > 0
-        ):
-            usage_dict["cache_read_input_tokens"] = chunk.usage._cache_read_input_tokens
+        usage_dict: UsageDelta = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
+            chunk.usage
+        )
         merged_chunk["usage"] = usage_dict
         if self.applied_edits and "context_management" not in merged_chunk:
             merged_chunk["context_management"] = ContextManagementResponse(

litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py

@@ -1399,6 +1399,95 @@ class LiteLLMAnthropicMessagesAdapter:
             return "tool_use"
         return "end_turn"
 
+    @staticmethod
+    def _positive_int(value: object) -> int:
+        if isinstance(value, bool):
+            return 0
+        if isinstance(value, int) and value > 0:
+            return value
+        if isinstance(value, float) and value.is_integer() and value > 0:
+            return int(value)
+        return 0
+
+    @classmethod
+    def _first_positive_usage_value(
+        cls, usage: Usage, field_names: tuple[str, ...]
+    ) -> int:
+        for field_name in field_names:
+            value = cls._positive_int(getattr(usage, field_name, None))
+            if value > 0:
+                return value
+        return 0
+
+    @classmethod
+    def _first_positive_prompt_tokens_detail_value(
+        cls, usage: Usage, field_names: tuple[str, ...]
+    ) -> int:
+        prompt_tokens_details = getattr(usage, "prompt_tokens_details", None)
+        if prompt_tokens_details is None:
+            return 0
+
+        for field_name in field_names:
+            if isinstance(prompt_tokens_details, dict):
+                value = cls._positive_int(prompt_tokens_details.get(field_name))
+            else:
+                value = cls._positive_int(
+                    getattr(prompt_tokens_details, field_name, None)
+                )
+            if value > 0:
+                return value
+        return 0
+
+    @classmethod
+    def _get_cache_read_input_tokens(cls, usage: Usage) -> int:
+        explicit_value = cls._first_positive_usage_value(
+            usage, ("cache_read_input_tokens", "_cache_read_input_tokens")
+        )
+        if explicit_value > 0:
+            return explicit_value
+        return cls._first_positive_prompt_tokens_detail_value(usage, ("cached_tokens",))
+
+    @classmethod
+    def _get_cache_creation_input_tokens(cls, usage: Usage) -> int:
+        explicit_value = cls._first_positive_usage_value(
+            usage, ("cache_creation_input_tokens", "_cache_creation_input_tokens")
+        )
+        if explicit_value > 0:
+            return explicit_value
+        return cls._first_positive_prompt_tokens_detail_value(
+            usage, ("cache_creation_tokens", "cache_write_tokens")
+        )
+
+    @classmethod
+    def _translate_openai_usage_to_anthropic_usage_delta(
+        cls, usage: Usage
+    ) -> UsageDelta:
+        cache_read_input_tokens = cls._get_cache_read_input_tokens(usage)
+        cache_creation_input_tokens = cls._get_cache_creation_input_tokens(usage)
+        input_tokens = max(
+            (usage.prompt_tokens or 0)
+            - cache_read_input_tokens
+            - cache_creation_input_tokens,
+            0,
+        )
+
+        usage_delta = UsageDelta(
+            input_tokens=input_tokens,
+            output_tokens=usage.completion_tokens or 0,
+        )
+        if cache_creation_input_tokens > 0:
+            usage_delta["cache_creation_input_tokens"] = cache_creation_input_tokens
+        if cache_read_input_tokens > 0:
+            usage_delta["cache_read_input_tokens"] = cache_read_input_tokens
+        return usage_delta
+
+    @classmethod
+    def _translate_openai_usage_to_anthropic_usage(cls, usage: Usage) -> AnthropicUsage:
+        return cast(
+            AnthropicUsage,
+            cls._translate_openai_usage_to_anthropic_usage_delta(usage),
+        )
+
     def translate_openai_response_to_anthropic(
         self,
         response: ModelResponse,
@@ -1430,32 +1519,12 @@ class LiteLLMAnthropicMessagesAdapter:
         )
         # extract usage
         usage: Usage = getattr(response, "usage")
-        uncached_input_tokens = usage.prompt_tokens or 0
-        cached_tokens = 0
-        if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
-            cached_tokens = (
-                getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
-            )
-            uncached_input_tokens -= cached_tokens
-
-        anthropic_usage = AnthropicUsage(
-            input_tokens=uncached_input_tokens,
-            output_tokens=usage.completion_tokens or 0,
-        )
-        if (
-            hasattr(usage, "_cache_creation_input_tokens")
-            and usage._cache_creation_input_tokens > 0
-        ):
-            anthropic_usage["cache_creation_input_tokens"] = (
-                usage._cache_creation_input_tokens
-            )
-        if cached_tokens > 0:
-            anthropic_usage["cache_read_input_tokens"] = cached_tokens
+        anthropic_usage = self._translate_openai_usage_to_anthropic_usage(usage)
 
         if polyfill_result is not None and polyfill_result.iterations_usage is not None:
             message_iteration: UsageIteration = {
                 "type": "message",
-                "input_tokens": uncached_input_tokens,
+                "input_tokens": anthropic_usage["input_tokens"],
                 "output_tokens": usage.completion_tokens or 0,
             }
             anthropic_usage["iterations"] = list(polyfill_result.iterations_usage) + [
@@ -1644,35 +1713,9 @@ class LiteLLMAnthropicMessagesAdapter:
             else:
                 litellm_usage_chunk = None
             if litellm_usage_chunk is not None:
-                uncached_input_tokens = litellm_usage_chunk.prompt_tokens or 0
-                cached_tokens = 0
-                if (
-                    hasattr(litellm_usage_chunk, "prompt_tokens_details")
-                    and litellm_usage_chunk.prompt_tokens_details
-                ):
-                    cached_tokens = (
-                        getattr(
-                            litellm_usage_chunk.prompt_tokens_details,
-                            "cached_tokens",
-                            0,
-                        )
-                        or 0
-                    )
-                    uncached_input_tokens -= cached_tokens
-
-                usage_delta = UsageDelta(
-                    input_tokens=uncached_input_tokens,
-                    output_tokens=litellm_usage_chunk.completion_tokens or 0,
+                usage_delta = self._translate_openai_usage_to_anthropic_usage_delta(
+                    litellm_usage_chunk
                 )
-                if (
-                    hasattr(litellm_usage_chunk, "_cache_creation_input_tokens")
-                    and litellm_usage_chunk._cache_creation_input_tokens > 0
-                ):
-                    usage_delta["cache_creation_input_tokens"] = (
-                        litellm_usage_chunk._cache_creation_input_tokens
-                    )
-                if cached_tokens > 0:
-                    usage_delta["cache_read_input_tokens"] = cached_tokens
             else:
                 usage_delta = UsageDelta(input_tokens=0, output_tokens=0)
             message_block = MessageBlockDelta(

litellm/llms/anthropic/files/handler.py

@@ -84,7 +84,7 @@ class AnthropicFilesHandler:
 
         # Get Anthropic API credentials
         api_base = self.anthropic_model_info.get_api_base(api_base)
-        auth_header = self.anthropic_model_info.get_auth_header(api_key)
+        auth_header = self.anthropic_model_info.get_auth_header(api_key, api_base)
 
         if auth_header is None:
             raise ValueError("Missing Anthropic API Key")

litellm/llms/anthropic/files/transformation.py

@@ -95,7 +95,9 @@ class AnthropicFilesConfig(BaseFilesConfig):
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
     ) -> dict:
-        auth_header = AnthropicModelInfo.get_auth_header(api_key)
+        if api_base is None and isinstance(litellm_params, dict):
+            api_base = litellm_params.get("api_base")
+        auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base)
         if auth_header is None:
             raise ValueError(
                 "Anthropic API key is required. Set ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN environment variable or pass api_key parameter."

litellm/llms/anthropic/skills/transformation.py

@@ -38,10 +38,12 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
 
         # Get API key from litellm_params if available
         api_key = None
+        api_base = None
         if litellm_params is not None:
             api_key = litellm_params.api_key
+            api_base = litellm_params.api_base
 
-        auth_header = AnthropicModelInfo.get_auth_header(api_key)
+        auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base)
         if auth_header is None:
             raise ValueError(
                 "ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is required for Skills API"

litellm/llms/base_llm/google_genai/transformation.py

@@ -62,6 +62,19 @@ class BaseGoogleGenAIGenerateContentConfig(ABC):
             "get_supported_generate_content_optional_params is not implemented"
         )
 
+    def get_generate_content_request_top_level_fields(self) -> tuple[str, ...]:
+        """
+        Native Google ``GenerateContentRequest`` fields that sit at the top level
+        (siblings of ``generationConfig``) rather than inside it. The proxy forwards
+        these verbatim from a native request so ``generateContent`` is a drop-in for
+        Google's REST API.
+
+        Excludes ``contents``, ``model`` and ``tools`` (dedicated params),
+        ``systemInstruction`` (dedicated extraction) and ``generationConfig`` (mapped
+        to ``config``).
+        """
+        return ("safetySettings", "toolConfig", "cachedContent", "labels")
+
     @abstractmethod
     def map_generate_content_optional_params(
         self,

litellm/llms/bedrock/common_utils.py

@@ -904,9 +904,12 @@ class BedrockModelInfo(BaseLLMModelInfo):
             "mantle/": "mantle",
         }
 
-        # Check explicit routes first
+        # Check explicit routes first. Match each prefix only as a leading path
+        # segment so the `bedrock_mantle/` provider prefix is never mistaken for
+        # the `mantle/` invoke route (which would mangle
+        # `bedrock_mantle/openai.gpt-5.5` into `bedrock_openai.gpt-5.5`).
         for prefix, route_type in route_mappings.items():
-            if prefix in model:
+            if BedrockModelInfo._model_has_route_prefix(model, prefix):
                 return route_type
 
         # Check for nova spec prefixes (nova/ and nova-2/)
@@ -930,14 +933,14 @@ class BedrockModelInfo(BaseLLMModelInfo):
         """
         Check if the model is an explicit converse route.
         """
-        return "converse/" in model
+        return BedrockModelInfo._model_has_route_prefix(model, "converse/")
 
     @staticmethod
     def _explicit_claude_platform_route(model: str) -> bool:
         """
         Check if the model is an explicit Claude Platform on AWS route.
         """
-        return "claude_platform/" in model
+        return BedrockModelInfo._model_has_route_prefix(model, "claude_platform/")
 
     @staticmethod
     def get_claude_platform_model(model: str) -> str:
@@ -967,42 +970,58 @@ class BedrockModelInfo(BaseLLMModelInfo):
         """
         Check if the model is an explicit invoke route.
         """
-        return "invoke/" in model
+        return BedrockModelInfo._model_has_route_prefix(model, "invoke/")
 
     @staticmethod
     def _explicit_agent_route(model: str) -> bool:
         """
         Check if the model is an explicit agent route.
         """
-        return "agent/" in model
+        return BedrockModelInfo._model_has_route_prefix(model, "agent/")
 
     @staticmethod
     def _explicit_agentcore_route(model: str) -> bool:
         """
         Check if the model is an explicit agentcore route.
         """
-        return "agentcore/" in model
+        return BedrockModelInfo._model_has_route_prefix(model, "agentcore/")
+
+    @staticmethod
+    def _model_has_route_prefix(model: str, prefix: str) -> bool:
+        """Whether a route prefix (e.g. ``mantle/``) appears as a leading path segment.
+
+        A route token is only valid at the start of the model id or immediately
+        after a ``/``. A plain substring check matches the ``bedrock_mantle/``
+        provider prefix against the ``mantle/`` route, so the body model gets
+        mangled to ``bedrock_openai.gpt-5.5``; anchoring to a segment boundary
+        keeps the bare model id intact.
+
+        ``f"/{prefix}" in model`` matches the token as a segment at any path
+        depth, not just the second segment; that is intentional and acceptable
+        for these short, unambiguous route tokens.
+        """
+        return model.startswith(prefix) or f"/{prefix}" in model
 
     @staticmethod
     def _explicit_mantle_route(model: str) -> bool:
         """
         Check if the model is an explicit mantle route (bedrock-mantle endpoint).
         """
-        return "mantle/" in model
+        return BedrockModelInfo._model_has_route_prefix(model, "mantle/")
 
     @staticmethod
     def _explicit_converse_like_route(model: str) -> bool:
         """
         Check if the model is an explicit converse like route.
         """
-        return "converse_like/" in model
+        return BedrockModelInfo._model_has_route_prefix(model, "converse_like/")
 
     @staticmethod
     def _explicit_async_invoke_route(model: str) -> bool:
         """
         Check if the model is an explicit async invoke route.
         """
-        return "async_invoke/" in model
+        return BedrockModelInfo._model_has_route_prefix(model, "async_invoke/")
 
     @staticmethod
     def _explicit_openai_route(model: str) -> bool:
@@ -1010,7 +1029,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
         Check if the model is an explicit openai route.
         Used for Bedrock imported models that use OpenAI Chat Completions format.
         """
-        return "openai/" in model
+        return BedrockModelInfo._model_has_route_prefix(model, "openai/")
 
     @staticmethod
     def get_bedrock_provider_config_for_messages_api(

litellm/llms/vertex_ai/gemini/transformation.py

@@ -1046,9 +1046,9 @@ def _gemini_convert_messages_with_history(
                                 if invocation.get("tool_type"):
                                     tr_dict["toolType"] = invocation["tool_type"]
                                 tr_part: Dict[str, Any] = {"toolResponse": tr_dict}
-                                if "thought_signature" in invocation:
+                                if "response_thought_signature" in invocation:
                                     tr_part["thoughtSignature"] = invocation[
-                                        "thought_signature"
+                                        "response_thought_signature"
                                     ]
                                 assistant_content.append(tr_part)  # type: ignore
 

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -1634,13 +1634,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
             resp = tool_responses_by_id.pop(call_id, None)
             if resp is not None:
                 merged["response"] = resp.get("response")
-                # Keep response signature if call didn't have one
-                if "thought_signature" not in merged and "thought_signature" in resp:
-                    merged["thought_signature"] = resp["thought_signature"]
+                if "thought_signature" in resp:
+                    merged["response_thought_signature"] = resp["thought_signature"]
             invocations.append(merged)
 
         # Any orphan responses (shouldn't happen, but be safe)
         for resp_id, resp_entry in tool_responses_by_id.items():
+            if "thought_signature" in resp_entry:
+                resp_entry["response_thought_signature"] = resp_entry[
+                    "thought_signature"
+                ]
             invocations.append(resp_entry)
 
         return invocations if invocations else None
@@ -3598,16 +3601,22 @@ class ModelResponseIterator:
         chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or ""
         message = chunk.replace("\n\n", "")
 
-        # Accumulate JSON data
         self.accumulated_json += message
 
-        # Try to parse the accumulated JSON
+        # json.loads on the whole buffer after every fragment is O(n^2) and
+        # holds the GIL, freezing the event loop for seconds on large responses
+        # (https://github.com/BerriAI/litellm/issues/26181). A complete Gemini
+        # chunk is a JSON object/array, so only attempt the parse once the
+        # buffer's last non-whitespace byte can close one.
+        stripped = self.accumulated_json.rstrip()
+        if not stripped or stripped[-1] not in "}]":
+            return None
+
         try:
             _data = json.loads(self.accumulated_json)
             self.accumulated_json = ""  # reset after successful parsing
             return self.chunk_parser(chunk=_data)
         except json.JSONDecodeError:
-            # If it's not valid JSON yet, continue to the next event
             return None
 
     def _common_chunk_parsing_logic(

litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py

@@ -274,6 +274,7 @@ class GoogleBatchEmbeddings(VertexLLM):
                 model_response=model_response,
                 model=model,
                 response_json=_json_response,
+                resolved_files=resolved_files,
             )
         else:
             _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response)  # type: ignore
@@ -377,6 +378,7 @@ class GoogleBatchEmbeddings(VertexLLM):
                 model_response=model_response,
                 model=model,
                 response_json=_json_response,
+                resolved_files=resolved_files,
             )
         else:
             _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response)  # type: ignore

litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py

@@ -4,7 +4,10 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc
 Why separate file? Make it easy to see how transformation works
 """
 
-from typing import Dict, List, Optional, Tuple
+from collections.abc import Mapping
+from typing import Dict, List, Optional, Sequence, Tuple
+
+from pydantic import TypeAdapter, ValidationError
 
 from litellm.types.llms.vertex_ai import (
     BlobType,
@@ -13,10 +16,17 @@ from litellm.types.llms.vertex_ai import (
     FileDataType,
     GeminiEmbeddingInput,
     PartType,
+    PromptTokensDetails,
+    UsageMetadata,
     VertexAIBatchEmbeddingsRequestBody,
     VertexAIBatchEmbeddingsResponseObject,
 )
-from litellm.types.utils import Embedding, EmbeddingResponse, Usage
+from litellm.types.utils import (
+    Embedding,
+    EmbeddingResponse,
+    PromptTokensDetailsWrapper,
+    Usage,
+)
 from litellm.utils import get_formatted_prompt, token_counter
 
 SUPPORTED_EMBEDDING_MIME_TYPES = {
@@ -294,11 +304,133 @@ def transform_openai_input_gemini_embed_content(
     return request_body
 
 
+_IMAGE_MIME_TYPES = frozenset({"image/png", "image/jpeg"})
+_VIDEO_TOKENS_PER_SECOND = 258.0
+_AUDIO_TOKENS_PER_SECOND = 32.0
+_usage_metadata_adapter = TypeAdapter(UsageMetadata)
+
+
+def _parse_usage_metadata(raw_usage_metadata: object) -> Optional[UsageMetadata]:
+    if not isinstance(raw_usage_metadata, dict):
+        return None
+    try:
+        return _usage_metadata_adapter.validate_python(raw_usage_metadata)
+    except ValidationError:
+        return None
+
+
+def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]:
+    if isinstance(input, str):
+        return (input,)
+    return tuple(
+        sub
+        for element in input
+        for sub in (element if isinstance(element, list) else [element])
+    )
+
+
+def _is_image_element(
+    element: str,
+    resolved_files: Mapping[str, Mapping[str, str]],
+) -> bool:
+    if element.startswith("data:") and ";base64," in element:
+        try:
+            mime_type, _ = _parse_data_url(element)
+        except ValueError:
+            return False
+        return mime_type in _IMAGE_MIME_TYPES
+    if _is_gcs_url(element):
+        try:
+            return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES
+        except ValueError:
+            return False
+    if _is_file_reference(element):
+        file_info = resolved_files.get(element)
+        return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES
+    return False
+
+
+def _count_input_images(
+    input: GeminiEmbeddingInput,
+    resolved_files: Mapping[str, Mapping[str, str]],
+) -> int:
+    return sum(
+        1
+        for element in _flatten_input(input)
+        if _is_image_element(element, resolved_files)
+    )
+
+
+def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int:
+    return sum(
+        detail["tokenCount"] for detail in details if detail["modality"] == modality
+    )
+
+
+def _fallback_usage(input: GeminiEmbeddingInput, model: str) -> Usage:
+    if _is_multimodal_input(input):
+        return Usage(prompt_tokens=0, total_tokens=0)
+    input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
+    prompt_tokens = token_counter(model=model, text=input_text)
+    return Usage(prompt_tokens=prompt_tokens, total_tokens=prompt_tokens)
+
+
+def _usage_from_embed_content_response(
+    input: GeminiEmbeddingInput,
+    model: str,
+    raw_usage_metadata: object,
+    resolved_files: Mapping[str, Mapping[str, str]],
+) -> Usage:
+    usage_metadata = _parse_usage_metadata(raw_usage_metadata)
+    if usage_metadata is None:
+        return _fallback_usage(input, model)
+
+    prompt_tokens = usage_metadata.get("promptTokenCount", 0)
+    total_tokens = usage_metadata.get("totalTokenCount") or prompt_tokens
+
+    details: Sequence[PromptTokensDetails] = (
+        usage_metadata.get("promptTokensDetails") or ()
+    )
+    text_tokens = _tokens_for_modality(details, "TEXT")
+    audio_tokens = _tokens_for_modality(details, "AUDIO")
+    video_tokens = _tokens_for_modality(details, "VIDEO")
+    image_count = _count_input_images(input, resolved_files)
+
+    video_length_seconds = (
+        video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0
+    )
+    audio_length_seconds = (
+        audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0
+    )
+
+    # generic_cost_per_token rewrites text_tokens to the full prompt minus
+    # other modalities when both text_tokens and image_count are zero. For
+    # video, that misallocates video tokens to text; a 1-token floor sidesteps
+    # the rewrite and keeps billing on input_cost_per_video_per_second.
+    needs_video_text_floor = (
+        video_length_seconds > 0 and text_tokens == 0 and image_count == 0
+    )
+    resolved_text_tokens = 1 if needs_video_text_floor else text_tokens
+
+    return Usage(
+        prompt_tokens=prompt_tokens,
+        total_tokens=total_tokens,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            text_tokens=resolved_text_tokens,
+            audio_tokens=audio_tokens,
+            image_count=image_count,
+            video_length_seconds=video_length_seconds,
+            audio_length_seconds=audio_length_seconds,
+        ),
+    )
+
+
 def process_embed_content_response(
     input: GeminiEmbeddingInput,
     model_response: EmbeddingResponse,
     model: str,
     response_json: dict,
+    resolved_files: Mapping[str, Mapping[str, str]] | None = None,
 ) -> EmbeddingResponse:
     """
     Process Gemini embedContent response (single embedding for multimodal input).
@@ -308,6 +440,8 @@ def process_embed_content_response(
         model_response: EmbeddingResponse to populate
         model: Model name
         response_json: Raw JSON response from embedContent endpoint
+        resolved_files: Mapping of file references (files/abc) to {mime_type, uri},
+            used to bill resolved image references at the per-image rate
 
     Returns:
         EmbeddingResponse with single embedding
@@ -327,14 +461,11 @@ def process_embed_content_response(
 
     model_response.data = [openai_embedding]
     model_response.model = model
-
-    if _is_multimodal_input(input):
-        prompt_tokens = 0
-    else:
-        input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
-        prompt_tokens = token_counter(model=model, text=input_text)
-    model_response.usage = Usage(
-        prompt_tokens=prompt_tokens, total_tokens=prompt_tokens
+    model_response.usage = _usage_from_embed_content_response(
+        input=input,
+        model=model,
+        raw_usage_metadata=response_json.get("usageMetadata"),
+        resolved_files=resolved_files or {},
     )
 
     return model_response

litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py

@@ -35,27 +35,21 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
 
         Validate the environment for the request
         """
+        # Work on a local copy — router shallow-copies litellm_params so the caller's
+        # headers dict may be the shared deployment extra_headers object.
+        headers = dict(headers)
         vertex_ai_project = VertexBase.safe_get_vertex_ai_project(litellm_params)
         vertex_ai_location = VertexBase.safe_get_vertex_ai_location(litellm_params)
 
-        project_id: Optional[str] = None
-        if "Authorization" not in headers:
-            vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(
-                litellm_params
-            )
+        vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params)
+        access_token, project_id = self._ensure_access_token(
+            credentials=vertex_credentials,
+            project_id=vertex_ai_project,
+            custom_llm_provider="vertex_ai",
+        )
+        headers["Authorization"] = f"Bearer {access_token}"
 
-            access_token, project_id = self._ensure_access_token(
-                credentials=vertex_credentials,
-                project_id=vertex_ai_project,
-                custom_llm_provider="vertex_ai",
-            )
-
-            headers["Authorization"] = f"Bearer {access_token}"
-        else:
-            # Authorization already in headers, but we still need project_id
-            project_id = vertex_ai_project
-
-        # Always calculate api_base if not provided, regardless of Authorization header
+        # Calculate api_base if not provided
         if api_base is None:
             api_base = self.get_complete_vertex_url(
                 custom_api_base=api_base,

litellm/llms/vertex_ai/vertex_ai_partner_models/main.py

@@ -194,9 +194,11 @@ class VertexAIPartnerModels(VertexBase):
                     encoding=encoding,
                 )
             elif "claude" in model:
-                if headers is None:
-                    headers = {}
-                headers.update({"Authorization": "Bearer {}".format(access_token)})
+                # Build a new dict so we never mutate the shared deployment extra_headers object.
+                headers = {
+                    **(headers or {}),
+                    "Authorization": "Bearer {}".format(access_token),
+                }
 
                 optional_params.update(
                     {

litellm/main.py

@@ -633,6 +633,7 @@ async def acompletion(
 
     try:
         # Use a partial function to pass your keyword arguments
+        kwargs.pop("acompletion", None)
         func = partial(completion, **completion_kwargs, **kwargs)
 
         # Add the context to the function
@@ -5063,6 +5064,12 @@ def completion(  # type: ignore
     ######### unpacking kwargs #####################
     args = locals()
 
+    # Set by the responses->completion fallback so completion() does not bridge
+    # back to the Responses API: that round-trip mutually recurses forever for a
+    # model whose model_cost mode is "responses" but whose provider has no
+    # Responses API config (get_provider_responses_api_config -> None).
+    skip_responses_api_bridge = kwargs.pop("_skip_responses_api_bridge", False)
+
     skip_mcp_handler = kwargs.pop("_skip_mcp_handler", False)
     if not skip_mcp_handler and tools:
         from litellm.responses.mcp.chat_completions_handler import acompletion_with_mcp
@@ -5358,7 +5365,7 @@ def completion(  # type: ignore
 
         messages = update_messages_with_model_file_ids(
             messages=messages,
-            model_id=kwargs.get("model_info", {}).get("id", None),
+            model_id=(kwargs.get("model_info") or {}).get("id", None),
             model_file_id_mapping=cast(
                 Dict[str, Dict[str, str]],
                 kwargs.get("model_file_id_mapping") or {},
@@ -5560,7 +5567,10 @@ def completion(  # type: ignore
         # detection when the deployment name differs from the model name.
         _azure_detection_model = base_model or model
 
-        if responses_api_model_info.get("mode") == "responses":
+        if (
+            responses_api_model_info.get("mode") == "responses"
+            and not skip_responses_api_bridge
+        ):
             from litellm.completion_extras import responses_api_bridge
 
             optional_params, rs_val = (

litellm/model_prices_and_context_window_backup.json

@@ -25793,7 +25793,7 @@
         "supports_response_schema": true,
         "supports_tool_choice": true
     },
-    "mistral/mistral-medium-latest": {
+    "mistral/mistral-medium-2508": {
         "input_cost_per_token": 4e-07,
         "litellm_provider": "mistral",
         "max_input_tokens": 131072,
@@ -25801,12 +25801,45 @@
         "max_tokens": 131072,
         "mode": "chat",
         "output_cost_per_token": 2e-06,
+        "source": "https://mistral.ai/news/mistral-medium-3",
         "supports_assistant_prefill": true,
         "supports_function_calling": true,
         "supports_response_schema": true,
         "supports_tool_choice": true,
         "supports_vision": true
     },
+    "mistral/mistral-medium-2604": {
+        "input_cost_per_token": 1.5e-06,
+        "litellm_provider": "mistral",
+        "max_input_tokens": 262144,
+        "max_output_tokens": 262144,
+        "max_tokens": 262144,
+        "mode": "chat",
+        "output_cost_per_token": 7.5e-06,
+        "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
+        "supports_assistant_prefill": true,
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "mistral/mistral-medium-latest": {
+        "input_cost_per_token": 1.5e-06,
+        "litellm_provider": "mistral",
+        "max_input_tokens": 262144,
+        "max_output_tokens": 262144,
+        "max_tokens": 262144,
+        "mode": "chat",
+        "output_cost_per_token": 7.5e-06,
+        "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
+        "supports_assistant_prefill": true,
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "mistral/mistral-medium-3-1-2508": {
         "input_cost_per_token": 4e-07,
         "litellm_provider": "mistral",
@@ -25833,6 +25866,7 @@
         "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
         "supports_assistant_prefill": true,
         "supports_function_calling": true,
+        "supports_reasoning": true,
         "supports_response_schema": true,
         "supports_tool_choice": true,
         "supports_vision": true

litellm/ocr/main.py

@@ -19,13 +19,7 @@ from litellm.constants import request_timeout
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
-from litellm.ocr.rust_bridge import (
-    RustAocr,
-    RustOcr,
-    load_rust_aocr,
-    load_rust_ocr,
-    rust_ocr_enabled,
-)
+from litellm.rust_bridge import ocr as rust_ocr_bridge
 from litellm.types.router import GenericLiteLLMParams
 from litellm.utils import ProviderConfigManager, client
 
@@ -60,27 +54,10 @@ class _PreparedRustOCRCall:
 _RUST_OCR_PROVIDERS = {
     "mistral",
     "azure_ai",
-    "azure_ai/doc-intelligence",
     "vertex_ai",
 }
 
 
-def _timeout_to_seconds(
-    timeout: Union[float, httpx.Timeout] | None,
-) -> float | None:
-    """Convert the Python OCR timeout to a single seconds value for the Rust bridge.
-
-    The Rust HTTP client takes one duration; ``httpx.Timeout`` carries separate
-    connect/read/write/pool values, so pick the read deadline as the closest
-    analog to a total-request timeout.
-    """
-    if timeout is None:
-        return None
-    if isinstance(timeout, httpx.Timeout):
-        return timeout.read
-    return float(timeout)
-
-
 def _prepare_ocr_request(
     model: str,
     document: dict[str, Any],
@@ -218,13 +195,9 @@ def _rust_bridge_api_base(
 ) -> str | None:
     if prepared_request.api_base is not None:
         return prepared_request.api_base
-    if prepared_request.custom_llm_provider == "azure_ai/doc-intelligence":
-        return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
     if prepared_request.custom_llm_provider == "azure_ai":
-        if (
-            "doc-intelligence" in prepared_request.model
-            or "documentintelligence" in prepared_request.model
-        ):
+        model = prepared_request.model.lower()
+        if "doc-intelligence" in model or "documentintelligence" in model:
             return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
         return resolve_secret("AZURE_AI_API_BASE")
     return None
@@ -278,57 +251,53 @@ def _prepare_rust_ocr_call(
 
 
 def _run_rust_ocr(
-    rust_ocr: RustOcr,
     prepared_request: _PreparedOCRRequest,
     resolve_api_key: Callable[[str], str | None],
-) -> OCRResponse:
-    """Run the Mistral OCR call through the Rust bridge and wrap the result.
-
-    Resolves the key the same way the Python path does so secret-manager backends
-    (AWS/Azure/GCP/Vault) work; the Rust bridge's own fallback only reads the
-    process environment. The request that Rust actually sends (resolved URL and
-    headers) is mirrored into pre_call so logs match the wire. Dependencies are
-    injected so this stays unit-testable without patching module globals.
-    """
+) -> OCRResponse | None:
+    if rust_ocr_bridge.load_rust_ocr() is None:
+        return None
     prepared = _prepare_rust_ocr_call(
         prepared_request=prepared_request,
         resolve_api_key=resolve_api_key,
     )
-    return OCRResponse.model_validate(
-        rust_ocr(
-            model=prepared_request.model,
-            document=cast(dict[str, object], prepared_request.document),
-            api_key=prepared.api_key,
-            api_base=prepared.api_base,
-            custom_llm_provider=prepared_request.custom_llm_provider,
-            extra_headers=prepared.headers,
-            optional_params=prepared.optional_params,
-            timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout),
-        )
+    rust_response = rust_ocr_bridge.ocr(
+        model=prepared_request.model,
+        document=prepared_request.document,
+        api_key=prepared.api_key,
+        api_base=prepared.api_base,
+        custom_llm_provider=prepared_request.custom_llm_provider,
+        extra_headers=prepared.headers,
+        optional_params=prepared.optional_params,
+        timeout=prepared_request.effective_timeout,
     )
+    if rust_response is None:
+        return None
+    return OCRResponse.model_validate(rust_response)
 
 
 async def _run_rust_aocr(
-    rust_aocr: RustAocr,
     prepared_request: _PreparedOCRRequest,
     resolve_api_key: Callable[[str], str | None],
-) -> OCRResponse:
+) -> OCRResponse | None:
+    if rust_ocr_bridge.load_rust_aocr() is None:
+        return None
     prepared = _prepare_rust_ocr_call(
         prepared_request=prepared_request,
         resolve_api_key=resolve_api_key,
     )
-    return OCRResponse.model_validate(
-        await rust_aocr(
-            model=prepared_request.model,
-            document=cast(dict[str, object], prepared_request.document),
-            api_key=prepared.api_key,
-            api_base=prepared.api_base,
-            custom_llm_provider=prepared_request.custom_llm_provider,
-            extra_headers=prepared.headers,
-            optional_params=prepared.optional_params,
-            timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout),
-        )
+    rust_response = await rust_ocr_bridge.aocr(
+        model=prepared_request.model,
+        document=prepared_request.document,
+        api_key=prepared.api_key,
+        api_base=prepared.api_base,
+        custom_llm_provider=prepared_request.custom_llm_provider,
+        extra_headers=prepared.headers,
+        optional_params=prepared.optional_params,
+        timeout=prepared_request.effective_timeout,
     )
+    if rust_response is None:
+        return None
+    return OCRResponse.model_validate(rust_response)
 
 
 @client
@@ -427,21 +396,19 @@ async def aocr(
             {"model": model, "custom_llm_provider": custom_llm_provider}
         )
 
-        if _rust_ocr_supported(prepared) and rust_ocr_enabled():
-            rust_aocr = load_rust_aocr()
-            if rust_aocr is None:
+        if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled():
+            from litellm.secret_managers.main import get_secret_str
+
+            rust_response = await _run_rust_aocr(
+                prepared_request=prepared,
+                resolve_api_key=get_secret_str,
+            )
+            if rust_response is None:
                 verbose_logger.debug(
                     "Async Rust OCR bridge unavailable; falling back to Python path"
                 )
             else:
-                from litellm.secret_managers.main import get_secret_str
-
-                response = await _run_rust_aocr(
-                    rust_aocr=rust_aocr,
-                    prepared_request=prepared,
-                    resolve_api_key=get_secret_str,
-                )
-                return response
+                return rust_response
 
         response = base_llm_http_handler.ocr(
             model=prepared.model,
@@ -704,21 +671,19 @@ def ocr(
             {"model": model, "custom_llm_provider": custom_llm_provider}
         )
 
-        # Optional Rust path: hand supported OCR calls to the Rust bridge.
-        if _rust_ocr_supported(prepared) and rust_ocr_enabled():
-            rust_ocr = load_rust_ocr()
-            if rust_ocr is None:
+        if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled():
+            from litellm.secret_managers.main import get_secret_str
+
+            rust_response = _run_rust_ocr(
+                prepared_request=prepared,
+                resolve_api_key=get_secret_str,
+            )
+            if rust_response is None:
                 verbose_logger.debug(
                     "Rust OCR bridge unavailable; falling back to Python path"
                 )
             else:
-                from litellm.secret_managers.main import get_secret_str
-
-                return _run_rust_ocr(
-                    rust_ocr=rust_ocr,
-                    prepared_request=prepared,
-                    resolve_api_key=get_secret_str,
-                )
+                return rust_response
 
         response = base_llm_http_handler.ocr(
             model=prepared.model,

litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py

@@ -624,7 +624,9 @@ class MCPRequestHandler:
 
         Permission hierarchy (all rules are intersections):
         1. Get allowed servers from key permissions
-        2. Get allowed servers from team permissions (key inherits from team, or intersection)
+        2. Get allowed servers from team permissions (key inherits from team, or
+           intersection; or inherits nothing when require_key_mcp_access_defined
+           is enabled, making the team a ceiling rather than a default)
         3. Get allowed servers from end_user permissions (intersected if set)
         4. Get allowed servers from agent permissions (intersected if set)
         5. Get allowed servers from org permissions — org acts as a ceiling: if the org
@@ -677,7 +679,16 @@ class MCPRequestHandler:
             if not team_set:
                 base = key_set  # no team restriction
             elif not key_set:
-                base = team_set  # key has no own perms → inherits team
+                # A key that grants no MCP servers of its own inherits the
+                # team's by default. With require_key_mcp_access_defined the
+                # team is a ceiling rather than a default, so the key must
+                # grant servers explicitly (or via an access group) to reach
+                # any — it inherits none.
+                base = (
+                    set()
+                    if general_settings.get("require_key_mcp_access_defined", False)
+                    else team_set
+                )
             else:
                 base = key_set & team_set  # both restrict → intersect
 

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -79,6 +79,7 @@ if MCP_AVAILABLE:
     )
     from litellm.proxy._experimental.mcp_server.server import (
         ListMCPToolsRestAPIResponseObject,
+        MCPInfo,
         MCPServer,
         _tool_name_matches,
         execute_mcp_tool,
@@ -238,14 +239,24 @@ if MCP_AVAILABLE:
             )
             return {}
 
-    def _create_tool_response_objects(tools, server_mcp_info):
-        """Helper function to create tool response objects."""
+    def _create_tool_response_objects(tools, server: MCPServer):
+        """Helper function to create tool response objects.
+
+        Enriches the server's ``mcp_info`` with ``server_id`` and ``alias`` so
+        REST clients can map the internal ``server_name`` to the user-facing
+        alias without needing access to the ``mcp_routes``-gated server listing.
+        """
+        enriched_mcp_info: MCPInfo = {
+            **(server.mcp_info or {}),
+            "server_id": server.server_id,
+            "alias": server.alias,
+        }
         return [
             ListMCPToolsRestAPIResponseObject(
                 name=tool.name,
                 description=tool.description,
                 inputSchema=tool.inputSchema,
-                mcp_info=server_mcp_info,
+                mcp_info=enriched_mcp_info,
             )
             for tool in tools
         ]
@@ -405,7 +416,7 @@ if MCP_AVAILABLE:
         )
 
         if not apply_tool_filters:
-            return _create_tool_response_objects(tools, server.mcp_info)
+            return _create_tool_response_objects(tools, server)
 
         # Always apply allowed_tools/disallowed_tools so the blacklist is
         # enforced even when no allowlist is set (matches the SSE/HTTP path).
@@ -436,7 +447,7 @@ if MCP_AVAILABLE:
                     if _tool_name_matches(tool.name, allowed_tools_for_server)
                 ]
 
-        return _create_tool_response_objects(tools, server.mcp_info)
+        return _create_tool_response_objects(tools, server)
 
     async def _resolve_allowed_mcp_servers_for_tool_call(
         user_api_key_dict: UserAPIKeyAuth,
@@ -587,6 +598,8 @@ if MCP_AVAILABLE:
                     "mcp_info": {
                         "server_name": "zapier",
                         "logo_url": "https://www.zapier.com/logo.png",
+                        "server_id": "a1b2c3d4-...",
+                        "alias": "zapier_prod",
                     }
                 }
             ],

litellm/proxy/_types.py

@@ -3181,6 +3181,7 @@ class SpendLogsMetadata(TypedDict):
         dict
     ]  # special param to log k,v pairs to spendlogs for a call
     requester_ip_address: Optional[str]
+    litellm_call_id: Optional[str]
     applied_guardrails: Optional[List[str]]
     mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall]
     vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]]

litellm/proxy/client/exceptions.py

@@ -2,18 +2,30 @@ from typing import Union
 
 import requests
 
+from litellm.litellm_core_utils.secret_redaction import redact_string
+
+
+def _redact_orig_exception(
+    orig_exception: Union[requests.exceptions.HTTPError, str],
+) -> Union[requests.exceptions.HTTPError, str]:
+    if isinstance(orig_exception, requests.exceptions.HTTPError):
+        return requests.exceptions.HTTPError(
+            redact_string(str(orig_exception)), response=orig_exception.response
+        )
+    return redact_string(str(orig_exception))
+
 
 class UnauthorizedError(Exception):
     """Exception raised when the API returns a 401 Unauthorized response."""
 
     def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]):
-        self.orig_exception = orig_exception
-        super().__init__(str(orig_exception))
+        self.orig_exception = _redact_orig_exception(orig_exception)
+        super().__init__(str(self.orig_exception))
 
 
 class NotFoundError(Exception):
     """Exception raised when the API returns a 404 Not Found response or indicates a resource was not found."""
 
     def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]):
-        self.orig_exception = orig_exception
-        super().__init__(str(orig_exception))
+        self.orig_exception = _redact_orig_exception(orig_exception)
+        super().__init__(str(self.orig_exception))

litellm/proxy/client/keys.py

@@ -2,6 +2,8 @@ from typing import Any, Dict, List, Optional, Union
 
 import requests
 
+from litellm.litellm_core_utils.secret_redaction import redact_string
+
 from .exceptions import UnauthorizedError
 
 
@@ -314,6 +316,9 @@ class KeysManagementClient:
             response.raise_for_status()
             return response.json()
         except requests.exceptions.HTTPError as e:
+            redacted_message = redact_string(str(e))
             if e.response.status_code == 401:
-                raise UnauthorizedError(e)
-            raise
+                raise UnauthorizedError(e) from None
+            raise requests.exceptions.HTTPError(
+                redacted_message, response=e.response
+            ) from None

litellm/proxy/common_request_processing.py

@@ -516,12 +516,13 @@ def _override_openai_response_model(
     LiteLLM internally prefixes some provider/deployment model identifiers (e.g. `hosted_vllm/...`).
     That internal identifier should not be returned to clients in the OpenAI `model` field.
 
-    Note: This is intentionally verbose. A model mismatch is a useful signal that an internal
-    model identifier is being stamped/preserved somewhere in the request/response pipeline.
-    We log mismatches as warnings (and then restamp to the client-requested value) so these
-    paths stay observable for maintainers/operators without breaking client compatibility.
+    Note: This is intentionally verbose at debug level. A model mismatch is a useful signal that an
+    internal model identifier is being stamped/preserved somewhere in the request/response pipeline.
+    We log mismatches as debug (and then restamp to the client-requested value) so these paths stay
+    observable for maintainers without breaking client compatibility or alarming operators.
 
-    Errors are reserved for cases where the proxy cannot read/override the response model field.
+    Responses that omit an OpenAI-style `model` field are left unchanged (silent return),
+    including dict responses with no `model` key.
 
     Exceptions:
     1. If a fallback occurred (indicated by x-litellm-attempted-fallbacks header),
@@ -577,6 +578,8 @@ def _override_openai_response_model(
         return
 
     if isinstance(response_obj, dict):
+        if "model" not in response_obj:
+            return
         downstream_model = response_obj.get("model")
         if downstream_model != requested_model:
             verbose_proxy_logger.debug(
@@ -589,11 +592,6 @@ def _override_openai_response_model(
         return
 
     if not hasattr(response_obj, "model"):
-        verbose_proxy_logger.error(
-            "%s: cannot override response model; missing `model` attribute. response_type=%s",
-            log_context,
-            type(response_obj),
-        )
         return
 
     downstream_model = getattr(response_obj, "model", None)
@@ -608,7 +606,7 @@ def _override_openai_response_model(
     try:
         setattr(response_obj, "model", requested_model)
     except Exception as e:
-        verbose_proxy_logger.error(
+        verbose_proxy_logger.debug(
             "%s: failed to override response.model=%r on response_type=%s. error=%s",
             log_context,
             requested_model,

litellm/proxy/health_endpoints/_health_endpoints.py

@@ -1849,7 +1849,10 @@ async def test_model_connection(
             "responses",
             "ocr",
         ]
-    ] = fastapi.Body("chat", description="The mode to test the model with"),
+    ] = fastapi.Body(
+        None,
+        description="The mode to test the model with. If not provided, auto-detected from model capabilities.",
+    ),
     litellm_params: Dict = fastapi.Body(
         None,
         description="Parameters for litellm.completion, litellm.embedding for the health check",

litellm/proxy/management_endpoints/common_utils.py

@@ -1,8 +1,27 @@
+import math
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 from fastapi import HTTPException, status
 from pydantic import BaseModel
 
+
+# Defined above the `litellm.proxy.*` imports so the name is bound even when
+# this module is imported first through the proxy import cycle (CodeQL:
+# module-level cyclic import). Depends only on `math` + `HTTPException`.
+def validate_finite_spend(spend: float | None) -> None:
+    """Reject NaN/±inf spend before it reaches the DB / spend counter.
+
+    A non-finite spend would otherwise slip past `spend >= max_budget`
+    enforcement, since any comparison with NaN (and `-inf >= max_budget`)
+    is False, letting the entity keep spending past its configured budget.
+    """
+    if spend is not None and not math.isfinite(spend):
+        raise HTTPException(
+            status_code=400,
+            detail={"error": f"spend must be a finite number. Received: {spend}"},
+        )
+
+
 from litellm._logging import verbose_proxy_logger
 from litellm.caching import DualCache
 from litellm.proxy._types import (

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -36,6 +36,7 @@ from litellm.proxy.management_endpoints.common_utils import (
     _is_user_team_admin,
     _user_has_admin_view,
     require_caller_user_id_for_non_admin,
+    validate_finite_spend,
 )
 from litellm.proxy.management_endpoints.key_management_endpoints import (
     generate_key_helper_fn,
@@ -1256,6 +1257,25 @@ def _check_user_update_authz(
         )
 
 
+async def _invalidate_user_spend_counter_if_changed(
+    non_default_values: dict[str, Any],
+) -> None:
+    """Invalidate the cross-pod spend counter after a direct ``spend`` change.
+
+    A direct ``spend`` change must also invalidate the cross-pod spend counter
+    enforcement reads; the DB write alone leaves a warm counter at the stale
+    value. ``non_default_values["user_id"]`` is populated in every branch of the
+    caller (incl. the email-new-user insert path, whose response is a bare model
+    and not safely subscriptable).
+    """
+    if non_default_values.get("spend") is not None:
+        from litellm.proxy.proxy_server import _invalidate_spend_counter
+
+        await _invalidate_spend_counter(
+            counter_key=f"spend:user:{non_default_values['user_id']}"
+        )
+
+
 async def _update_single_user_helper(
     user_request: UpdateUserRequest,
     user_api_key_dict: UserAPIKeyAuth,
@@ -1336,6 +1356,9 @@ async def _update_single_user_helper(
         existing_metadata=existing_metadata or {},
     )
 
+    # Reject NaN/±inf spend before it can reach the DB / spend counter.
+    validate_finite_spend(non_default_values.get("spend"))
+
     # Perform the update
     response: Optional[Dict[str, Any]] = None
 
@@ -1384,6 +1407,8 @@ async def _update_single_user_helper(
             litellm_proxy_admin_name=litellm_proxy_admin_name,
         )
 
+        await _invalidate_user_spend_counter_if_changed(non_default_values)
+
     if response is None:
         raise HTTPException(
             status_code=400,

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -69,6 +69,7 @@ from litellm.proxy.management_endpoints.common_utils import (
     _is_user_team_admin,
     _set_object_metadata_field,
     _team_member_has_permission,
+    validate_finite_spend,
 )
 from litellm.proxy.management_endpoints.model_management_endpoints import (
     _add_model_to_db,
@@ -2210,6 +2211,9 @@ async def _validate_update_key_data(
     user_api_key_cache: Any,
 ) -> None:
     """Validate permissions and constraints for key update."""
+    # Reject NaN/±inf spend before it can reach the DB / spend counter.
+    validate_finite_spend(data.spend)
+
     _is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
 
     _check_allowed_routes_caller_permission(
@@ -2269,12 +2273,14 @@ async def _validate_update_key_data(
     #   existing admin-only budget semantics).  budget_limits uses
     #   model_fields_set because an explicit null/[] clears the field
     #   and must gate the same as setting or changing it.
+    # - spend gates on presence alone (not a value diff): the DB spend
+    #   lags the live cross-pod counter, so letting an "unchanged" spend
+    #   through the non-admin path would let a key owner / team member
+    #   overwrite the live counter below real usage and silently weaken
+    #   enforcement.
     _is_budget_change = (
         (data.max_budget is not None and data.max_budget != existing_key_row.max_budget)
-        or (
-            data.spend is not None
-            and data.spend != getattr(existing_key_row, "spend", None)
-        )
+        or data.spend is not None
         or "budget_limits" in data.model_fields_set
     )
 
@@ -2609,15 +2615,24 @@ async def update_key_fn(
         )
 
         if data.spend is not None:
-            try:
-                from litellm.proxy.proxy_server import _invalidate_spend_counter
+            from litellm.proxy.proxy_server import spend_counter_cache
 
-                token_to_invalidate = _hash_token_if_needed(key)
-                await _invalidate_spend_counter(
-                    counter_key=f"spend:key:{token_to_invalidate}"
-                )
-            except Exception:
-                pass
+            counter_key = f"spend:key:{_hash_token_if_needed(key)}"
+            spend_counter_cache.in_memory_cache.set_cache(
+                key=counter_key, value=data.spend, ttl=60
+            )
+            if spend_counter_cache.redis_cache is not None:
+                try:
+                    await spend_counter_cache.redis_cache.async_set_cache(
+                        key=counter_key, value=data.spend, ttl=60
+                    )
+                except Exception as redis_err:
+                    verbose_proxy_logger.warning(
+                        "Failed to update spend counter %s in Redis after key spend update: %s. "
+                        "Budget checks may use stale value until counter expires.",
+                        counter_key,
+                        redis_err,
+                    )
 
         asyncio.create_task(
             KeyManagementEventHooks.async_key_updated_hook(

litellm/proxy/proxy_server.py

@@ -1114,6 +1114,33 @@ _OPENAPI_HTTP_METHODS = {
 # `_SSO_SENSITIVE_FIELDS` / `_CACHE_SENSITIVE_FIELDS` constants in the SSO
 # and cache endpoint files.
 _ALERTING_SENSITIVE_VARS: Set[str] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"}
+_DB_LITELLM_PARAM_ENV_REF_KEYS = frozenset(
+    {
+        "api_key",
+        "client_secret",
+        "vertex_credentials",
+        "vertex_ai_credentials",
+        "aws_access_key_id",
+        "aws_secret_access_key",
+    }
+)
+
+
+def _db_model_is_team_scoped(model: object) -> bool:
+    model_info = getattr(model, "model_info", None)
+    if isinstance(model_info, BaseModel):
+        return getattr(model_info, "team_id", None) is not None
+    if isinstance(model_info, str):
+        try:
+            model_info = json.loads(model_info)
+        except (TypeError, ValueError):
+            model_info = None
+    if isinstance(model_info, dict) and model_info.get("team_id") is not None:
+        return True
+    if getattr(model_info, "team_id", None) is not None:
+        return True
+    model_name = getattr(model, "model_name", None)
+    return isinstance(model_name, str) and model_name.startswith("model_name_")
 
 
 def _strip_operation_id_method_suffix(operation_id: str) -> str:
@@ -5255,6 +5282,24 @@ class ProxyConfig:
                     deleted_deployments += 1
         return deleted_deployments
 
+    def _resolve_db_litellm_param(
+        self, key: str, value: object, resolve_env_refs: bool = True
+    ) -> object:
+        if not isinstance(value, str):
+            return value
+
+        decrypted_value = decrypt_value_helper(
+            value=value, key=key, return_original_value=True
+        )
+        if (
+            resolve_env_refs
+            and key in _DB_LITELLM_PARAM_ENV_REF_KEYS
+            and isinstance(decrypted_value, str)
+            and decrypted_value.startswith("os.environ/")
+        ):
+            return get_secret(decrypted_value)
+        return decrypted_value
+
     def _add_deployment(self, db_models: list) -> int:
         """
         Iterate through db models
@@ -5272,15 +5317,13 @@ class ProxyConfig:
         ## ADD MODEL LOGIC
         for m in db_models:
             _litellm_params = m.litellm_params
+            resolve_env_refs = not _db_model_is_team_scoped(m)
             if isinstance(_litellm_params, dict):
                 # decrypt values
                 for k, v in _litellm_params.items():
-                    if isinstance(v, str):
-                        # decrypt value - returns original value if decryption fails or no key is set
-                        _value = decrypt_value_helper(
-                            value=v, key=k, return_original_value=True
-                        )
-                        _litellm_params[k] = _value
+                    _litellm_params[k] = self._resolve_db_litellm_param(
+                        key=k, value=v, resolve_env_refs=resolve_env_refs
+                    )
                 _litellm_params = LiteLLM_Params(**_litellm_params)
 
             else:
@@ -5308,15 +5351,15 @@ class ProxyConfig:
         _model_list: list = []
         for m in new_models:
             _litellm_params = m.litellm_params
+            resolve_env_refs = not _db_model_is_team_scoped(m)
             if isinstance(_litellm_params, BaseModel):
                 _litellm_params = _litellm_params.model_dump()
             if isinstance(_litellm_params, dict):
                 # decrypt values
                 for k, v in _litellm_params.items():
-                    decrypted_value = decrypt_value_helper(
-                        value=v, key=k, return_original_value=True
+                    _litellm_params[k] = self._resolve_db_litellm_param(
+                        key=k, value=v, resolve_env_refs=resolve_env_refs
                     )
-                    _litellm_params[k] = decrypted_value
                 _litellm_params = LiteLLM_Params(**_litellm_params)
             else:
                 verbose_proxy_logger.error(

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -79,6 +79,7 @@ def _get_spend_logs_metadata(
     cold_storage_object_key: Optional[str] = None,
     litellm_overhead_time_ms: Optional[float] = None,
     cost_breakdown: Optional[CostBreakdown] = None,
+    litellm_call_id: Optional[str] = None,
 ) -> SpendLogsMetadata:
     if metadata is None:
         return SpendLogsMetadata(
@@ -109,6 +110,7 @@ def _get_spend_logs_metadata(
             attempted_retries=None,
             max_retries=None,
             cost_breakdown=None,
+            litellm_call_id=litellm_call_id,
         )
     verbose_proxy_logger.debug(
         "getting payload for SpendLogs, available keys in metadata: "
@@ -133,6 +135,7 @@ def _get_spend_logs_metadata(
     clean_metadata["cold_storage_object_key"] = cold_storage_object_key
     clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms
     clean_metadata["cost_breakdown"] = cost_breakdown
+    clean_metadata["litellm_call_id"] = litellm_call_id
 
     return clean_metadata
 
@@ -383,6 +386,10 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
             if standard_logging_payload is not None
             else None
         ),
+        litellm_call_id=cast(
+            Optional[str],
+            kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"),
+        ),
     )
 
     special_usage_fields = ["completion_tokens", "prompt_tokens", "total_tokens"]

litellm/responses/litellm_completion_transformation/handler.py

@@ -59,6 +59,7 @@ class LiteLLMCompletionTransformationHandler:
         completion_args = {}
         completion_args.update(kwargs)
         completion_args.update(litellm_completion_request)
+        completion_args["_skip_responses_api_bridge"] = True
 
         litellm_completion_response: Union[
             ModelResponse, litellm.CustomStreamWrapper
@@ -107,6 +108,7 @@ class LiteLLMCompletionTransformationHandler:
         acompletion_args = {}
         acompletion_args.update(kwargs)
         acompletion_args.update(litellm_completion_request)
+        acompletion_args["_skip_responses_api_bridge"] = True
 
         litellm_completion_response: Union[
             ModelResponse, litellm.CustomStreamWrapper

litellm/router.py

@@ -2323,6 +2323,14 @@ class Router:
                     verbose_router_logger.error(
                         f"Fallback also failed: {fallback_error}"
                     )
+                    # No fallback handled the mid-stream error, so surface the
+                    # real provider exception (e.g. RateLimitError) instead of
+                    # leaking the internal MidStreamFallbackError to the client
+                    if (
+                        isinstance(fallback_error, MidStreamFallbackError)
+                        and fallback_error.original_exception is not None
+                    ):
+                        raise fallback_error.original_exception from fallback_error
                     raise fallback_error
             finally:
                 # Close the underlying streams to release HTTP connections
@@ -2754,6 +2762,11 @@ class Router:
                     verbose_router_logger.error(
                         f"Responses streaming fallback also failed: {fallback_error}"
                     )
+                    if (
+                        isinstance(fallback_error, MidStreamFallbackError)
+                        and fallback_error.original_exception is not None
+                    ):
+                        raise fallback_error.original_exception from fallback_error
                     raise fallback_error
             finally:
                 with anyio.CancelScope(shield=True):
@@ -2890,6 +2903,11 @@ class Router:
                     verbose_router_logger.error(
                         f"Fallback also failed: {fallback_error}"
                     )
+                    if (
+                        isinstance(fallback_error, MidStreamFallbackError)
+                        and fallback_error.original_exception is not None
+                    ):
+                        raise fallback_error.original_exception from fallback_error
                     raise fallback_error
             finally:
                 if hasattr(model_response, "close"):

litellm/rust_bridge/__init__.py

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

litellm/rust_bridge/ocr.py

@@ -1,30 +1,21 @@
-"""
-Optional Rust-backed OCR path.
-
-Enable with ``litellm.use_litellm_rust()``; the sync ``litellm.ocr()`` entrypoint
-then routes supported Mistral calls through the compiled ``litellm.rust_bridge._native``
-extension, which performs the whole OCR call (URL, headers, HTTP, parse) in Rust.
-
-No module-level ``litellm`` imports keep this a leaf so ``litellm/ocr/main.py``
-can import it statically without forming an import cycle.
-"""
+"""Thin Python wrapper for the native Rust OCR bridge."""
 
 from __future__ import annotations
 
 import os
-from typing import Awaitable, Final, Protocol, cast
+from typing import Any, Awaitable, Final, Protocol, Union, cast
+
+import httpx
 
 
 class RustOcr(Protocol):
-    """Signature of the compiled Rust OCR entrypoint."""
-
     def __call__(
         self,
         model: str,
         document: dict[str, object],
         api_key: str | None,
         api_base: str | None,
-        custom_llm_provider: str,
+        custom_llm_provider: str | None,
         extra_headers: dict[str, object] | None,
         optional_params: dict[str, object],
         timeout_seconds: float | None,
@@ -33,15 +24,13 @@ class RustOcr(Protocol):
 
 
 class RustAocr(Protocol):
-    """Signature of the compiled ``litellm_python_bridge.aocr`` entrypoint."""
-
     def __call__(
         self,
         model: str,
         document: dict[str, object],
         api_key: str | None,
         api_base: str | None,
-        custom_llm_provider: str,
+        custom_llm_provider: str | None,
         extra_headers: dict[str, object] | None,
         optional_params: dict[str, object],
         timeout_seconds: float | None,
@@ -50,7 +39,7 @@ class RustAocr(Protocol):
 
 
 class _Unset:
-    """Sentinel type so ``ocr=None`` can clear a prior injection while omission preserves it."""
+    pass
 
 
 _UNSET: Final[_Unset] = _Unset()
@@ -76,12 +65,6 @@ def use_litellm_rust(
     ocr: RustOcr | None | _Unset = _UNSET,
     aocr: RustAocr | None | _Unset = _UNSET,
 ) -> None:
-    """Route supported OCR calls through the packaged Rust extension.
-
-    ``ocr`` and ``aocr`` inject bridge callables; when omitted the compiled
-    extension is loaded on demand and any previously injected bridge is
-    preserved. Pass ``None`` explicitly to clear a prior injection.
-    """
     global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl
     _rust_ocr_enabled = enabled
     if not isinstance(ocr, _Unset):
@@ -91,17 +74,10 @@ def use_litellm_rust(
 
 
 def rust_ocr_enabled() -> bool:
-    """Whether the Rust OCR path has been turned on via ``use_litellm_rust()``."""
     return _rust_ocr_enabled
 
 
 def load_rust_ocr() -> RustOcr | None:
-    """Return the Rust OCR callable, or ``None`` when no bridge is available.
-
-    Prefers an injected implementation, otherwise loads the compiled
-    ``litellm.rust_bridge._native`` extension; a missing extension yields ``None`` so
-    the caller can fall back to the Python path instead of hard-failing.
-    """
     if _rust_ocr_impl is not None:
         return _rust_ocr_impl
     from litellm.rust_bridge import get_native_bridge
@@ -113,7 +89,6 @@ def load_rust_ocr() -> RustOcr | None:
 
 
 def load_rust_aocr() -> RustAocr | None:
-    """Return the async Rust OCR callable, or ``None`` when unavailable."""
     if _rust_aocr_impl is not None:
         return _rust_aocr_impl
     from litellm.rust_bridge import get_native_bridge
@@ -122,3 +97,63 @@ def load_rust_aocr() -> RustAocr | None:
     if native_bridge is None:
         return None
     return cast(RustAocr, getattr(native_bridge, "aocr", None))
+
+
+def _timeout_to_seconds(timeout: Union[float, httpx.Timeout] | None) -> float | None:
+    if timeout is None:
+        return None
+    if isinstance(timeout, httpx.Timeout):
+        return timeout.read
+    return float(timeout)
+
+
+def ocr(
+    *,
+    model: str,
+    document: dict[str, Any],
+    api_key: str | None,
+    api_base: str | None,
+    custom_llm_provider: str | None,
+    extra_headers: dict[str, Any] | None,
+    optional_params: dict[str, object],
+    timeout: Union[float, httpx.Timeout] | None,
+) -> dict[str, object] | None:
+    rust_ocr = load_rust_ocr()
+    if rust_ocr is None:
+        return None
+    return rust_ocr(
+        model=model,
+        document=cast(dict[str, object], document),
+        api_key=api_key,
+        api_base=api_base,
+        custom_llm_provider=custom_llm_provider,
+        extra_headers=cast(dict[str, object] | None, extra_headers),
+        optional_params=optional_params,
+        timeout_seconds=_timeout_to_seconds(timeout),
+    )
+
+
+async def aocr(
+    *,
+    model: str,
+    document: dict[str, Any],
+    api_key: str | None,
+    api_base: str | None,
+    custom_llm_provider: str | None,
+    extra_headers: dict[str, Any] | None,
+    optional_params: dict[str, object],
+    timeout: Union[float, httpx.Timeout] | None,
+) -> dict[str, object] | None:
+    rust_aocr = load_rust_aocr()
+    if rust_aocr is None:
+        return None
+    return await rust_aocr(
+        model=model,
+        document=cast(dict[str, object], document),
+        api_key=api_key,
+        api_base=api_base,
+        custom_llm_provider=custom_llm_provider,
+        extra_headers=cast(dict[str, object] | None, extra_headers),
+        optional_params=optional_params,
+        timeout_seconds=_timeout_to_seconds(timeout),
+    )

litellm/types/utils.py

@@ -1539,6 +1539,9 @@ class PromptTokensDetailsWrapper(
     video_length_seconds: Optional[float] = None
     """Length of videos sent to the model. Used for Vertex AI multimodal embeddings."""
 
+    audio_length_seconds: Optional[float] = None
+    """Length of audio sent to the model. Used for multimodal embeddings priced per audio-second."""
+
     cache_creation_tokens: Optional[int] = None
     """Number of cache creation tokens sent to the model. Used for Anthropic prompt caching."""
 
@@ -1553,6 +1556,8 @@ class PromptTokensDetailsWrapper(
             del self.image_count
         if self.video_length_seconds is None:
             del self.video_length_seconds
+        if self.audio_length_seconds is None:
+            del self.audio_length_seconds
         if self.web_search_requests is None:
             del self.web_search_requests
         if self.cache_creation_tokens is None:

model_prices_and_context_window.json

@@ -25965,7 +25965,7 @@
         "supports_response_schema": true,
         "supports_tool_choice": true
     },
-    "mistral/mistral-medium-latest": {
+    "mistral/mistral-medium-2508": {
         "input_cost_per_token": 4e-07,
         "litellm_provider": "mistral",
         "max_input_tokens": 131072,
@@ -25973,12 +25973,45 @@
         "max_tokens": 131072,
         "mode": "chat",
         "output_cost_per_token": 2e-06,
+        "source": "https://mistral.ai/news/mistral-medium-3",
         "supports_assistant_prefill": true,
         "supports_function_calling": true,
         "supports_response_schema": true,
         "supports_tool_choice": true,
         "supports_vision": true
     },
+    "mistral/mistral-medium-2604": {
+        "input_cost_per_token": 1.5e-06,
+        "litellm_provider": "mistral",
+        "max_input_tokens": 262144,
+        "max_output_tokens": 262144,
+        "max_tokens": 262144,
+        "mode": "chat",
+        "output_cost_per_token": 7.5e-06,
+        "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
+        "supports_assistant_prefill": true,
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "mistral/mistral-medium-latest": {
+        "input_cost_per_token": 1.5e-06,
+        "litellm_provider": "mistral",
+        "max_input_tokens": 262144,
+        "max_output_tokens": 262144,
+        "max_tokens": 262144,
+        "mode": "chat",
+        "output_cost_per_token": 7.5e-06,
+        "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
+        "supports_assistant_prefill": true,
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "mistral/mistral-medium-3-1-2508": {
         "input_cost_per_token": 4e-07,
         "litellm_provider": "mistral",
@@ -26005,6 +26038,7 @@
         "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
         "supports_assistant_prefill": true,
         "supports_function_calling": true,
+        "supports_reasoning": true,
         "supports_response_schema": true,
         "supports_tool_choice": true,
         "supports_vision": true

tests/logging_callback_tests/test_gcs_pub_sub.py

@@ -31,6 +31,7 @@ verbose_logger.setLevel(logging.DEBUG)
 
 ignored_keys = [
     "request_id",
+    "metadata.litellm_call_id",
     "session_id",
     "startTime",
     "endTime",

tests/mcp_tests/test_mcp_server.py

@@ -1883,10 +1883,11 @@ def test_get_server_auth_header_no_auth_headers():
 
 
 def test_create_tool_response_objects():
-    """Test _create_tool_response_objects function."""
+    """Test _create_tool_response_objects enriches mcp_info with server_id and alias."""
     from litellm.proxy._experimental.mcp_server.rest_endpoints import (
         _create_tool_response_objects,
     )
+    from litellm.types.mcp_server.mcp_server_manager import MCPServer
     from mcp.types import Tool as MCPTool
 
     # Create mock tools
@@ -1903,20 +1904,32 @@ def test_create_tool_response_objects():
         ),
     ]
 
-    server_mcp_info = {
-        "server_name": "zapier",
+    server = MCPServer(
+        server_id="a1b2c3d4",
+        name="zapier_internal",
+        alias="zapier",
+        transport="http",
+        mcp_info={
+            "server_name": "zapier_internal",
+            "logo_url": "https://zapier.com/logo.png",
+        },
+    )
+
+    result = _create_tool_response_objects(mock_tools, server)
+
+    expected_mcp_info = {
+        "server_name": "zapier_internal",
         "logo_url": "https://zapier.com/logo.png",
+        "server_id": "a1b2c3d4",
+        "alias": "zapier",
     }
-
-    result = _create_tool_response_objects(mock_tools, server_mcp_info)
-
     assert len(result) == 2
     assert result[0].name == "send_email"
     assert result[0].description == "Send an email"
-    assert result[0].mcp_info == server_mcp_info
+    assert result[0].mcp_info == expected_mcp_info
     assert result[1].name == "create_event"
     assert result[1].description == "Create a calendar event"
-    assert result[1].mcp_info == server_mcp_info
+    assert result[1].mcp_info == expected_mcp_info
 
 
 @pytest.mark.asyncio
@@ -1930,6 +1943,8 @@ async def test_get_tools_for_single_server():
     # Create a mock server (pin allowlist fields; MagicMock auto-attrs are truthy)
     mock_server = MagicMock()
     mock_server.mcp_info = {"server_name": "zapier"}
+    mock_server.server_id = "zapier_id"
+    mock_server.alias = "zapier_alias"
     mock_server.allowed_tools = None
     mock_server.disallowed_tools = None
 
@@ -1963,7 +1978,11 @@ async def test_get_tools_for_single_server():
         # Verify the result
         assert len(result) == 1
         assert result[0].name == "send_email"
-        assert result[0].mcp_info == {"server_name": "zapier"}
+        assert result[0].mcp_info == {
+            "server_name": "zapier",
+            "server_id": "zapier_id",
+            "alias": "zapier_alias",
+        }
 
 
 @pytest.mark.asyncio

tests/test_litellm/caching/test_embedding_router.py

@@ -0,0 +1,67 @@
+import os
+import sys
+from unittest.mock import MagicMock
+
+sys.path.insert(0, os.path.abspath("../../.."))
+
+from litellm.caching._embedding_router import (
+    build_router_embedding_metadata,
+    resolve_embedding_router,
+)
+
+
+def test_resolve_returns_router_when_model_is_a_deployment():
+    router = MagicMock()
+    assert (
+        resolve_embedding_router("sem-embed", router, [{"model_name": "sem-embed"}])
+        is router
+    )
+
+
+def test_resolve_returns_none_when_model_not_in_router():
+    router = MagicMock()
+    assert (
+        resolve_embedding_router("sem-embed", router, [{"model_name": "other"}]) is None
+    )
+
+
+def test_resolve_returns_none_when_router_is_none():
+    assert (
+        resolve_embedding_router("sem-embed", None, [{"model_name": "sem-embed"}])
+        is None
+    )
+
+
+def test_resolve_returns_none_when_model_list_is_none():
+    router = MagicMock()
+    assert resolve_embedding_router("sem-embed", router, None) is None
+
+
+def test_resolve_skips_entries_missing_model_name():
+    router = MagicMock()
+    model_list = [
+        {"litellm_params": {"model": "bedrock/x"}},
+        {"model_name": "sem-embed"},
+    ]
+    assert resolve_embedding_router("sem-embed", router, model_list) is router
+    assert resolve_embedding_router("other", router, [{"litellm_params": {}}]) is None
+
+
+def test_build_metadata_preserves_request_fields_and_adds_flag():
+    md = build_router_embedding_metadata(
+        {"user_api_key": "sk-x", "user_api_key_team_id": "team-1", "trace_id": "t-1"}
+    )
+    assert md == {
+        "user_api_key": "sk-x",
+        "user_api_key_team_id": "team-1",
+        "trace_id": "t-1",
+        "semantic-cache-embedding": True,
+    }
+
+
+def test_build_metadata_handles_none_and_does_not_mutate_input():
+    original = {"user_api_key": "sk-x"}
+    md = build_router_embedding_metadata(original)
+    assert md == {"user_api_key": "sk-x", "semantic-cache-embedding": True}
+    assert original == {"user_api_key": "sk-x"}
+    assert build_router_embedding_metadata(None) == {"semantic-cache-embedding": True}

tests/test_litellm/caching/test_qdrant_semantic_cache.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import types
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -806,3 +807,104 @@ def test_qdrant_semantic_cache_large_vector_size():
         )
         create_payload = put_call.kwargs["json"]
         assert create_payload["vectors"]["size"] == 4096
+
+
+def _router_proxy_module(router, model_name):
+    mod = types.ModuleType("litellm.proxy.proxy_server")
+    mod.llm_router = router
+    mod.llm_model_list = [{"model_name": model_name}]
+    return mod
+
+
+def test_qdrant_sync_get_cache_routes_through_router(monkeypatch):
+    from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
+
+    cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
+    cache.embedding_model = "sem-embed"
+    cache.qdrant_api_base = "http://test.qdrant.local"
+    cache.collection_name = "test_collection"
+    cache.headers = {"Content-Type": "application/json", "api-key": "test_key"}
+    cache.similarity_threshold = 0.8
+    cache.sync_client = MagicMock()
+    search_response = MagicMock()
+    search_response.status_code = 200
+    search_response.json.return_value = {"result": []}
+    cache.sync_client.post.return_value = search_response
+
+    router = MagicMock()
+    router.embedding = MagicMock(
+        return_value={"data": [{"embedding": [0.3, 0.3, 0.3]}]}
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        "litellm.proxy.proxy_server",
+        _router_proxy_module(router, "sem-embed"),
+    )
+
+    with patch("litellm.embedding") as direct_embed:
+        result = cache.get_cache(
+            key="test_key",
+            messages=[{"content": "What is the capital of France?"}],
+            metadata={},
+        )
+
+    assert result is None
+    router.embedding.assert_called_once()
+    assert router.embedding.call_args.kwargs["model"] == "sem-embed"
+    direct_embed.assert_not_called()
+
+
+def test_qdrant_sync_set_cache_falls_back_to_direct(monkeypatch):
+    from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
+
+    cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
+    cache.embedding_model = "text-embedding-ada-002"
+    cache.qdrant_api_base = "http://test.qdrant.local"
+    cache.collection_name = "test_collection"
+    cache.headers = {"Content-Type": "application/json", "api-key": "test_key"}
+    cache.sync_client = MagicMock()
+    put_response = MagicMock()
+    put_response.status_code = 200
+    cache.sync_client.put.return_value = put_response
+
+    fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
+    fake_proxy.llm_router = None
+    fake_proxy.llm_model_list = None
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
+
+    with patch(
+        "litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.1, 0.1]}]}
+    ) as direct_embed:
+        cache.set_cache(
+            key="test_key",
+            value={"content": "Paris"},
+            messages=[{"content": "What is the capital of France?"}],
+        )
+
+    direct_embed.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_qdrant_async_embedding_forwards_full_metadata(monkeypatch):
+    from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
+
+    cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
+    cache.embedding_model = "sem-embed"
+
+    router = MagicMock()
+    router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]})
+    monkeypatch.setitem(
+        sys.modules,
+        "litellm.proxy.proxy_server",
+        _router_proxy_module(router, "sem-embed"),
+    )
+
+    await cache._get_async_embedding(
+        "hello",
+        metadata={"user_api_key": "sk-x", "user_api_key_team_id": "team-1"},
+    )
+
+    md = router.aembedding.call_args.kwargs["metadata"]
+    assert md["user_api_key"] == "sk-x"
+    assert md["user_api_key_team_id"] == "team-1"
+    assert md["semantic-cache-embedding"] is True

tests/test_litellm/caching/test_redis_cache.py

@@ -3,7 +3,6 @@ import sys
 from unittest.mock import MagicMock, patch
 
 import pytest
-from fastapi.testclient import TestClient
 
 sys.path.insert(
     0, os.path.abspath("../../..")
@@ -94,6 +93,42 @@ async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl(
     mock_redis_instance.expire.assert_not_awaited()
 
 
+@pytest.mark.parametrize("namespace", [None, "litellm"])
+@pytest.mark.asyncio
+async def test_async_delete_cache_applies_namespace(
+    namespace, monkeypatch, redis_no_ping
+):
+    """async_delete_cache must prefix keys with the namespace, matching every
+    other cache operation. Without this, Redis NOPERM errors occur when an
+    ACL restricts DEL to the litellm:* pattern."""
+    monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
+    redis_cache = RedisCache(namespace=namespace)
+    mock_redis_instance = AsyncMock()
+
+    with patch.object(
+        redis_cache, "init_async_client", return_value=mock_redis_instance
+    ):
+        await redis_cache.async_delete_cache(key="3997c4abcdef")
+
+    expected_key = "litellm:3997c4abcdef" if namespace else "3997c4abcdef"
+    mock_redis_instance.delete.assert_awaited_once_with(expected_key)
+
+
+@pytest.mark.parametrize("namespace", [None, "litellm"])
+def test_delete_cache_applies_namespace(namespace, monkeypatch, redis_no_ping):
+    """delete_cache must prefix keys with the namespace, matching every other
+    cache operation."""
+    monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
+    redis_cache = RedisCache(namespace=namespace)
+    mock_redis_client = MagicMock()
+    redis_cache.redis_client = mock_redis_client
+
+    redis_cache.delete_cache(key="3997c4abcdef")
+
+    expected_key = "litellm:3997c4abcdef" if namespace else "3997c4abcdef"
+    mock_redis_client.delete.assert_called_once_with(expected_key)
+
+
 @pytest.mark.asyncio
 async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping):
     monkeypatch.setenv("REDIS_HOST", "my-fake-host")

tests/test_litellm/caching/test_redis_semantic_cache.py

@@ -104,6 +104,7 @@ def test_redis_semantic_cache_get_cache(monkeypatch):
             # Verify llmcache.check was called
             redis_semantic_cache.llmcache.check.assert_called_once_with(
                 prompt="What is the capital of France?",
+                vector=[0.1, 0.2, 0.3],
                 filter_expression="cache-key-filter",
             )
 
@@ -138,10 +139,16 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch):
             ]
         )
 
-        with patch.object(
-            redis_semantic_cache,
-            "_get_cache_key_filter_expression",
-            return_value="cache-key-filter",
+        with (
+            patch(
+                "litellm.embedding",
+                return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]},
+            ),
+            patch.object(
+                redis_semantic_cache,
+                "_get_cache_key_filter_expression",
+                return_value="cache-key-filter",
+            ),
         ):
             metadata = {}
             result = redis_semantic_cache.get_cache(
@@ -176,16 +183,21 @@ def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch):
         redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
         redis_semantic_cache.llmcache.store = MagicMock()
 
-        redis_semantic_cache.set_cache(
-            key="test_key",
-            value={"content": "Paris"},
-            messages=[{"content": "What is the capital of France?"}],
-            ttl=60,
-        )
+        with patch(
+            "litellm.embedding",
+            return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]},
+        ):
+            redis_semantic_cache.set_cache(
+                key="test_key",
+                value={"content": "Paris"},
+                messages=[{"content": "What is the capital of France?"}],
+                ttl=60,
+            )
 
         redis_semantic_cache.llmcache.store.assert_called_once_with(
             "What is the capital of France?",
             "{'content': 'Paris'}",
+            vector=[0.1, 0.2, 0.3],
             filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
             ttl=60,
         )
@@ -299,10 +311,11 @@ def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypat
         monkeypatch.setenv("REDIS_PASSWORD", "test_password")
 
         with pytest.raises(ValueError, match="connection failed"):
-            RedisSemanticCache(
+            cache = RedisSemanticCache(
                 similarity_threshold=0.8,
                 index_name="existing_index",
             )
+            _ = cache.llmcache
 
 
 def test_redis_semantic_cache_reraises_unexpected_index_error():
@@ -534,6 +547,7 @@ def test_redis_semantic_cache_set_cache_uses_responses_string_input():
         return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
     )
     redis_semantic_cache._get_ttl = MagicMock(return_value=None)
+    redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
 
     redis_semantic_cache.set_cache(
         key="test_key",
@@ -544,6 +558,7 @@ def test_redis_semantic_cache_set_cache_uses_responses_string_input():
     redis_semantic_cache.llmcache.store.assert_called_once_with(
         "What is the capital of France?",
         "{'content': 'Paris'}",
+        vector=[0.1, 0.2, 0.3],
         filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
     )
 
@@ -564,6 +579,7 @@ def test_redis_semantic_cache_get_cache_uses_responses_string_input():
             }
         ]
     )
+    redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
 
     with patch.object(
         redis_semantic_cache,
@@ -581,6 +597,7 @@ def test_redis_semantic_cache_get_cache_uses_responses_string_input():
     assert metadata["semantic-similarity"] == pytest.approx(0.9)
     redis_semantic_cache.llmcache.check.assert_called_once_with(
         prompt="What is the capital of France?",
+        vector=[0.1, 0.2, 0.3],
         filter_expression="cache-key-filter",
     )
 
@@ -594,6 +611,7 @@ def test_redis_semantic_cache_set_cache_flattens_structured_responses_input():
         return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
     )
     redis_semantic_cache._get_ttl = MagicMock(return_value=None)
+    redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
 
     redis_semantic_cache.set_cache(
         key="test_key",
@@ -616,6 +634,7 @@ def test_redis_semantic_cache_set_cache_flattens_structured_responses_input():
     redis_semantic_cache.llmcache.store.assert_called_once_with(
         "What is the capital of France?\nAnswer briefly.",
         "{'content': 'Paris'}",
+        vector=[0.1, 0.2, 0.3],
         filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
     )
 
@@ -740,6 +759,7 @@ def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results():
     redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
     redis_semantic_cache.llmcache = MagicMock()
     redis_semantic_cache.llmcache.check = MagicMock(return_value=[])
+    redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
 
     with patch.object(
         redis_semantic_cache,
@@ -757,6 +777,7 @@ def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results():
     assert metadata["semantic-similarity"] == 0.0
     redis_semantic_cache.llmcache.check.assert_called_once_with(
         prompt="What is the capital of France?",
+        vector=[0.1, 0.2, 0.3],
         filter_expression="cache-key-filter",
     )
 
@@ -870,6 +891,63 @@ async def test_redis_semantic_cache_async_paths_set_similarity_on_misses():
     )
 
 
+def test_redis_get_embedding_routes_through_router(monkeypatch):
+    import sys
+    import types
+
+    from litellm.caching.redis_semantic_cache import RedisSemanticCache
+
+    cache = RedisSemanticCache.__new__(RedisSemanticCache)
+    cache.embedding_model = "sem-embed"
+
+    router = MagicMock()
+    router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]})
+    fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
+    fake_proxy.llm_router = router
+    fake_proxy.llm_model_list = [{"model_name": "sem-embed"}]
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
+
+    with patch("litellm.embedding") as direct_embed:
+        vec = cache._get_embedding("hello", metadata={"user_api_key": "sk-x"})
+
+    assert vec == [0.5, 0.6]
+    router.embedding.assert_called_once()
+    assert router.embedding.call_args.kwargs["model"] == "sem-embed"
+    assert router.embedding.call_args.kwargs["input"] == "hello"
+    assert router.embedding.call_args.kwargs["cache"] == {
+        "no-store": True,
+        "no-cache": True,
+    }
+    assert router.embedding.call_args.kwargs["metadata"] == {
+        "user_api_key": "sk-x",
+        "semantic-cache-embedding": True,
+    }
+    direct_embed.assert_not_called()
+
+
+def test_redis_get_embedding_falls_back_to_direct(monkeypatch):
+    import sys
+    import types
+
+    from litellm.caching.redis_semantic_cache import RedisSemanticCache
+
+    cache = RedisSemanticCache.__new__(RedisSemanticCache)
+    cache.embedding_model = "text-embedding-ada-002"
+
+    fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
+    fake_proxy.llm_router = None
+    fake_proxy.llm_model_list = None
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
+
+    with patch(
+        "litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2]}]}
+    ) as direct_embed:
+        vec = cache._get_embedding("hello")
+
+    assert vec == [0.1, 0.2]
+    direct_embed.assert_called_once()
+
+
 def test_cache_get_cache_passes_responses_input_to_backend_cache():
     from litellm.caching.caching import Cache
 
@@ -893,7 +971,7 @@ def test_cache_get_cache_passes_responses_input_to_backend_cache():
     )
 
 
-def test_cache_get_cache_filters_sensitive_kwargs_from_backend_cache():
+def test_cache_get_cache_filters_non_lookup_kwargs_from_backend_cache():
     from litellm.caching.caching import Cache
 
     cache = Cache.__new__(Cache)
@@ -927,7 +1005,11 @@ def test_cache_get_cache_filters_sensitive_kwargs_from_backend_cache():
     forwarded_kwargs = cache.cache.get_cache.call_args.kwargs
     assert forwarded_kwargs == {
         "input": "What is the capital of France?",
-        "metadata": {"semantic-similarity": 0.7},
+        "metadata": {
+            "user_api_key": "sk-secret",
+            "trace_id": "trace-id",
+            "semantic-similarity": 0.7,
+        },
     }
     assert forwarded_kwargs["metadata"] is not metadata
     cache._get_cache_logic.assert_called_once_with(
@@ -988,3 +1070,166 @@ def test_cache_get_cache_passes_responses_input_to_dynamic_cache():
         cached_result={"content": "Paris"},
         max_age=float("inf"),
     )
+
+
+def test_redis_sync_set_cache_passes_precomputed_vector():
+    from litellm.caching.redis_semantic_cache import RedisSemanticCache
+
+    cache = RedisSemanticCache.__new__(RedisSemanticCache)
+    cache.llmcache = MagicMock()
+    cache._get_cache_filters = MagicMock(
+        return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
+    )
+    cache._get_ttl = MagicMock(return_value=None)
+    cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
+
+    cache.set_cache(
+        key="test_key",
+        value={"content": "Paris"},
+        messages=[{"content": "What is the capital of France?"}],
+    )
+
+    cache._get_embedding.assert_called_once()
+    cache.llmcache.store.assert_called_once_with(
+        "What is the capital of France?",
+        "{'content': 'Paris'}",
+        vector=[0.1, 0.2, 0.3],
+        filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
+    )
+
+
+def test_redis_sync_get_cache_passes_precomputed_vector():
+    from litellm.caching.redis_semantic_cache import RedisSemanticCache
+
+    cache = RedisSemanticCache.__new__(RedisSemanticCache)
+    cache.similarity_threshold = 0.8
+    cache.llmcache = MagicMock()
+    cache.llmcache.check = MagicMock(
+        return_value=[
+            {
+                "prompt": "What is the capital of France?",
+                "response": '{"content": "Paris"}',
+                "vector_distance": 0.1,
+                RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
+            }
+        ]
+    )
+    cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
+
+    with patch.object(
+        cache, "_get_cache_key_filter_expression", return_value="cache-key-filter"
+    ):
+        result = cache.get_cache(
+            key="test_key",
+            messages=[{"content": "What is the capital of France?"}],
+            metadata={},
+        )
+
+    assert result == {"content": "Paris"}
+    cache._get_embedding.assert_called_once()
+    cache.llmcache.check.assert_called_once_with(
+        prompt="What is the capital of France?",
+        vector=[0.1, 0.2, 0.3],
+        filter_expression="cache-key-filter",
+    )
+
+
+@pytest.mark.asyncio
+async def test_redis_async_embedding_forwards_full_metadata(monkeypatch):
+    import sys
+    import types
+
+    from litellm.caching.redis_semantic_cache import RedisSemanticCache
+
+    cache = RedisSemanticCache.__new__(RedisSemanticCache)
+    cache.embedding_model = "sem-embed"
+
+    router = MagicMock()
+    router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]})
+    fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
+    fake_proxy.llm_router = router
+    fake_proxy.llm_model_list = [{"model_name": "sem-embed"}]
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
+
+    await cache._get_async_embedding(
+        "hello",
+        metadata={"user_api_key": "sk-x", "user_api_key_team_id": "team-1"},
+    )
+
+    md = router.aembedding.call_args.kwargs["metadata"]
+    assert md["user_api_key"] == "sk-x"
+    assert md["user_api_key_team_id"] == "team-1"  # FAILS today: team_id is dropped
+    assert md["semantic-cache-embedding"] is True
+
+
+def test_redis_init_defers_redisvl_construction(monkeypatch):
+    semantic_cache_mock = MagicMock()
+    custom_vectorizer_mock = MagicMock()
+
+    with patch.dict(
+        "sys.modules",
+        {
+            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
+            "redisvl.utils.vectorize": MagicMock(
+                CustomTextVectorizer=custom_vectorizer_mock
+            ),
+        },
+    ):
+        from litellm.caching.redis_semantic_cache import RedisSemanticCache
+
+        monkeypatch.setenv("REDIS_HOST", "localhost")
+        monkeypatch.setenv("REDIS_PORT", "6379")
+        monkeypatch.setenv("REDIS_PASSWORD", "test_password")
+
+        cache = RedisSemanticCache(similarity_threshold=0.8)
+
+        semantic_cache_mock.assert_not_called()
+        custom_vectorizer_mock.assert_not_called()
+
+        first = cache.llmcache
+        semantic_cache_mock.assert_called_once()
+        custom_vectorizer_mock.assert_called_once()
+
+        second = cache.llmcache
+        assert first is second
+        semantic_cache_mock.assert_called_once()
+
+
+def test_redis_failed_llmcache_build_is_not_memoized(monkeypatch):
+    built_cache = MagicMock()
+    semantic_cache_mock = MagicMock(
+        side_effect=[ConnectionError("redis down"), built_cache]
+    )
+    custom_vectorizer_mock = MagicMock()
+
+    with patch.dict(
+        "sys.modules",
+        {
+            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
+            "redisvl.utils.vectorize": MagicMock(
+                CustomTextVectorizer=custom_vectorizer_mock
+            ),
+        },
+    ):
+        from litellm.caching.redis_semantic_cache import RedisSemanticCache
+
+        monkeypatch.setenv("REDIS_HOST", "localhost")
+        monkeypatch.setenv("REDIS_PORT", "6379")
+        monkeypatch.setenv("REDIS_PASSWORD", "test_password")
+
+        cache = RedisSemanticCache(similarity_threshold=0.8)
+
+        with pytest.raises(ConnectionError, match="redis down"):
+            _ = cache.llmcache
+
+        assert cache.llmcache is built_cache
+        assert semantic_cache_mock.call_count == 2
+
+
+def test_redis_llmcache_setter_supported():
+    from litellm.caching.redis_semantic_cache import RedisSemanticCache
+
+    cache = RedisSemanticCache.__new__(RedisSemanticCache)
+    sentinel = MagicMock()
+    cache.llmcache = sentinel
+    assert cache.llmcache is sentinel

tests/test_litellm/google_genai/test_google_genai_main.py

@@ -2,6 +2,7 @@
 """
 Test to verify the Google GenAI generate_content adapter functionality
 """
+
 import json
 import os
 import sys
@@ -42,4 +43,228 @@ async def test_agenerate_content_stream():
             stream=True,
         )
         mock_post.assert_called_once()
-        mock_post.call_args.kwargs["stream"] == True
+        assert mock_post.call_args.kwargs["stream"] is True
+
+
+def _mock_gemini_post_response():
+    """A minimal stand-in for a successful Gemini generateContent HTTP response."""
+    from unittest.mock import MagicMock
+
+    resp = MagicMock()
+    resp.status_code = 200
+    resp.headers = {}
+    resp.json.return_value = {
+        "candidates": [
+            {
+                "content": {"parts": [{"text": "hi"}], "role": "model"},
+                "finishReason": "STOP",
+            }
+        ]
+    }
+    return resp
+
+
+NATIVE_TOP_LEVEL_FIELD_CASES = [
+    (
+        "safetySettings",
+        [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}],
+    ),
+    ("toolConfig", {"functionCallingConfig": {"mode": "AUTO"}}),
+    ("cachedContent", "cachedContents/abc123"),
+    ("labels", {"team": "search"}),
+]
+
+
+@pytest.mark.parametrize("field_name, field_value", NATIVE_TOP_LEVEL_FIELD_CASES)
+def test_native_top_level_field_forwarded_to_request_body(field_name, field_value):
+    """
+    Regression for https://github.com/BerriAI/litellm/issues/12671
+
+    Google's native generateContent body carries fields like safetySettings at the
+    top level (siblings of generationConfig). The proxy spreads them as loose kwargs
+    into generate_content. They must reach Google's request body at the top level and
+    must NOT be silently dropped nor nested under generationConfig.
+    """
+    from unittest.mock import patch
+
+    from litellm.google_genai.main import generate_content
+    from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+    with patch.object(
+        HTTPHandler, "post", return_value=_mock_gemini_post_response()
+    ) as mock_post:
+        generate_content(
+            model="gemini/gemini-2.0-flash",
+            contents=[{"role": "user", "parts": [{"text": "Say hi"}]}],
+            custom_llm_provider="gemini",
+            api_key="test-key",
+            **{field_name: field_value},
+        )
+
+    assert mock_post.called, "expected the request to reach the HTTP client"
+    body = mock_post.call_args.kwargs["json"]
+    assert (
+        body[field_name] == field_value
+    ), f"{field_name} should be forwarded to Google at the top level"
+    assert field_name not in body.get("generationConfig", {}), (
+        f"{field_name} must be a top-level sibling of generationConfig, "
+        "not nested inside it"
+    )
+
+
+@pytest.mark.asyncio
+async def test_native_safety_settings_forwarded_async():
+    """The async path (used by the proxy's :generateContent route) must also forward
+    native top-level fields."""
+    from unittest.mock import AsyncMock, patch
+
+    from litellm.google_genai.main import agenerate_content
+    from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+
+    safety_settings = [
+        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}
+    ]
+
+    with patch.object(
+        AsyncHTTPHandler,
+        "post",
+        new_callable=AsyncMock,
+        return_value=_mock_gemini_post_response(),
+    ) as mock_post:
+        await agenerate_content(
+            model="gemini/gemini-2.0-flash",
+            contents=[{"role": "user", "parts": [{"text": "Say hi"}]}],
+            custom_llm_provider="gemini",
+            api_key="test-key",
+            safetySettings=safety_settings,
+        )
+
+    assert mock_post.called
+    body = mock_post.call_args.kwargs["json"]
+    assert body["safetySettings"] == safety_settings
+    assert "safetySettings" not in body.get("generationConfig", {})
+
+
+def test_native_fields_coexist_with_generation_config():
+    """Forwarding native top-level fields must not regress the already-working
+    generationConfig path; both must land in their correct positions."""
+    from unittest.mock import patch
+
+    from litellm.google_genai.main import generate_content
+    from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+    safety_settings = [
+        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}
+    ]
+
+    with patch.object(
+        HTTPHandler, "post", return_value=_mock_gemini_post_response()
+    ) as mock_post:
+        generate_content(
+            model="gemini/gemini-2.0-flash",
+            contents=[{"role": "user", "parts": [{"text": "Say hi"}]}],
+            custom_llm_provider="gemini",
+            api_key="test-key",
+            safetySettings=safety_settings,
+            generationConfig={"temperature": 0, "responseMimeType": "application/json"},
+        )
+
+    body = mock_post.call_args.kwargs["json"]
+    assert body["safetySettings"] == safety_settings
+    generation_config = body["generationConfig"]
+    assert generation_config["temperature"] == 0
+    assert generation_config["responseMimeType"] == "application/json"
+
+
+def test_explicit_extra_body_overrides_native_top_level_field():
+    """An explicit extra_body value takes precedence over the same top-level field."""
+    from unittest.mock import patch
+
+    from litellm.google_genai.main import generate_content
+    from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+    native = [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}]
+    override = [
+        {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"}
+    ]
+
+    with patch.object(
+        HTTPHandler, "post", return_value=_mock_gemini_post_response()
+    ) as mock_post:
+        generate_content(
+            model="gemini/gemini-2.0-flash",
+            contents=[{"role": "user", "parts": [{"text": "Say hi"}]}],
+            custom_llm_provider="gemini",
+            api_key="test-key",
+            safetySettings=native,
+            extra_body={"safetySettings": override},
+        )
+
+    body = mock_post.call_args.kwargs["json"]
+    assert body["safetySettings"] == override
+
+
+def test_native_fields_and_system_instruction_forwarded_on_sync_stream():
+    """The sync streaming path (generate_content_stream) must forward native top-level
+    fields AND systemInstruction. The PR changed the merge here and newly added the
+    systemInstruction kwarg; without coverage a regression on either ships green."""
+    from unittest.mock import patch
+
+    from litellm.google_genai.main import generate_content_stream
+    from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+    safety_settings = [
+        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}
+    ]
+    system_instruction = {"parts": [{"text": "Be terse"}]}
+
+    with patch.object(
+        HTTPHandler, "post", return_value=_mock_gemini_post_response()
+    ) as mock_post:
+        generate_content_stream(
+            model="gemini/gemini-2.0-flash",
+            contents=[{"role": "user", "parts": [{"text": "Say hi"}]}],
+            custom_llm_provider="gemini",
+            api_key="test-key",
+            safetySettings=safety_settings,
+            systemInstruction=system_instruction,
+        )
+
+    assert mock_post.called
+    body = mock_post.call_args.kwargs["json"]
+    assert body["safetySettings"] == safety_settings
+    assert body["systemInstruction"] == system_instruction
+    assert "safetySettings" not in body.get("generationConfig", {})
+
+
+@pytest.mark.asyncio
+async def test_native_fields_forwarded_on_async_stream():
+    """The async streaming path (agenerate_content_stream) backs the proxy's
+    :streamGenerateContent route and must forward native top-level fields too."""
+    from unittest.mock import AsyncMock, patch
+
+    from litellm.google_genai.main import agenerate_content_stream
+    from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+
+    safety_settings = [
+        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}
+    ]
+
+    with patch.object(
+        AsyncHTTPHandler,
+        "post",
+        new_callable=AsyncMock,
+        return_value=_mock_gemini_post_response(),
+    ) as mock_post:
+        await agenerate_content_stream(
+            model="gemini/gemini-2.0-flash",
+            contents=[{"role": "user", "parts": [{"text": "Say hi"}]}],
+            custom_llm_provider="gemini",
+            api_key="test-key",
+            safetySettings=safety_settings,
+        )
+
+    assert mock_post.called
+    body = mock_post.call_args.kwargs["json"]
+    assert body["safetySettings"] == safety_settings
+    assert "safetySettings" not in body.get("generationConfig", {})

tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py

@@ -951,6 +951,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details():
         "character_count": 0,
         "image_count": 0,
         "video_length_seconds": 0.0,
+        "audio_length_seconds": 0.0,
     }
 
     model_info: ModelInfo = {}

tests/test_litellm/litellm_core_utils/test_litellm_logging.py

@@ -11,7 +11,9 @@ sys.path.insert(
 
 import time
 
+import litellm
 from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST
+from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
 from litellm.litellm_core_utils.litellm_logging import set_callbacks
 from litellm.types.utils import ModelResponse, TextCompletionResponse
@@ -3408,6 +3410,140 @@ def test_handle_anthropic_messages_response_logging_degrades_on_unparseable_resp
     assert result.usage.prompt_tokens == 4  # type: ignore[attr-defined]
 
 
+class _SuccessCapturingLogger(CustomLogger):
+    """Records the success payload. success_payload is populated only in
+    async_log_success_event, so it stays None when the buggy no-op
+    async_log_stream_event path runs for streaming."""
+
+    def __init__(self):
+        super().__init__()
+        self.success_payload = None
+        self.success_calls = 0
+        self.stream_event_calls = 0
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self.success_calls += 1
+        self.success_payload = kwargs.get("standard_logging_object")
+
+    async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time):
+        self.stream_event_calls += 1
+
+
+def _responses_stream_sse_bytes():
+    """A full Responses stream: an opened message item, two text deltas, then the
+    terminal response.completed carrying usage. Exercises mid-stream delta handling
+    in addition to end-of-stream success logging."""
+    import json
+
+    events = [
+        {
+            "type": "response.output_item.added",
+            "output_index": 0,
+            "item": {
+                "id": "msg-1",
+                "type": "message",
+                "status": "in_progress",
+                "role": "assistant",
+                "content": [],
+            },
+        },
+        {
+            "type": "response.output_text.delta",
+            "item_id": "msg-1",
+            "output_index": 0,
+            "content_index": 0,
+            "delta": "hello ",
+            "sequence_number": 2,
+        },
+        {
+            "type": "response.output_text.delta",
+            "item_id": "msg-1",
+            "output_index": 0,
+            "content_index": 0,
+            "delta": "world",
+            "sequence_number": 3,
+        },
+        {
+            "type": "response.completed",
+            "sequence_number": 4,
+            "response": _responses_api_response_with_text("hello world").model_dump(),
+        },
+    ]
+    return [f"data: {json.dumps(e)}\n\n".encode("utf-8") for e in events]
+
+
+def _fake_streaming_responses_http_response():
+    sse_chunks = _responses_stream_sse_bytes()
+
+    async def aiter_bytes(*args, **kwargs):
+        for chunk in sse_chunks:
+            yield chunk
+
+    resp = MagicMock()
+    resp.status_code = 200
+    resp.headers = {}
+    resp.aiter_bytes = aiter_bytes
+    return resp
+
+
+def _chunk_text(chunk):
+    if isinstance(chunk, (bytes, bytearray)):
+        return chunk.decode("utf-8", "ignore")
+    return str(chunk)
+
+
+async def _drain_until_logged(logger, max_iter=30):
+    for _ in range(max_iter):
+        if logger.success_payload is not None:
+            break
+        await asyncio.sleep(0.1)
+
+
+@pytest.mark.asyncio
+async def test_streaming_anthropic_messages_openai_bridge_fires_success_logging(
+    monkeypatch,
+):
+    """Regression for #28595 / #28943. The existing tests above call
+    _handle_anthropic_messages_response_logging directly; they do not cover the
+    streaming wiring that originally broke. Drive a real streaming
+    anthropic_messages call routed to the OpenAI Responses backend (upstream SSE
+    mocked) and assert the bridge surfaces delta chunks and fires success logging
+    exactly once with real cost. On the broken version the stream ran but only the
+    no-op async_log_stream_event was called, so success_payload stayed None and the
+    SpendLogs row never landed."""
+    logger = _SuccessCapturingLogger()
+    monkeypatch.setattr(litellm, "callbacks", [logger])
+
+    chunks = []
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        new=AsyncMock(return_value=_fake_streaming_responses_http_response()),
+    ):
+        stream = await litellm.anthropic_messages(
+            model="openai/gpt-4o",
+            api_key="sk-test-28595",
+            messages=[{"role": "user", "content": "ping"}],
+            max_tokens=16,
+            stream=True,
+        )
+        async for chunk in stream:  # logging fires on stream end; must drain fully
+            chunks.append(chunk)
+
+    await _drain_until_logged(logger)
+
+    assert chunks, "stream yielded no chunks"
+    assert any("content_block_delta" in _chunk_text(c) for c in chunks), (
+        "no delta chunks surfaced; the streaming text deltas were not forwarded"
+    )
+    assert logger.success_payload is not None, (
+        "async_log_success_event never fired for streaming /v1/messages -> openai "
+        "Responses bridge; the no-op stream path dropped the spend row"
+    )
+    assert logger.success_calls == 1, "bridge call must log success exactly once"
+    assert logger.success_payload["response_cost"] > 0
+    assert logger.success_payload["call_type"] == "anthropic_messages"
+
+
 def test_failure_handler_records_recovered_partial_spend(logging_obj):
     """A stream interrupted mid-flight still billed the provider for the chunks
     already delivered. When the router stashes that recovered usage as

tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py

@@ -0,0 +1,321 @@
+"""
+Regression tests for the Anthropic message_start cursor=1 bug in
+ChunkProcessor._calculate_usage_per_chunk.
+
+Background
+----------
+Anthropic streams a `message_start` event that carries
+`usage.output_tokens=1` as a placeholder ("cursor"). The real cumulative
+output count only arrives in the final `message_delta` event. When a
+stream is cancelled before `message_delta` lands (very common for
+thinking models on long-tail prompts), the last-wins accumulator in
+ChunkProcessor leaves completion_tokens stuck at 1. Because 1 is
+truthy, the `completion_tokens or token_counter(text=...)` fallback in
+calculate_usage() never fires, and the request is billed for 1 output
+token even when several thousand tokens of text were actually streamed.
+
+These tests pin the post-fix behavior: completion_tokens should reset
+to 0 when the only update we saw was the cursor, allowing the
+text-based fallback to estimate from the real completion text.
+"""
+
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../.."))
+
+from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
+from litellm.types.utils import (
+    Delta,
+    ModelResponseStream,
+    StreamingChoices,
+    Usage,
+)
+
+
+def _make_chunk(
+    *,
+    content: str = "",
+    usage: Usage = None,
+    finish_reason: str = None,
+    custom_llm_provider: str = "anthropic",
+) -> ModelResponseStream:
+    chunk = ModelResponseStream(
+        id="msg_test",
+        created=1738900000,
+        model="claude-sonnet-4-6",
+        object="chat.completion.chunk",
+        choices=[
+            StreamingChoices(
+                finish_reason=finish_reason,
+                index=0,
+                delta=Delta(content=content, role="assistant"),
+            )
+        ],
+        usage=usage,
+    )
+    # The cursor reset is now gated on provider; populate the same field the
+    # real streaming_handler sets (see litellm/litellm_core_utils/streaming_handler.py).
+    chunk._hidden_params = {"custom_llm_provider": custom_llm_provider}
+    return chunk
+
+
+class TestAnthropicCursorBug:
+    """The core regression: completion_tokens=1 cursor must not leak through."""
+
+    def test_only_message_start_cursor_resets_completion_to_zero(self):
+        """
+        Stream cancelled before message_delta — only the message_start cursor
+        (output_tokens=1) was seen. Per-chunk accumulator must reset to 0 so
+        token_counter fallback can estimate from completion text.
+        """
+        # Anthropic message_start: input_tokens accurate, output_tokens=1 cursor
+        message_start = _make_chunk(
+            usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
+        )
+        # Several content_block_delta chunks (no usage attached)
+        text_chunks = [
+            _make_chunk(content="Hello"),
+            _make_chunk(content=" world"),
+            _make_chunk(content=" this is partial."),
+        ]
+        chunks = [message_start, *text_chunks]
+
+        processor = ChunkProcessor(chunks=chunks, messages=[])
+        result = processor._calculate_usage_per_chunk(chunks=chunks)
+
+        assert result["prompt_tokens"] == 1024
+        # The cursor value of 1 must NOT leak through — should be reset to 0
+        # so the text-based fallback estimates the real completion length.
+        assert result["completion_tokens"] == 0, (
+            "completion_tokens=1 from message_start cursor leaked through. "
+            "Should reset to 0 when only cursor was seen, so token_counter "
+            "fallback in calculate_usage() can estimate from completion text."
+        )
+
+    def test_message_start_plus_message_delta_uses_delta_value(self):
+        """
+        Normal complete stream: message_start cursor=1, then message_delta=3847.
+        Last-wins must give 3847 (the real value).
+        """
+        message_start = _make_chunk(
+            usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
+        )
+        text_chunks = [_make_chunk(content=t) for t in ["Hello", " world", "!"]]
+        # message_delta with the real cumulative output_tokens
+        message_delta = _make_chunk(
+            usage=Usage(prompt_tokens=1024, completion_tokens=3847, total_tokens=4871),
+            finish_reason="stop",
+        )
+        chunks = [message_start, *text_chunks, message_delta]
+
+        processor = ChunkProcessor(chunks=chunks, messages=[])
+        result = processor._calculate_usage_per_chunk(chunks=chunks)
+
+        assert result["prompt_tokens"] == 1024
+        assert result["completion_tokens"] == 3847
+
+    def test_calculate_usage_falls_back_to_token_counter_for_cursor_only(self):
+        """
+        End-to-end via calculate_usage(): cursor-only stream + real completion
+        text should produce a token-counter estimate, NOT 1.
+        """
+        message_start = _make_chunk(
+            usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
+        )
+        # ~50 visible chars ≈ ~12 tokens (anthropic-style tokenizer ballpark)
+        text_chunks = [
+            _make_chunk(content="Based on your question, I think the answer is "),
+            _make_chunk(content="forty-two. Here is my reasoning: "),
+        ]
+        chunks = [message_start, *text_chunks]
+        completion_output = (
+            "Based on your question, I think the answer is forty-two. "
+            "Here is my reasoning: "
+        )
+
+        processor = ChunkProcessor(chunks=chunks, messages=[])
+        usage = processor.calculate_usage(
+            chunks=chunks,
+            model="claude-sonnet-4-6",
+            completion_output=completion_output,
+            messages=[],
+        )
+
+        # Should be a token_counter estimate of the text, not the cursor 1
+        assert usage.completion_tokens > 1, (
+            f"Expected token_counter estimate of completion text, got "
+            f"completion_tokens={usage.completion_tokens} (likely stuck at cursor)"
+        )
+
+    def test_cache_fields_preserved_from_message_start(self):
+        """cache_read / cache_creation come from message_start and must survive."""
+        message_start_usage = Usage(
+            prompt_tokens=1024, completion_tokens=1, total_tokens=1025
+        )
+        # Anthropic puts these in message_start
+        message_start_usage.cache_read_input_tokens = 512
+        message_start_usage.cache_creation_input_tokens = 128
+        message_start = _make_chunk(usage=message_start_usage)
+
+        chunks = [message_start, _make_chunk(content="hi")]
+        processor = ChunkProcessor(chunks=chunks, messages=[])
+        result = processor._calculate_usage_per_chunk(chunks=chunks)
+
+        assert result["cache_read_input_tokens"] == 512
+        assert result["cache_creation_input_tokens"] == 128
+
+    def test_openai_streaming_unaffected(self):
+        """
+        OpenAI's only usage chunk is the penultimate one (with
+        stream_options.include_usage=true), and it carries the real value
+        directly. Our cursor fix must not break this path — output > 1
+        means saw_non_cursor_completion=True so no reset happens.
+        """
+        # Simulate OpenAI: content chunks first, then ONE usage chunk at the end
+        text_chunks = [_make_chunk(content=t) for t in ["The", " answer", " is 42"]]
+        usage_chunk = _make_chunk(
+            usage=Usage(prompt_tokens=42, completion_tokens=15, total_tokens=57),
+            finish_reason="stop",
+        )
+        chunks = [*text_chunks, usage_chunk]
+
+        processor = ChunkProcessor(chunks=chunks, messages=[])
+        result = processor._calculate_usage_per_chunk(chunks=chunks)
+
+        assert result["prompt_tokens"] == 42
+        assert result["completion_tokens"] == 15
+
+    def test_single_token_completion_legitimate_case(self):
+        """
+        Edge case: a stream that legitimately completes with output_tokens=1
+        (e.g., model returns just "Yes."). Without saw_non_cursor_completion
+        we'd reset to 0 and fall through to token_counter — but token_counter
+        on a 1-token string also gives ~1, so billing is still approximately
+        correct. This test pins that the result is sane (1 or 0).
+        """
+        message_start = _make_chunk(
+            usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21)
+        )
+        text_chunk = _make_chunk(content="Yes.")
+        # Anthropic's message_delta also gives output_tokens=1 in this case
+        message_delta = _make_chunk(
+            usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21),
+            finish_reason="stop",
+        )
+        chunks = [message_start, text_chunk, message_delta]
+
+        processor = ChunkProcessor(chunks=chunks, messages=[])
+        usage = processor.calculate_usage(
+            chunks=chunks,
+            model="claude-sonnet-4-6",
+            completion_output="Yes.",
+            messages=[],
+        )
+
+        # Two completion-bearing usage events (message_start AND message_delta
+        # both with output_tokens=1) is positive evidence that message_delta
+        # arrived — saw_non_cursor_completion goes True via the count >= 2
+        # branch and the reset is suppressed. Result: completion_tokens stays
+        # at the legitimate value of 1.
+        assert usage.completion_tokens == 1, (
+            f"Legitimate single-token completion should bill exactly 1 token "
+            f"(message_start + message_delta both saw output_tokens=1, "
+            f"confirming message_delta arrived), got {usage.completion_tokens}"
+        )
+
+    def test_anthropic_cache_only_chunks_after_message_start_still_resets(self):
+        """
+        Cache-only chunks (cache_read_input_tokens > 0 but completion_tokens=0)
+        following message_start should not be mistaken for completion progress.
+        The cursor=1 from message_start stays the only completion update; reset
+        must fire so token_counter estimates from completion text instead of
+        billing the placeholder.
+        """
+        message_start_usage = Usage(
+            prompt_tokens=1024, completion_tokens=1, total_tokens=1025
+        )
+        message_start_usage.cache_read_input_tokens = 4096
+        message_start = _make_chunk(usage=message_start_usage)
+        # Subsequent chunks with cache fields but no completion_tokens
+        cache_chunk_usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
+        cache_chunk_usage.cache_read_input_tokens = 4096
+        cache_chunk = _make_chunk(content="partial", usage=cache_chunk_usage)
+        # No message_delta — stream was cancelled
+        chunks = [message_start, cache_chunk]
+
+        processor = ChunkProcessor(chunks=chunks, messages=[])
+        result = processor._calculate_usage_per_chunk(chunks=chunks)
+
+        assert result["cache_read_input_tokens"] == 4096
+        assert result["completion_tokens"] == 0, (
+            "cache chunks alone don't count as completion progress — only "
+            "completion_tokens > 0 in a usage event proves real output happened. "
+            "Reset to 0 forces token_counter fallback."
+        )
+
+
+class TestProviderGuard:
+    """Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic
+    providers, even if they happen to report completion_tokens=1."""
+
+    def test_non_anthropic_provider_completion_tokens_one_not_reset(self):
+        """
+        Some non-Anthropic provider legitimately reports completion_tokens=1
+        in its single usage chunk. Without the provider guard the cursor
+        heuristic would silently reset it to 0 and bill via token_counter,
+        producing a different (often inflated) number than what the provider
+        actually charged.
+        """
+        chunks = [
+            _make_chunk(
+                usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11),
+                finish_reason="stop",
+                custom_llm_provider="openai",
+            ),
+        ]
+        processor = ChunkProcessor(chunks=chunks, messages=[])
+        result = processor._calculate_usage_per_chunk(chunks=chunks)
+        assert result["completion_tokens"] == 1, (
+            "Non-Anthropic providers must not be subject to the message_start "
+            "cursor reset — their completion_tokens=1 is the real value."
+        )
+
+    def test_unknown_provider_completion_tokens_one_not_reset(self):
+        """No custom_llm_provider on hidden_params (older path or custom
+        plugin) — heuristic must not fire."""
+        chunk = _make_chunk(
+            usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11),
+        )
+        # Explicitly clear hidden_params to simulate the unknown-provider case
+        chunk._hidden_params = {}
+        processor = ChunkProcessor(chunks=[chunk], messages=[])
+        result = processor._calculate_usage_per_chunk(chunks=[chunk])
+        assert result["completion_tokens"] == 1
+
+
+class TestNonAnthropicStreamingIntact:
+    """Make sure providers without cursor pattern still work."""
+
+    def test_completion_tokens_above_one_never_resets(self):
+        """Any chunk reporting completion_tokens > 1 sets saw_non_cursor
+        and prevents the reset."""
+        chunks = [
+            _make_chunk(
+                usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
+            ),
+        ]
+        processor = ChunkProcessor(chunks=chunks, messages=[])
+        result = processor._calculate_usage_per_chunk(chunks=chunks)
+        assert result["completion_tokens"] == 5
+
+    def test_no_usage_chunks_leaves_zero(self):
+        """Stream with zero usage info → completion_tokens stays 0
+        (token_counter fallback will handle it)."""
+        chunks = [_make_chunk(content="hi"), _make_chunk(content=" there")]
+        processor = ChunkProcessor(chunks=chunks, messages=[])
+        result = processor._calculate_usage_per_chunk(chunks=chunks)
+        assert result["prompt_tokens"] == 0
+        assert result["completion_tokens"] == 0

tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py

@@ -2175,6 +2175,283 @@ def test_translate_openai_response_to_anthropic_cache_tokens_from_prompt_tokens_
     assert anthropic_response["usage"]["cache_read_input_tokens"] == 30
 
 
+def test_translate_openai_usage_to_anthropic_cache_tokens_from_dict_details_with_integral_floats():
+    usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=50,
+        total_tokens=170,
+    )
+    usage.prompt_tokens_details = {
+        "cached_tokens": 30.0,
+        "cache_write_tokens": 20.0,
+    }
+
+    anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
+        usage
+    )
+
+    assert anthropic_usage["input_tokens"] == 70
+    assert anthropic_usage["output_tokens"] == 50
+    assert anthropic_usage["cache_read_input_tokens"] == 30
+    assert anthropic_usage["cache_creation_input_tokens"] == 20
+
+
+def test_translate_openai_usage_to_anthropic_ignores_fractional_cache_tokens():
+    usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=50,
+        total_tokens=170,
+    )
+    usage.prompt_tokens_details = {
+        "cached_tokens": 30.5,
+        "cache_creation_tokens": 20.25,
+    }
+
+    anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
+        usage
+    )
+
+    assert anthropic_usage["input_tokens"] == 120
+    assert anthropic_usage["output_tokens"] == 50
+    assert "cache_read_input_tokens" not in anthropic_usage
+    assert "cache_creation_input_tokens" not in anthropic_usage
+
+
+def test_translate_openai_usage_to_anthropic_ignores_bool_cache_tokens():
+    usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=50,
+        total_tokens=170,
+    )
+    usage.cache_read_input_tokens = True
+    usage.cache_creation_input_tokens = True
+
+    anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
+        usage
+    )
+
+    assert anthropic_usage["input_tokens"] == 120
+    assert anthropic_usage["output_tokens"] == 50
+    assert "cache_read_input_tokens" not in anthropic_usage
+    assert "cache_creation_input_tokens" not in anthropic_usage
+
+
+def test_translate_openai_response_to_anthropic_cache_creation_from_prompt_tokens_details():
+    from litellm.types.utils import PromptTokensDetailsWrapper
+
+    usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=50,
+        total_tokens=170,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            cached_tokens=30,
+            cache_creation_tokens=20,
+        ),
+    )
+
+    response = ModelResponse(
+        id="test-id",
+        choices=[
+            Choices(
+                index=0,
+                finish_reason="stop",
+                message=Message(
+                    role="assistant",
+                    content="Test response",
+                ),
+            )
+        ],
+        model="gpt-4o-2024-08-06",
+        usage=usage,
+    )
+
+    adapter = LiteLLMAnthropicMessagesAdapter()
+    anthropic_response = adapter.translate_openai_response_to_anthropic(
+        response=response,
+        tool_name_mapping=None,
+    )
+
+    assert anthropic_response["usage"]["input_tokens"] == 70
+    assert anthropic_response["usage"]["output_tokens"] == 50
+    assert anthropic_response["usage"]["cache_read_input_tokens"] == 30
+    assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20
+
+
+def test_translate_openai_response_to_anthropic_cache_tokens_from_usage_fields():
+    usage = Usage(prompt_tokens=120, completion_tokens=50, total_tokens=170)
+    usage.cache_read_input_tokens = 30
+    usage.cache_creation_input_tokens = 20
+
+    response = ModelResponse(
+        id="test-id",
+        choices=[
+            Choices(
+                index=0,
+                finish_reason="stop",
+                message=Message(
+                    role="assistant",
+                    content="Test response",
+                ),
+            )
+        ],
+        model="claude-3-sonnet-20240229",
+        usage=usage,
+    )
+
+    adapter = LiteLLMAnthropicMessagesAdapter()
+    anthropic_response = adapter.translate_openai_response_to_anthropic(
+        response=response,
+        tool_name_mapping=None,
+    )
+
+    assert anthropic_response["usage"]["input_tokens"] == 70
+    assert anthropic_response["usage"]["output_tokens"] == 50
+    assert anthropic_response["usage"]["cache_read_input_tokens"] == 30
+    assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20
+
+
+def test_translate_openai_response_to_anthropic_cache_tokens_from_private_usage_fields():
+    usage = Usage(prompt_tokens=120, completion_tokens=50, total_tokens=170)
+
+    response = ModelResponse(
+        id="test-id",
+        choices=[
+            Choices(
+                index=0,
+                finish_reason="stop",
+                message=Message(
+                    role="assistant",
+                    content="Test response",
+                ),
+            )
+        ],
+        model="claude-3-sonnet-20240229",
+        usage=usage,
+    )
+    response.usage._cache_read_input_tokens = 30
+    response.usage._cache_creation_input_tokens = 20
+
+    adapter = LiteLLMAnthropicMessagesAdapter()
+    anthropic_response = adapter.translate_openai_response_to_anthropic(
+        response=response,
+        tool_name_mapping=None,
+    )
+
+    assert anthropic_response["usage"]["input_tokens"] == 70
+    assert anthropic_response["usage"]["output_tokens"] == 50
+    assert anthropic_response["usage"]["cache_read_input_tokens"] == 30
+    assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20
+
+
+def test_translate_streaming_openai_response_to_anthropic_cache_tokens_from_prompt_tokens_details():
+    from litellm.types.utils import PromptTokensDetailsWrapper
+
+    usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=50,
+        total_tokens=170,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            cached_tokens=30,
+            cache_creation_tokens=20,
+        ),
+    )
+    response = ModelResponseStream(
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(),
+                finish_reason="stop",
+            )
+        ],
+        usage=usage,
+    )
+
+    adapter = LiteLLMAnthropicMessagesAdapter()
+    message_delta = adapter.translate_streaming_openai_response_to_anthropic(
+        response=response,
+        current_content_block_index=0,
+    )
+
+    assert message_delta["usage"]["input_tokens"] == 70
+    assert message_delta["usage"]["output_tokens"] == 50
+    assert message_delta["usage"]["cache_read_input_tokens"] == 30
+    assert message_delta["usage"]["cache_creation_input_tokens"] == 20
+
+
+def test_translate_streaming_openai_response_to_anthropic_cache_tokens_from_hidden_params_usage():
+    from litellm.types.utils import PromptTokensDetailsWrapper
+
+    usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=50,
+        total_tokens=170,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            cached_tokens=30,
+            cache_creation_tokens=20,
+        ),
+    )
+    response = ModelResponseStream(
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(),
+                finish_reason="stop",
+            )
+        ],
+    )
+    response._hidden_params = {"usage": usage}
+
+    adapter = LiteLLMAnthropicMessagesAdapter()
+    message_delta = adapter.translate_streaming_openai_response_to_anthropic(
+        response=response,
+        current_content_block_index=0,
+    )
+
+    assert message_delta["usage"]["input_tokens"] == 70
+    assert message_delta["usage"]["output_tokens"] == 50
+    assert message_delta["usage"]["cache_read_input_tokens"] == 30
+    assert message_delta["usage"]["cache_creation_input_tokens"] == 20
+
+
+def test_translate_streaming_openai_response_to_anthropic_cache_tokens_with_applied_edits():
+    from litellm.types.utils import PromptTokensDetailsWrapper
+
+    usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=50,
+        total_tokens=170,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            cached_tokens=30,
+            cache_creation_tokens=20,
+        ),
+    )
+    response = ModelResponseStream(
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(),
+                finish_reason="stop",
+            )
+        ],
+        usage=usage,
+    )
+
+    adapter = LiteLLMAnthropicMessagesAdapter()
+    message_delta = adapter.translate_streaming_openai_response_to_anthropic(
+        response=response,
+        current_content_block_index=0,
+        applied_edits=[{"type": "compact_20260112"}],
+    )
+
+    assert message_delta["usage"]["input_tokens"] == 70
+    assert message_delta["usage"]["output_tokens"] == 50
+    assert message_delta["usage"]["cache_read_input_tokens"] == 30
+    assert message_delta["usage"]["cache_creation_input_tokens"] == 20
+    assert message_delta["context_management"]["applied_edits"][0]["type"] == (
+        "compact_20260112"
+    )
+
+
 # =====================================================================
 # Web Search Tool Transformation Tests
 # =====================================================================

tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py

@@ -24,6 +24,7 @@ from litellm.types.utils import (
     Message,
     ModelResponse,
     ModelResponseStream,
+    PromptTokensDetailsWrapper,
     StreamingChoices,
     Usage,
 )
@@ -88,6 +89,59 @@ def test_fake_stream_usage_preserved():
     assert message_delta["usage"]["input_tokens"] == 10
 
 
+def test_delayed_usage_chunk_preserves_cache_tokens():
+    usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=5,
+        total_tokens=125,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            cached_tokens=30,
+            cache_creation_tokens=20,
+        ),
+    )
+    chunks = [
+        ModelResponseStream(
+            choices=[
+                StreamingChoices(
+                    index=0,
+                    delta=Delta(content="Two."),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ModelResponseStream(
+            choices=[
+                StreamingChoices(
+                    index=0,
+                    delta=Delta(),
+                    finish_reason="stop",
+                )
+            ],
+        ),
+        ModelResponseStream(
+            choices=[
+                StreamingChoices(
+                    index=0,
+                    delta=Delta(),
+                    finish_reason=None,
+                )
+            ],
+            usage=usage,
+        ),
+    ]
+    wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="gpt-4o")
+    events = list(wrapper)
+
+    message_delta = next(
+        event for event in events if event.get("type") == "message_delta"
+    )
+
+    assert message_delta["usage"]["input_tokens"] == 70
+    assert message_delta["usage"]["output_tokens"] == 5
+    assert message_delta["usage"]["cache_read_input_tokens"] == 30
+    assert message_delta["usage"]["cache_creation_input_tokens"] == 20
+
+
 def test_splitter_passes_through_non_combined_chunks():
     """A chunk with content but no finish_reason is not split."""
     chunk = ModelResponseStream(

tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py

@@ -30,7 +30,9 @@ from litellm.types.utils import (
     ChatCompletionDeltaToolCall,
     Delta,
     Function,
+    PromptTokensDetailsWrapper,
     StreamingChoices,
+    Usage,
 )
 
 
@@ -107,6 +109,34 @@ def _input_json_deltas(events: List[dict]) -> List[str]:
     ]
 
 
+def test_held_stop_reason_usage_merge_preserves_openai_cache_token_details():
+    """OpenAI-compatible usage chunks carry cache reads in prompt_tokens_details."""
+    wrapper = AnthropicStreamWrapper(completion_stream=iter([]), model="claude-x")
+    wrapper.holding_stop_reason_chunk = {
+        "type": "message_delta",
+        "delta": {"stop_reason": "end_turn"},
+        "usage": {"input_tokens": 0, "output_tokens": 0},
+    }
+
+    usage_chunk = MagicMock()
+    usage_chunk.usage = Usage(
+        prompt_tokens=120,
+        completion_tokens=50,
+        total_tokens=170,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            cached_tokens=30,
+            cache_creation_tokens=20,
+        ),
+    )
+
+    merged_chunk = wrapper._merge_usage_into_held_stop_reason_chunk(usage_chunk)
+
+    assert merged_chunk["usage"]["input_tokens"] == 70
+    assert merged_chunk["usage"]["output_tokens"] == 50
+    assert merged_chunk["usage"]["cache_read_input_tokens"] == 30
+    assert merged_chunk["usage"]["cache_creation_input_tokens"] == 20
+
+
 def test_first_text_delta_after_tool_use_is_not_dropped_sync():
     """A tool_use -> text transition (text resuming after a tool call) carries
     the resumed text's first token in the trigger chunk. Without the fix it was

tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py

@@ -159,6 +159,59 @@ class TestGetAnthropicHeaders:
         assert "authorization" not in headers
         assert "anthropic-dangerous-direct-browser-access" not in headers
 
+    def test_custom_api_base_uses_bearer_header(self):
+        """Custom api_base and non-standard API key should produce Authorization: Bearer header when opted in."""
+        from litellm.llms.anthropic.common_utils import AnthropicModelInfo
+
+        config = AnthropicModelInfo()
+        headers = config.get_anthropic_headers(
+            api_key="my-custom-ollama-token",
+            computer_tool_used=False,
+            prompt_caching_set=False,
+            pdf_used=False,
+            is_vertex_request=False,
+            api_base="https://ollama.com/",
+            use_bearer_for_custom_base=True,
+        )
+
+        assert headers["authorization"] == "Bearer my-custom-ollama-token"
+        assert "x-api-key" not in headers
+
+    def test_custom_api_base_uses_bearer_header_already_starts_with_bearer(self):
+        """If the key already starts with Bearer and Bearer opt-in is enabled, use it directly."""
+        from litellm.llms.anthropic.common_utils import AnthropicModelInfo
+
+        config = AnthropicModelInfo()
+        headers = config.get_anthropic_headers(
+            api_key="Bearer my-custom-ollama-token",
+            computer_tool_used=False,
+            prompt_caching_set=False,
+            pdf_used=False,
+            is_vertex_request=False,
+            api_base="https://ollama.com/",
+            use_bearer_for_custom_base=True,
+        )
+
+        assert headers["authorization"] == "Bearer my-custom-ollama-token"
+        assert "x-api-key" not in headers
+
+    def test_custom_api_base_uses_x_api_key_when_standard_key(self):
+        """If the key is standard sk-ant- key, use x-api-key even with custom api_base."""
+        from litellm.llms.anthropic.common_utils import AnthropicModelInfo
+
+        config = AnthropicModelInfo()
+        headers = config.get_anthropic_headers(
+            api_key=FAKE_REGULAR_KEY,
+            computer_tool_used=False,
+            prompt_caching_set=False,
+            pdf_used=False,
+            is_vertex_request=False,
+            api_base="https://ollama.com/",
+        )
+
+        assert headers["x-api-key"] == FAKE_REGULAR_KEY
+        assert "authorization" not in headers
+
     def test_oauth_includes_standard_headers(self):
         """OAuth path should still include standard Anthropic headers."""
         from litellm.llms.anthropic.common_utils import AnthropicModelInfo
@@ -242,6 +295,46 @@ class TestValidateEnvironmentOAuth:
 
         assert updated_headers["x-api-key"] == FAKE_REGULAR_KEY
         assert "authorization" not in updated_headers
+
+    def test_custom_api_base_via_param(self):
+        """validate_environment uses Bearer when use_bearer_for_custom_base is set in litellm_params."""
+        from litellm.llms.anthropic.common_utils import AnthropicModelInfo
+
+        config = AnthropicModelInfo()
+        headers = {}
+
+        updated_headers = config.validate_environment(
+            headers=headers,
+            model="claude-sonnet-4-5-20250929",
+            messages=[{"role": "user", "content": "Hello"}],
+            optional_params={},
+            litellm_params={"use_bearer_for_custom_base": True},
+            api_key="custom-api-key",
+            api_base="https://custom-gateway.com",
+        )
+
+        assert updated_headers["authorization"] == "Bearer custom-api-key"
+        assert "x-api-key" not in updated_headers
+
+    def test_custom_api_base_via_litellm_params(self):
+        """validate_environment uses Bearer when api_base and use_bearer_for_custom_base are in litellm_params."""
+        from litellm.llms.anthropic.common_utils import AnthropicModelInfo
+
+        config = AnthropicModelInfo()
+        headers = {}
+
+        updated_headers = config.validate_environment(
+            headers=headers,
+            model="claude-sonnet-4-5-20250929",
+            messages=[{"role": "user", "content": "Hello"}],
+            optional_params={},
+            litellm_params={"api_base": "https://custom-gateway.com", "use_bearer_for_custom_base": True},
+            api_key="custom-api-key",
+            api_base=None,
+        )
+
+        assert updated_headers["authorization"] == "Bearer custom-api-key"
+        assert "x-api-key" not in updated_headers
         assert "anthropic-dangerous-direct-browser-access" not in updated_headers
 
 
@@ -1004,6 +1097,20 @@ class TestGetAuthHeader:
             result = AnthropicModelInfo.get_auth_header()
             assert result == {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"}
 
+    def test_custom_api_base_get_auth_header_uses_bearer(self):
+        """Non-standard API key and custom api_base returns Bearer when use_bearer_for_custom_base=True."""
+        from litellm.llms.anthropic.common_utils import AnthropicModelInfo
+
+        result = AnthropicModelInfo.get_auth_header(api_key="my-custom-key", api_base="https://custom-gateway.com", use_bearer_for_custom_base=True)
+        assert result == {"authorization": "Bearer my-custom-key"}
+
+    def test_custom_api_base_get_auth_header_uses_x_api_key_when_standard(self):
+        """Standard sk-ant- key with custom api_base should still return x-api-key."""
+        from litellm.llms.anthropic.common_utils import AnthropicModelInfo
+
+        result = AnthropicModelInfo.get_auth_header(api_key=FAKE_REGULAR_KEY, api_base="https://custom-gateway.com")
+        assert result == {"x-api-key": FAKE_REGULAR_KEY}
+
 
 class TestGetApiBaseFallbackChain:
     """Tests for AnthropicModelInfo.get_api_base() fallback to ANTHROPIC_BASE_URL."""