refactor(rust): run OCR through a route-neutral callback contract and a legacy Logging adapter

Extracted from #41733 without the router loop, the cache machine layer, streaming, or the
error, timeout and route-pruning work that moved to #41745

litellm-callbacks holds the contract a native call and its host share: Machine, HostOp,
CallEvent, the in-process run loop, and Passthrough, which is built only by comparing the
caller's inputs with the body the route sends, so a route can never mark a key it rewrote.
litellm-host-python (formerly python-interop) owns the CPython driver and the Execution
handle, and litellm-callbacks-legacy is the @client wrapper as the native call sees it:
function_setup, the deployment hooks, pre_call and post_call, the success and failure fan-out
and the deferred proxy release. OCR is the one route on it, and the old core and bridge
lifecycles are gone

The passthrough rule is the structural fix for the bug #41719 patched in core and #41716
reworks: an inlined remote document no longer counts as the caller's value, so the legacy
adapter never hands the caller's URL back into the body. core/tests/ocr/passthrough.rs pins
it for every route and document source, including that unchanged values stay passthrough,
and callbacks-legacy/tests/payload.rs pins the adapter side with a real pre_call callback

Python OCR integration tests that only exercised core behavior now live as Rust tests, so
tests/test_litellm_rust keeps the cases that need the full Python stack

litellm-rust/Cargo.lock

@@ -2001,6 +2001,26 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "litellm-callbacks"
+version = "0.1.0"
+dependencies = [
+ "rstest",
+ "serde_json",
+ "tokio",
+]
+
+[[package]]
+name = "litellm-callbacks-legacy"
+version = "0.1.0"
+dependencies = [
+ "litellm-callbacks",
+ "litellm-host-python",
+ "pyo3",
+ "rstest",
+ "serde_json",
+]
+
 [[package]]
 name = "litellm-core"
 version = "0.1.0"
@@ -2015,6 +2035,7 @@ dependencies = [
  "litellm-auth-aws",
  "litellm-auth-azure",
  "litellm-auth-gcp",
+ "litellm-callbacks",
  "litellm-framing",
  "litellm-providers",
  "mime_guess",
@@ -2022,6 +2043,7 @@ dependencies = [
  "rand 0.8.7",
  "reqwest 0.12.28",
  "rstest",
+ "rstest_reuse",
  "rustls 0.23.42",
  "rustls-native-certs",
  "serde",
@@ -2053,6 +2075,21 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "litellm-host-python"
+version = "0.1.0"
+dependencies = [
+ "futures-util",
+ "litellm-callbacks",
+ "pyo3",
+ "pyo3-async-runtimes",
+ "pythonize",
+ "rstest",
+ "serde",
+ "serde_json",
+ "tokio",
+]
+
 [[package]]
 name = "litellm-providers"
 version = "0.1.0"
@@ -2073,29 +2110,18 @@ dependencies = [
  "criterion",
  "futures-util",
  "litellm-auth",
+ "litellm-callbacks-legacy",
  "litellm-core",
- "litellm-python-interop",
+ "litellm-host-python",
  "litellm-token-counter",
  "pyo3",
  "pyo3-async-runtimes",
  "rstest",
- "serde",
  "serde_json",
  "tokio",
  "tokio-tungstenite",
 ]
 
-[[package]]
-name = "litellm-python-interop"
-version = "0.1.0"
-dependencies = [
- "pyo3",
- "pythonize",
- "rstest",
- "serde",
- "serde_json",
-]
-
 [[package]]
 name = "litellm-token-counter"
 version = "0.1.0"
@@ -3009,6 +3035,17 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "rstest_reuse"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3a8fb4672e840a587a66fc577a5491375df51ddb88f2a2c2a792598c326fe14"
+dependencies = [
+ "quote",
+ "rand 0.8.7",
+ "syn 2.0.119",
+]
+
 [[package]]
 name = "rustc-hash"
 version = "2.1.3"

litellm-rust/Cargo.toml

@@ -9,8 +9,9 @@ license = "MIT"
 repository = "https://github.com/BerriAI/litellm"
 
 [workspace.dependencies]
-bytes = "1"
 litellm-core = { path = "crates/core" }
+litellm-callbacks = { path = "crates/callbacks" }
+litellm-callbacks-legacy = { path = "crates/callbacks-legacy" }
 litellm-framing = { path = "crates/framer" }
 litellm-auth = { path = "crates/auth" }
 litellm-auth-aws = { path = "crates/auth-aws" }
@@ -20,13 +21,16 @@ litellm-providers = { path = "crates/providers" }
 litellm-cache = { path = "crates/cache" }
 litellm-cache-memory = { path = "crates/cache-memory" }
 litellm-token-counter = { path = "crates/token-counter" }
-litellm-python-interop = { path = "crates/python-interop" }
+litellm-host-python = { path = "crates/host-python" }
+
+bytes = "1"
 pyo3 = "0.29.2"
 pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
 pythonize = "0.29.0"
 rand = "0.8"
 reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
 rstest = "0.26.1"
+rstest_reuse = "0.7.0"
 rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
 rustls-native-certs = "0.8"
 serde = { version = "1.0", features = ["derive"] }

litellm-rust/crates/auth-azure/src/resolve.rs

@@ -657,4 +657,49 @@ mod tests {
 
         assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2));
     }
+
+    #[derive(Debug)]
+    struct CallerToken(&'static str);
+
+    impl litellm_auth::TokenProvider for CallerToken {
+        fn acquire(&self) -> litellm_auth::TokenFuture<'_> {
+            Box::pin(async move {
+                Ok(ResolvedCredential::AccessToken {
+                    token: SecretValue::new(self.0),
+                    expires_on: None,
+                })
+            })
+        }
+    }
+
+    fn caller_inputs(token: &'static str) -> AzureAuthInputs {
+        let params = json!({"azure_ad_token": "static-token"});
+        AzureAuthInputs {
+            azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new(
+                CallerToken(token),
+            ))),
+            ..AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap()
+        }
+    }
+
+    #[tokio::test]
+    async fn caller_token_is_chosen_over_supplied_static_token() {
+        let credential = AzureAuthService::default()
+            .get_azure_ad_token(&caller_inputs("caller-token"), &|_| None)
+            .await
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(credential.value().secret().expose(), "caller-token");
+    }
+
+    #[tokio::test]
+    async fn empty_caller_token_is_rejected() {
+        let error = AzureAuthService::default()
+            .get_azure_ad_token(&caller_inputs(""), &|_| None)
+            .await
+            .unwrap_err();
+
+        assert!(matches!(error, Error::EmptyAzureToken));
+    }
 }

litellm-rust/crates/auth/src/credential.rs

@@ -9,21 +9,6 @@ use crate::Error;
 
 use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
 
-pub fn credential_index(requested: &str, names: &[String]) -> Option<usize> {
-    names.iter().position(|name| name == requested)
-}
-
-pub fn credential_default_fields<'a>(
-    supplied: &[String],
-    credential_fields: &'a [String],
-) -> Vec<&'a str> {
-    credential_fields
-        .iter()
-        .filter(|name| !supplied.contains(name))
-        .map(String::as_str)
-        .collect()
-}
-
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum CredentialFileRef {
     Path(PathBuf),

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

@@ -47,7 +47,6 @@ impl<T> Sourced<T> {
 pub use credential::{
     CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
     CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
-    credential_default_fields, credential_index,
 };
 pub use error::Error;
 pub use http::{CredentialPlacement, RequestAuth};

litellm-rust/crates/callbacks-legacy/AGENTS.md

@@ -0,0 +1,17 @@
+- Target invariants, not completion claims
+- Keep this crate the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) plus the kwargs rewrites the wrapper makes on the way in (credential-name inheritance, the budget and retry-count limits)
+  - The driver in `litellm-host-python`, the routes and core see one `CallbackAdapter`; they never learn which Python objects consume a call
+  - `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the legacy path rewrites it (setup, deployment hook, prepare) and the bound request object whose attributes back keywords the caller omitted; routes hand it over through `run_legacy_call` and keep no copy
+- `setup` decides once who owns the `Logging` instance and returns it as `CallSetup.bridge_owned`; `PythonLogger` carries it and nothing on the instance records it
+  - A logger the caller passed as `litellm_logging_obj` is caller-owned and observed in full, because the caller reads it after the call; the proxy is the live case
+  - A logger `function_setup` built for this call is bridge-owned, so each fan-out phase is skipped when `callbacks_needed` finds no registry, dynamic callback, `logger_fn` or debug switch for it; cost, timing and response metadata still run
+- Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation
+  - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view
+  - Re-alias every `passthrough_fields` body key to the caller's object before `pre_call`; a keyword wins over the request attribute even when it is an explicit `None`
+  - Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only
+  - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only facts that cross from the route are the prepared keyword view and `RequestContext.passthrough_fields`
+- Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts
+  - Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch
+  - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once
+  - Delivery follows the registry, not the callable's type: direct, awaited, executor-submitted, logging-worker and deferred paths stay distinct
+- Traverse every retained Python edge; `close` is idempotent and restores the correlation context once

litellm-rust/crates/callbacks-legacy/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "litellm-callbacks-legacy"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+autotests = false
+
+[dependencies]
+litellm-callbacks.workspace = true
+litellm-host-python.workspace = true
+pyo3.workspace = true
+
+[dev-dependencies]
+rstest.workspace = true
+serde_json.workspace = true

litellm-rust/crates/callbacks-legacy/src/adapter.rs

@@ -0,0 +1,385 @@
+//! The legacy `Logging` contract as one adapter: every event and interception the driver
+//! raises is answered with the same `Logging` calls, in the same order, as the Python
+//! `@client` path makes them.
+
+use litellm_callbacks::event::{CallEvent, FailureOrigin, RequestContext, Timing, WireRequest};
+use litellm_host_python::{
+    AdapterStep, CallbackAdapter, PublicValue, from_py, missing_state, to_py,
+};
+use pyo3::{
+    exceptions::{PyBaseException, PyException},
+    gc::{PyTraverseError, PyVisit},
+    prelude::*,
+    types::PyDict,
+};
+
+use crate::{
+    DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
+    deferred::{PendingLogging, PendingSuccess},
+    finalize, is_internal_call, prepare, setup,
+};
+
+/// What the legacy contract needs to know about the route it is logging.
+#[derive(Clone, Copy, Debug)]
+pub struct LegacySurface {
+    pub call_type: &'static str,
+    /// What `Logging.pre_call` is told the input was.
+    pub input_description: &'static str,
+}
+
+enum Pending {
+    DeploymentPreCall,
+    DeploymentPostCall,
+    DeploymentFailure,
+    AsyncFailure,
+}
+
+pub struct LegacyLogging {
+    surface: LegacySurface,
+    call: PublicCall,
+    logger: Option<PythonLogger>,
+    start: Py<PyAny>,
+    end: Option<Py<PyAny>>,
+    response: Option<Py<PyAny>>,
+    error: Option<Py<PyBaseException>>,
+    body: Option<Py<PyDict>>,
+    headers: Option<Py<PyDict>>,
+    asynchronous: bool,
+    internal: bool,
+    pending: Option<Pending>,
+}
+
+fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
+    py.import("datetime")?
+        .getattr("datetime")?
+        .call_method1("fromtimestamp", (epoch_seconds,))
+        .map(Bound::unbind)
+}
+
+fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {
+    !error.is_instance_of::<PyException>(py)
+}
+
+impl LegacyLogging {
+    pub fn new(
+        py: Python<'_>,
+        surface: LegacySurface,
+        call: PublicCall,
+        asynchronous: bool,
+    ) -> Self {
+        Self {
+            surface,
+            call,
+            logger: None,
+            start: py.None(),
+            end: None,
+            response: None,
+            error: None,
+            body: None,
+            headers: None,
+            asynchronous,
+            internal: false,
+            pending: None,
+        }
+    }
+
+    /// Deployment hooks are awaited, and Python's synchronous `@client` wrapper never
+    /// runs them.
+    fn deployment_hooks(&self, py: Python<'_>) -> PyResult<bool> {
+        Ok(self.asynchronous && DeploymentHooks::needed(py)?)
+    }
+
+    fn logger(&self) -> PyResult<&PythonLogger> {
+        self.logger.as_ref().ok_or_else(|| {
+            pyo3::exceptions::PyRuntimeError::new_err("call logging is not initialized")
+        })
+    }
+
+    fn prepare(&mut self, py: Python<'_>) -> PyResult<AdapterStep> {
+        let prepared = prepare(py, self.call.kwargs().bind(py), self.logger()?)?.unbind();
+        self.call.set_kwargs(prepared);
+        Ok(AdapterStep::Arguments(self.call.kwargs().clone_ref(py)))
+    }
+
+    fn finalize(&mut self, py: Python<'_>) -> PyResult<AdapterStep> {
+        finalize(
+            py,
+            &self.response,
+            self.logger()?,
+            self.call.kwargs(),
+            &self.start,
+            &self.end,
+        )?;
+        self.response
+            .as_ref()
+            .map(|response| AdapterStep::Response(response.clone_ref(py)))
+            .ok_or_else(missing_state)
+    }
+
+    fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> {
+        match self.try_dispatch_success(py) {
+            Err(error) if error.is_instance_of::<PyException>(py) => {
+                error.write_unraisable(py, self.logger.as_ref().map(|logger| logger.object(py)));
+                Ok(())
+            }
+            result => result,
+        }
+    }
+
+    fn try_dispatch_success(&self, py: Python<'_>) -> PyResult<()> {
+        let logger = self.logger()?;
+        let pending = || PendingSuccess {
+            logger: logger.clone_ref(py),
+            response: self.response.as_ref().map(|value| value.clone_ref(py)),
+            start: self.start.clone_ref(py),
+            end: self.end.as_ref().map(|value| value.clone_ref(py)),
+        };
+        if !self.asynchronous {
+            return pending().sync(py);
+        }
+        if !self.internal
+            && self
+                .call
+                .kwargs()
+                .bind(py)
+                .get_item("fallbacks")?
+                .is_none_or(|value| value.is_none())
+        {
+            if !logger.callbacks_needed(py, "async_success")? {
+                logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?;
+            } else if logger.defers_async_logging(py) {
+                let pending = Py::new(
+                    py,
+                    PendingLogging {
+                        pending: Some(pending()),
+                    },
+                )?;
+                logger.defer_success(py, pending.bind(py).as_any())?;
+            } else {
+                pending().asynchronous(py)?;
+            }
+        }
+        logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
+    }
+
+    /// The sync failure handler, then the async one for async calls. Ordinary handler
+    /// errors never replace the selected failure or suppress the other family; a
+    /// cancellation does end the call.
+    fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult<AdapterStep> {
+        let (Some(logger), Some(error)) = (&self.logger, &self.error) else {
+            return Ok(AdapterStep::Done);
+        };
+        if self.asynchronous && self.internal {
+            return Ok(AdapterStep::Done);
+        }
+        if let Err(failure) = logger.failure(py, error, &self.start, &self.end, false)
+            && is_cancellation(py, &failure)
+        {
+            return Err(failure);
+        }
+        if !self.asynchronous {
+            return Ok(AdapterStep::Done);
+        }
+        match logger.failure(py, error, &self.start, &self.end, true) {
+            Ok(Some(awaitable)) => {
+                self.pending = Some(Pending::AsyncFailure);
+                Ok(AdapterStep::Await(awaitable))
+            }
+            Ok(None) => Ok(AdapterStep::Done),
+            Err(failure) if is_cancellation(py, &failure) => Err(failure),
+            Err(_) => Ok(AdapterStep::Done),
+        }
+    }
+}
+
+impl CallbackAdapter for LegacyLogging {
+    fn begin(
+        &mut self,
+        py: Python<'_>,
+        arguments: Py<PyDict>,
+        started_at: f64,
+    ) -> PyResult<AdapterStep> {
+        self.call.set_kwargs(arguments);
+        self.start = datetime(py, started_at)?;
+        self.internal = is_internal_call(py)?;
+        let result = setup(
+            py,
+            self.surface.call_type,
+            self.call.args(),
+            self.call.kwargs(),
+            &self.start,
+            self.asynchronous,
+        )?;
+        self.logger = Some(result.logger()?);
+        self.call.set_kwargs(result.kwargs()?);
+        if self.deployment_hooks(py)? {
+            self.pending = Some(Pending::DeploymentPreCall);
+            return Ok(AdapterStep::Await(DeploymentHooks::before_call(
+                py,
+                self.call.kwargs(),
+                self.surface.call_type,
+            )?));
+        }
+        self.prepare(py)
+    }
+
+    fn before_send(
+        &mut self,
+        py: Python<'_>,
+        wire: Box<WireRequest>,
+        context: &RequestContext,
+    ) -> PyResult<AdapterStep> {
+        let logger = self.logger()?;
+        logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?;
+        if !logger.callbacks_needed(py, "payload")? {
+            logger.record_api_call_start(py)?;
+            return Ok(AdapterStep::Wire(wire));
+        }
+        let body = to_py(py, &wire.body)?
+            .into_bound(py)
+            .cast_into::<PyDict>()?;
+        for name in context.passthrough_fields.iter() {
+            if let Some(value) = self.call.lookup(py, name)? {
+                body.set_item(name, value)?;
+            }
+        }
+        let headers = PyDict::new(py);
+        for (name, value) in &wire.headers {
+            headers.set_item(name, value)?;
+        }
+        self.body = Some(body.clone().unbind());
+        self.headers = Some(headers.clone().unbind());
+        let api_key = self.call.lookup(py, "api_key")?;
+        self.logger()?.pre_call(
+            py,
+            self.surface.input_description,
+            api_key.as_ref(),
+            &body,
+            &headers,
+            &wire.url,
+        )?;
+        let headers = headers
+            .iter()
+            .map(|(name, value)| Ok((name.extract::<String>()?, value.extract::<String>()?)))
+            .collect::<PyResult<Vec<_>>>()?;
+        Ok(AdapterStep::Wire(Box::new(WireRequest {
+            body: from_py(&body)?,
+            headers,
+            ..*wire
+        })))
+    }
+
+    fn after_success(
+        &mut self,
+        py: Python<'_>,
+        response: Py<PyAny>,
+        timing: Timing,
+    ) -> PyResult<AdapterStep> {
+        self.end = Some(datetime(py, timing.end_time)?);
+        self.response = Some(response);
+        if self.deployment_hooks(py)? {
+            self.pending = Some(Pending::DeploymentPostCall);
+            return Ok(AdapterStep::Await(DeploymentHooks::after_success(
+                py,
+                self.call.kwargs(),
+                &self.response,
+                self.surface.call_type,
+            )?));
+        }
+        self.finalize(py)
+    }
+
+    fn emit(
+        &mut self,
+        py: Python<'_>,
+        event: &CallEvent,
+        public: Option<PublicValue<'_>>,
+    ) -> PyResult<AdapterStep> {
+        match (event, public) {
+            (CallEvent::ResponseReceived { raw }, _) => {
+                let logger = self.logger()?;
+                if logger.callbacks_needed(py, "payload")? {
+                    logger.post_call(py, &raw.body, self.body.as_ref(), self.headers.as_ref())?;
+                }
+                Ok(AdapterStep::Done)
+            }
+            (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => {
+                self.end = Some(datetime(py, timing.end_time)?);
+                self.response = Some(response.clone_ref(py));
+                self.dispatch_success(py)?;
+                Ok(AdapterStep::Done)
+            }
+            (CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => {
+                self.end = Some(datetime(py, timing.end_time)?);
+                self.error = Some(error.clone_ref(py).into_value(py));
+                if *origin == FailureOrigin::Call
+                    && self.logger.is_some()
+                    && self.deployment_hooks(py)?
+                {
+                    let error = self.error.as_ref().ok_or_else(missing_state)?;
+                    self.pending = Some(Pending::DeploymentFailure);
+                    return Ok(AdapterStep::Await(DeploymentHooks::after_failure(
+                        py,
+                        self.call.kwargs(),
+                        error,
+                        self.surface.call_type,
+                    )?));
+                }
+                self.dispatch_failure(py)
+            }
+            _ => Err(missing_state()),
+        }
+    }
+
+    fn resume(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<AdapterStep> {
+        match self.pending.take().ok_or_else(missing_state)? {
+            Pending::DeploymentPreCall => {
+                self.call
+                    .set_kwargs(result?.into_bound(py).cast_into::<PyDict>()?.unbind());
+                self.prepare(py)
+            }
+            Pending::DeploymentPostCall => {
+                self.response = Some(result?);
+                self.finalize(py)
+            }
+            Pending::DeploymentFailure => self.dispatch_failure(py),
+            Pending::AsyncFailure => match result {
+                Err(failure) if is_cancellation(py, &failure) => Err(failure),
+                _ => Ok(AdapterStep::Done),
+            },
+        }
+    }
+
+    fn close(&mut self, py: Python<'_>) {
+        if let Some(logger) = self.logger.take()
+            && let Err(error) = logger.restore_context(py)
+        {
+            error.write_unraisable(py, None);
+        }
+        self.body = None;
+        self.headers = None;
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        self.call.traverse(visit)?;
+        if let Some(logger) = &self.logger {
+            logger.traverse(visit)?;
+        }
+        visit.call(&self.start)?;
+        visit.call(&self.end)?;
+        visit.call(&self.response)?;
+        visit.call(&self.error)?;
+        visit.call(&self.body)?;
+        visit.call(&self.headers)
+    }
+}
+
+#[cfg(test)]
+#[path = "../tests/deployment_hooks.rs"]
+mod deployment_hooks_tests;
+#[cfg(test)]
+#[path = "../tests/payload.rs"]
+mod payload_tests;
+#[cfg(test)]
+#[path = "../tests/terminal.rs"]
+mod terminal_tests;

litellm-rust/crates/callbacks-legacy/src/call.rs

@@ -0,0 +1,179 @@
+//! The caller's public call as the legacy `Logging` contract sees it. Legacy callbacks
+//! receive these exact objects and may mutate them, so the call keeps them for its whole
+//! lifetime. No other callback host has that obligation, which is why nothing outside
+//! this crate holds them.
+
+use litellm_callbacks::{machine::Machine, route::Route};
+use litellm_host_python::{RouteHost, run_call};
+use pyo3::{
+    gc::{PyTraverseError, PyVisit},
+    prelude::*,
+    types::{PyDict, PyTuple},
+};
+
+use crate::{LegacyLogging, LegacySurface};
+
+pub struct PublicCall {
+    args: Py<PyTuple>,
+    kwargs: Py<PyDict>,
+    request: Py<PyAny>,
+}
+
+impl PublicCall {
+    /// Copies the keyword arguments once, so the legacy path's rewrites never reach the
+    /// caller's own dict while every value keeps its identity.
+    pub fn capture(
+        request: &Bound<'_, PyAny>,
+        args: &Bound<'_, PyTuple>,
+        kwargs: &Bound<'_, PyDict>,
+    ) -> PyResult<Self> {
+        Ok(Self {
+            args: args.clone().unbind(),
+            kwargs: kwargs.copy()?.unbind(),
+            request: request.clone().unbind(),
+        })
+    }
+
+    pub(crate) fn args(&self) -> &Py<PyTuple> {
+        &self.args
+    }
+
+    /// The keyword view the legacy path currently reads: the caller's copy until
+    /// `function_setup`, then each rewrite (setup, deployment hook, prepare) in turn.
+    pub(crate) fn kwargs(&self) -> &Py<PyDict> {
+        &self.kwargs
+    }
+
+    pub(crate) fn set_kwargs(&mut self, kwargs: Py<PyDict>) {
+        self.kwargs = kwargs;
+    }
+
+    pub(crate) fn lookup<'py>(
+        &self,
+        py: Python<'py>,
+        name: &str,
+    ) -> PyResult<Option<Bound<'py, PyAny>>> {
+        lookup(self.kwargs.bind(py), self.request.bind(py), name)
+    }
+
+    pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.args)?;
+        visit.call(&self.kwargs)?;
+        visit.call(&self.request)
+    }
+}
+
+/// The caller's own object for a public argument, as every legacy reader resolves it: the
+/// keyword if given, even an explicit `None`, else the bound request's attribute. A route
+/// host projecting from the prepared keyword view uses the same rule, so the callbacks
+/// and the provider see one object per argument.
+pub fn lookup<'py>(
+    kwargs: &Bound<'py, PyDict>,
+    request: &Bound<'py, PyAny>,
+    name: &str,
+) -> PyResult<Option<Bound<'py, PyAny>>> {
+    if let Some(value) = kwargs.get_item(name)? {
+        return Ok(Some(value));
+    }
+    request.getattr_opt(name)
+}
+
+/// Runs one native call under the legacy `Logging` contract: the route host projects from
+/// the keyword view the contract prepares, and the contract observes the call.
+pub fn run_legacy_call<H, M>(
+    py: Python<'_>,
+    surface: LegacySurface,
+    call: PublicCall,
+    machine: M,
+    route: H,
+    asynchronous: bool,
+) -> PyResult<Py<PyAny>>
+where
+    H: RouteHost + 'static,
+    M: Machine<Route = H::Route, Complete = <H::Route as Route>::Response> + 'static,
+{
+    let arguments = call.kwargs.clone_ref(py);
+    run_call(
+        py,
+        machine,
+        route,
+        Box::new(LegacyLogging::new(py, surface, call, asynchronous)),
+        arguments,
+        asynchronous,
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn capture<'py>(py: Python<'py>, source: &std::ffi::CStr) -> (PublicCall, Bound<'py, PyDict>) {
+        let locals = PyDict::new(py);
+        py.run(source, Some(&locals), Some(&locals)).unwrap();
+        let request = locals.get_item("request").unwrap().unwrap();
+        let kwargs = locals
+            .get_item("kwargs")
+            .unwrap()
+            .unwrap()
+            .cast_into::<PyDict>()
+            .unwrap();
+        let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap();
+        (call, locals)
+    }
+
+    #[test]
+    fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() {
+        Python::initialize();
+        Python::attach(|py| {
+            let (call, locals) = capture(
+                py,
+                c"
+key = object()
+document = {'type': 'document_url'}
+class Request:
+    api_key = 'from-request'
+    api_base = 'from-request'
+    document = document
+request = Request()
+kwargs = {'api_key': key, 'api_base': None}
+",
+            );
+            let key = locals.get_item("key").unwrap().unwrap();
+            let document = locals.get_item("document").unwrap().unwrap();
+            assert!(call.lookup(py, "api_key").unwrap().unwrap().is(&key));
+            assert!(call.lookup(py, "api_base").unwrap().unwrap().is_none());
+            assert!(call.lookup(py, "document").unwrap().unwrap().is(&document));
+            assert!(call.lookup(py, "model").unwrap().is_none());
+        });
+    }
+
+    #[test]
+    fn capture_copies_the_keyword_dict_without_copying_its_values() {
+        Python::initialize();
+        Python::attach(|py| {
+            let (call, locals) = capture(
+                py,
+                c"
+pages = [0]
+class Request:
+    pass
+request = Request()
+kwargs = {'pages': pages}
+",
+            );
+            let caller = locals
+                .get_item("kwargs")
+                .unwrap()
+                .unwrap()
+                .cast_into::<PyDict>()
+                .unwrap();
+            call.kwargs()
+                .bind(py)
+                .set_item("litellm_call_id", "call")
+                .unwrap();
+            assert!(!caller.contains("litellm_call_id").unwrap());
+            let pages = locals.get_item("pages").unwrap().unwrap();
+            assert!(call.lookup(py, "pages").unwrap().unwrap().is(&pages));
+        });
+    }
+}

litellm-rust/crates/callbacks-legacy/src/callbacks.rs

@@ -0,0 +1,404 @@
+//! Callback fan-out over litellm's `Logging` object: which callbacks are registered,
+//! the deferred and worker-submitted success paths, and the sync-callbacks-for-async-calls
+//! duplication. All of it expires with the legacy callback contract.
+
+use litellm_callbacks::event::{RequestContext, WireRequest};
+use litellm_host_python::to_py;
+use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
+
+use crate::logger::PythonLogger;
+
+pub trait LegacyCallbacks {
+    fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult<bool>;
+
+    /// `Logging.update_from_kwargs`: what the logger is told about the request it is
+    /// about to see, with consumed credentials redacted.
+    fn update_from_kwargs(
+        &self,
+        py: Python<'_>,
+        kwargs: &Py<PyDict>,
+        wire: &WireRequest,
+        context: &RequestContext,
+    ) -> PyResult<()>;
+
+    fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()>;
+
+    /// `Logging.pre_call`, or its payload-free shortcut when no input callback listens.
+    fn pre_call(
+        &self,
+        py: Python<'_>,
+        input: &str,
+        api_key: Option<&Bound<'_, PyAny>>,
+        body: &Bound<'_, PyDict>,
+        headers: &Bound<'_, PyDict>,
+        url: &str,
+    ) -> PyResult<()>;
+
+    /// `Logging.post_call`, or its payload-free shortcut when no input callback listens.
+    fn post_call(
+        &self,
+        py: Python<'_>,
+        original_response: &str,
+        body: Option<&Py<PyDict>>,
+        headers: Option<&Py<PyDict>>,
+    ) -> PyResult<()>;
+
+    fn defers_async_logging(&self, py: Python<'_>) -> bool;
+
+    fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>;
+
+    fn sync_success_for_async_call(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+    ) -> PyResult<()>;
+
+    fn failure(
+        &self,
+        py: Python<'_>,
+        error: &Py<PyBaseException>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+        asynchronous: bool,
+    ) -> PyResult<Option<Py<PyAny>>>;
+
+    fn submit_success(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+    ) -> PyResult<()>;
+
+    fn enqueue_success(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+    ) -> PyResult<()>;
+}
+
+impl LegacyCallbacks for PythonLogger {
+    fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult<bool> {
+        if !self.bridge_owned() {
+            return Ok(true);
+        }
+        py.import("litellm.rust_bridge.legacy_callbacks")?
+            .getattr("callbacks_needed")?
+            .call1((self.object(py), phase))?
+            .extract()
+    }
+
+    fn update_from_kwargs(
+        &self,
+        py: Python<'_>,
+        kwargs: &Py<PyDict>,
+        wire: &WireRequest,
+        context: &RequestContext,
+    ) -> PyResult<()> {
+        let secret_fields: Vec<&str> = context.secret_fields.iter().map(String::as_str).collect();
+        let update = PyDict::new(py);
+        update.set_item("kwargs", redact(py, kwargs.bind(py), &secret_fields)?)?;
+        update.set_item("model", &context.model)?;
+        update.set_item(
+            "optional_params",
+            redact(
+                py,
+                &to_py(py, &context.optional_params)?
+                    .into_bound(py)
+                    .cast_into::<PyDict>()?,
+                &secret_fields,
+            )?,
+        )?;
+        let params = PyDict::new(py);
+        params.set_item(
+            "litellm_call_id",
+            kwargs.bind(py).get_item("litellm_call_id")?,
+        )?;
+        params.set_item("api_base", &wire.url)?;
+        for name in ["logger_fn", "litellm_request_debug"] {
+            if let Some(value) = kwargs.bind(py).get_item(name)? {
+                params.set_item(name, value)?;
+            }
+        }
+        for name in custom_pricing_fields(py)? {
+            if let Some(value) = kwargs.bind(py).get_item(&name)?
+                && !value.is_none()
+            {
+                params.set_item(name, value)?;
+            }
+        }
+        update.set_item("litellm_params", params)?;
+        update.set_item("custom_llm_provider", &context.custom_llm_provider)?;
+        self.object(py)
+            .call_method("update_from_kwargs", (), Some(&update))?;
+        Ok(())
+    }
+
+    fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()> {
+        self.object(py).call_method0("record_api_call_start_time")?;
+        Ok(())
+    }
+
+    fn pre_call(
+        &self,
+        py: Python<'_>,
+        input: &str,
+        api_key: Option<&Bound<'_, PyAny>>,
+        body: &Bound<'_, PyDict>,
+        headers: &Bound<'_, PyDict>,
+        url: &str,
+    ) -> PyResult<()> {
+        let additional = PyDict::new(py);
+        additional.set_item("complete_input_dict", body)?;
+        additional.set_item("headers", headers)?;
+        additional.set_item("api_base", url)?;
+        let kwargs = PyDict::new(py);
+        kwargs.set_item("input", input)?;
+        kwargs.set_item("api_key", api_key)?;
+        kwargs.set_item("additional_args", &additional)?;
+        if self.callbacks_needed(py, "input")? {
+            self.object(py).call_method("pre_call", (), Some(&kwargs))?;
+        } else {
+            self.object(py)
+                .call_method("_pre_call", (), Some(&kwargs))?;
+            self.record_api_call_start(py)?;
+        }
+        Ok(())
+    }
+
+    fn post_call(
+        &self,
+        py: Python<'_>,
+        original_response: &str,
+        body: Option<&Py<PyDict>>,
+        headers: Option<&Py<PyDict>>,
+    ) -> PyResult<()> {
+        let additional = PyDict::new(py);
+        additional.set_item("complete_input_dict", body)?;
+        additional.set_item("headers", headers)?;
+        if self.callbacks_needed(py, "input")? {
+            let kwargs = PyDict::new(py);
+            kwargs.set_item("original_response", original_response)?;
+            kwargs.set_item("additional_args", &additional)?;
+            self.object(py)
+                .call_method("post_call", (), Some(&kwargs))?;
+        } else {
+            let response = py
+                .import("json")?
+                .call_method1("dumps", (original_response,))?;
+            self.object(py).call_method1(
+                "record_post_call",
+                (response, py.None(), py.None(), additional),
+            )?;
+        }
+        Ok(())
+    }
+    fn defers_async_logging(&self, py: Python<'_>) -> bool {
+        self.object(py)
+            .getattr("_defer_async_logging")
+            .is_ok_and(|value| value.is_truthy().unwrap_or(false))
+    }
+
+    fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> {
+        self.object(py).setattr("_native_pending_logging", pending)
+    }
+
+    fn sync_success_for_async_call(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+    ) -> PyResult<()> {
+        if !self.callbacks_needed(py, "sync_success_async")? {
+            return Ok(());
+        }
+        self.object(py).call_method1(
+            "handle_sync_success_callbacks_for_async_calls",
+            (response, start, end),
+        )?;
+        Ok(())
+    }
+
+    fn failure(
+        &self,
+        py: Python<'_>,
+        error: &Py<PyBaseException>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+        asynchronous: bool,
+    ) -> PyResult<Option<Py<PyAny>>> {
+        if !self.callbacks_needed(
+            py,
+            if asynchronous {
+                "async_failure"
+            } else {
+                "sync_failure"
+            },
+        )? {
+            py.import("litellm.rust_bridge.legacy_callbacks")?
+                .getattr("failure_bookkeeping")?
+                .call1((self.object(py), error, start, end, asynchronous))?;
+            return Ok(None);
+        }
+        let trace = py
+            .import("traceback")?
+            .getattr("format_exception")?
+            .call1((error,))?;
+        let trace = pyo3::types::PyString::new(py, "").call_method1("join", (trace,))?;
+        let value = self.object(py).call_method1(
+            if asynchronous {
+                "async_failure_handler"
+            } else {
+                "failure_handler"
+            },
+            (error, trace, start, end),
+        )?;
+        Ok(asynchronous.then(|| value.unbind()))
+    }
+    fn submit_success(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+    ) -> PyResult<()> {
+        if !self.callbacks_needed(py, "sync_success")? {
+            return self.success_bookkeeping(py, response, start, end, false);
+        }
+        let context = py.import("contextvars")?.call_method0("copy_context")?;
+        py.import("litellm.litellm_core_utils.litellm_logging")?
+            .getattr("executor")?
+            .call_method1(
+                "submit",
+                (
+                    context.getattr("run")?,
+                    self.object(py).getattr("success_handler")?,
+                    response,
+                    start,
+                    end,
+                ),
+            )?;
+        Ok(())
+    }
+
+    fn enqueue_success(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+    ) -> PyResult<()> {
+        if !self.callbacks_needed(py, "async_success")? {
+            return self.success_bookkeeping(py, response, start, end, true);
+        }
+        let context = py.import("contextvars")?.call_method0("copy_context")?;
+        let worker = py
+            .import("litellm.litellm_core_utils.logging_worker")?
+            .getattr("GLOBAL_LOGGING_WORKER")?
+            .getattr("ensure_initialized_and_enqueue")?;
+        let coroutine = self
+            .object(py)
+            .call_method1("async_success_handler", (response, start, end))?;
+        let enqueue = context.call_method1("run", (worker, &coroutine));
+        if enqueue.is_err()
+            && let Err(error) = coroutine.call_method0("close")
+        {
+            error.write_unraisable(py, Some(&coroutine));
+        }
+        enqueue.map(|_| ())
+    }
+}
+
+fn custom_pricing_fields(py: Python<'_>) -> PyResult<Vec<String>> {
+    py.import("litellm.types.utils")?
+        .getattr("CustomPricingLiteLLMParams")?
+        .getattr("model_fields")?
+        .cast_into::<PyDict>()?
+        .keys()
+        .iter()
+        .map(|name| name.extract::<String>())
+        .collect()
+}
+
+fn redact(
+    py: Python<'_>,
+    params: &Bound<'_, PyDict>,
+    secret_fields: &[&str],
+) -> PyResult<Py<PyDict>> {
+    let redacted = PyDict::new(py);
+    for (name, value) in params {
+        let name = name.extract::<String>()?;
+        if name == "proxy_server_request" {
+            continue;
+        }
+        if secret_fields.contains(&name.as_str()) {
+            redacted.set_item(name, "****")?;
+        } else {
+            redacted.set_item(name, value)?;
+        }
+    }
+    Ok(redacted.unbind())
+}
+
+/// Proxy-internal calls skip the legacy success fan-out.
+pub fn is_internal_call(py: Python<'_>) -> PyResult<bool> {
+    py.import("litellm._internal_context")?
+        .getattr("is_internal_call")?
+        .call_method0("get")?
+        .extract()
+}
+
+#[cfg(test)]
+mod tests {
+    use pyo3::types::PyDict;
+
+    use super::*;
+
+    fn logger_whose_registries_need_no_input(py: Python<'_>, bridge_owned: bool) -> PythonLogger {
+        let locals = PyDict::new(py);
+        py.run(
+            c"
+import sys
+import types
+for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'):
+    sys.modules.setdefault(name, types.ModuleType(name))
+legacy = sys.modules['litellm.rust_bridge.legacy_callbacks']
+legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True)
+class Logger:
+    needed = {'input': False}
+logger = Logger()
+",
+            Some(&locals),
+            Some(&locals),
+        )
+        .unwrap();
+        PythonLogger::new(
+            locals.get_item("logger").unwrap().unwrap().unbind(),
+            bridge_owned,
+        )
+    }
+
+    #[test]
+    fn a_caller_owned_logger_is_observed_in_full() {
+        Python::initialize();
+        Python::attach(|py| {
+            let logger = logger_whose_registries_need_no_input(py, false);
+            assert!(logger.callbacks_needed(py, "input").unwrap());
+        });
+    }
+
+    #[test]
+    fn a_bridge_owned_logger_is_elided_where_no_registry_needs_it() {
+        Python::initialize();
+        Python::attach(|py| {
+            let logger = logger_whose_registries_need_no_input(py, true);
+            assert!(!logger.callbacks_needed(py, "input").unwrap());
+            assert!(logger.callbacks_needed(py, "payload").unwrap());
+        });
+    }
+}

litellm-rust/crates/callbacks-legacy/src/deferred.rs

@@ -0,0 +1,67 @@
+//! The proxy's deferred success release: the async success handler is queued only once
+//! the proxy accepts the response, and at most once.
+
+use pyo3::{exceptions::PyException, prelude::*};
+
+use crate::{LegacyCallbacks, PythonLogger};
+
+pub(crate) struct PendingSuccess {
+    pub(crate) logger: PythonLogger,
+    pub(crate) response: Option<Py<PyAny>>,
+    pub(crate) start: Py<PyAny>,
+    pub(crate) end: Option<Py<PyAny>>,
+}
+
+impl PendingSuccess {
+    pub(crate) fn sync(&self, py: Python<'_>) -> PyResult<()> {
+        self.logger
+            .submit_success(py, &self.response, &self.start, &self.end)
+    }
+
+    pub(crate) fn asynchronous(&self, py: Python<'_>) -> PyResult<()> {
+        self.logger
+            .enqueue_success(py, &self.response, &self.start, &self.end)
+    }
+}
+
+#[pyclass]
+pub(crate) struct PendingLogging {
+    pub(crate) pending: Option<PendingSuccess>,
+}
+
+#[pymethods]
+impl PendingLogging {
+    fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> {
+        let pending = slf.borrow_mut().pending.take();
+        if let Some(pending) = pending
+            && success
+        {
+            match pending.asynchronous(py) {
+                Err(error) if error.is_instance_of::<PyException>(py) => {
+                    error.write_unraisable(py, Some(pending.logger.object(py)));
+                }
+                result => return result,
+            }
+        }
+        Ok(())
+    }
+
+    fn __traverse__(&self, visit: pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> {
+        if let Some(pending) = &self.pending {
+            pending.logger.traverse(&visit)?;
+            visit.call(&pending.response)?;
+            visit.call(&pending.start)?;
+            visit.call(&pending.end)?;
+        }
+        Ok(())
+    }
+
+    fn __clear__(slf: &Bound<'_, Self>) {
+        let pending = slf.borrow_mut().pending.take();
+        drop(pending);
+    }
+}
+
+#[cfg(test)]
+#[path = "../tests/deferred.rs"]
+mod tests;

litellm-rust/crates/callbacks-legacy/src/lib.rs

@@ -0,0 +1,27 @@
+//! The legacy `@client` wrapper as the native call sees it: litellm's `Logging` object, the
+//! sync and async callback registries it fans out to, the deployment hooks, the deferred
+//! proxy release, and the kwargs rewrites the wrapper makes on the way in (credential-name
+//! inheritance, budget and retry-count limits). All of it sits behind one
+//! [`CallbackAdapter`](litellm_host_python::CallbackAdapter), so the driver, the routes and
+//! core never learn which Python object is on the other end.
+//!
+//! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`]
+//! is where those objects live, and [`run_legacy_call`] is how a route hands them over
+//! without keeping a copy.
+
+mod adapter;
+mod call;
+mod callbacks;
+mod deferred;
+mod logger;
+mod preparation;
+#[cfg(test)]
+#[path = "../tests/support.rs"]
+mod test_support;
+
+pub(crate) use adapter::LegacyLogging;
+pub use adapter::LegacySurface;
+pub use call::{PublicCall, lookup, run_legacy_call};
+pub(crate) use callbacks::{LegacyCallbacks, is_internal_call};
+pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup};
+pub(crate) use preparation::prepare;

litellm-rust/crates/callbacks-legacy/src/logger.rs

@@ -0,0 +1,236 @@
+use pyo3::{
+    exceptions::PyBaseException,
+    gc::{PyTraverseError, PyVisit},
+    prelude::*,
+    types::{PyDict, PyTuple},
+};
+
+/// The `Logging` instance one call fans out through, and who owns it. A logger the caller
+/// handed in is observed in full, because the caller reads it after the call; one this
+/// crate built through `function_setup` is elided wherever no registry needs it.
+pub struct PythonLogger {
+    object: Py<PyAny>,
+    bridge_owned: bool,
+}
+
+impl PythonLogger {
+    pub(crate) fn new(object: Py<PyAny>, bridge_owned: bool) -> Self {
+        Self {
+            object,
+            bridge_owned,
+        }
+    }
+
+    pub(crate) fn object<'py>(&self, py: Python<'py>) -> &Bound<'py, PyAny> {
+        self.object.bind(py)
+    }
+
+    pub(crate) fn bridge_owned(&self) -> bool {
+        self.bridge_owned
+    }
+
+    pub fn clone_ref(&self, py: Python<'_>) -> Self {
+        Self {
+            object: self.object.clone_ref(py),
+            bridge_owned: self.bridge_owned,
+        }
+    }
+
+    pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.object)
+    }
+
+    pub fn success_bookkeeping(
+        &self,
+        py: Python<'_>,
+        response: &Option<Py<PyAny>>,
+        start: &Py<PyAny>,
+        end: &Option<Py<PyAny>>,
+        asynchronous: bool,
+    ) -> PyResult<()> {
+        py.import("litellm.rust_bridge.legacy_callbacks")?
+            .getattr("success_bookkeeping")?
+            .call1((self.object(py), response, start, end, asynchronous))?;
+        Ok(())
+    }
+
+    pub fn restore_context(&self, py: Python<'_>) -> PyResult<()> {
+        py.import("litellm.utils")?
+            .getattr("_restore_correlation_context_if_supported")?
+            .call1((self.object(py),))?;
+        Ok(())
+    }
+}
+
+/// A bare Python object was not obtained from `setup`, so it is caller-owned.
+impl FromPyObject<'_, '_> for PythonLogger {
+    type Error = PyErr;
+
+    fn extract(object: Borrowed<'_, '_, PyAny>) -> PyResult<Self> {
+        Ok(Self::new(object.to_owned().unbind(), false))
+    }
+}
+
+pub struct SetupResult<'py>(Bound<'py, PyAny>);
+
+impl SetupResult<'_> {
+    pub fn logger(&self) -> PyResult<PythonLogger> {
+        let object = self.0.getattr("logger")?.unbind();
+        let bridge_owned = self.0.getattr("bridge_owned")?.extract()?;
+        Ok(PythonLogger::new(object, bridge_owned))
+    }
+
+    pub fn kwargs(&self) -> PyResult<Py<PyDict>> {
+        Ok(self.0.getattr("kwargs")?.extract()?)
+    }
+}
+
+pub fn setup<'py>(
+    py: Python<'py>,
+    call_type: &str,
+    args: &Py<PyTuple>,
+    kwargs: &Py<PyDict>,
+    start: &Py<PyAny>,
+    asynchronous: bool,
+) -> PyResult<SetupResult<'py>> {
+    py.import("litellm.rust_bridge.legacy_callbacks")?
+        .getattr("setup")?
+        .call1((call_type, args, kwargs, start, asynchronous))
+        .map(SetupResult)
+}
+
+pub fn finalize(
+    py: Python<'_>,
+    response: &Option<Py<PyAny>>,
+    logger: &PythonLogger,
+    kwargs: &Py<PyDict>,
+    start: &Py<PyAny>,
+    end: &Option<Py<PyAny>>,
+) -> PyResult<()> {
+    py.import("litellm.rust_bridge.legacy_callbacks")?
+        .getattr("finalize")?
+        .call1((response, logger.object(py), kwargs, start, end))?;
+    Ok(())
+}
+
+pub struct DeploymentHooks;
+
+impl DeploymentHooks {
+    pub fn needed(py: Python<'_>) -> PyResult<bool> {
+        py.import("litellm.rust_bridge.legacy_callbacks")?
+            .getattr("deployment_callbacks_needed")?
+            .call0()?
+            .extract()
+    }
+
+    pub fn before_call(
+        py: Python<'_>,
+        kwargs: &Py<PyDict>,
+        call_type: &str,
+    ) -> PyResult<Py<PyAny>> {
+        py.import("litellm.utils")?
+            .getattr("async_pre_call_deployment_hook")?
+            .call1((kwargs, call_type))
+            .map(Bound::unbind)
+    }
+
+    pub fn after_success(
+        py: Python<'_>,
+        kwargs: &Py<PyDict>,
+        response: &Option<Py<PyAny>>,
+        call_type: &str,
+    ) -> PyResult<Py<PyAny>> {
+        py.import("litellm.utils")?
+            .getattr("async_post_call_success_deployment_hook")?
+            .call1((kwargs, response, call_type))
+            .map(Bound::unbind)
+    }
+
+    pub fn after_failure(
+        py: Python<'_>,
+        kwargs: &Py<PyDict>,
+        error: &Py<PyBaseException>,
+        call_type: &str,
+    ) -> PyResult<Py<PyAny>> {
+        py.import("litellm.utils")?
+            .getattr("async_post_call_failure_deployment_hook")?
+            .call1((kwargs, error, call_type))
+            .map(Bound::unbind)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use pyo3::exceptions::PyTypeError;
+
+    use super::*;
+
+    #[test]
+    fn setup_fields_are_checked_lazily() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            py.run(
+                pyo3::ffi::c_str!(
+                    r#"
+reads = []
+class Logger:
+    def __getattribute__(self, name):
+        reads.append(name)
+        raise AssertionError('logger methods must remain lazy')
+logger = Logger()
+class Setup:
+    @property
+    def logger(self):
+        reads.append('logger')
+        return logger
+    @property
+    def bridge_owned(self):
+        reads.append('bridge_owned')
+        return True
+    @property
+    def kwargs(self):
+        reads.append('kwargs')
+        return []
+result = Setup()
+"#
+                ),
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+            let result = SetupResult(locals.get_item("result").unwrap().unwrap());
+            let logger = result.logger().unwrap();
+            assert!(
+                logger
+                    .object(py)
+                    .is(locals.get_item("logger").unwrap().unwrap())
+            );
+            assert!(logger.bridge_owned());
+            assert!(
+                result
+                    .kwargs()
+                    .unwrap_err()
+                    .is_instance_of::<PyTypeError>(py)
+            );
+            assert_eq!(
+                locals
+                    .get_item("reads")
+                    .unwrap()
+                    .unwrap()
+                    .extract::<Vec<String>>()
+                    .unwrap(),
+                ["logger", "bridge_owned", "kwargs"]
+            );
+        });
+    }
+
+    #[test]
+    fn a_logger_extracted_from_a_bare_object_is_caller_owned() {
+        Python::initialize();
+        Python::attach(|py| {
+            let logger: PythonLogger = py.None().into_bound(py).extract().unwrap();
+            assert!(!logger.bridge_owned());
+        });
+    }
+}

litellm-rust/crates/callbacks-legacy/src/preparation.rs

@@ -1,6 +1,7 @@
-use litellm_auth::{credential_default_fields, credential_index};
-use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyList};
+use pyo3::{
+    prelude::*,
+    types::{PyDict, PyList},
+};
 
 struct CredentialEntry<'py>(Bound<'py, PyAny>);
 
@@ -14,16 +15,16 @@ impl<'py> CredentialEntry<'py> {
     }
 }
 
-pub(super) fn prepare<'py>(
+pub fn prepare<'py>(
     py: Python<'py>,
     kwargs: &Bound<'py, PyDict>,
-    logger: &super::PythonLogger,
+    logger: &crate::PythonLogger,
 ) -> PyResult<Bound<'py, PyDict>> {
     let arguments = kwargs.copy()?;
     arguments.set_item("litellm_logging_obj", logger.object(py))?;
     let litellm = py.import("litellm")?;
     inherit_credentials(py, &litellm, &arguments)?;
-    py.import("litellm.rust_bridge.lifecycle")?
+    py.import("litellm.rust_bridge.legacy_callbacks")?
         .getattr("check_limits")?
         .call1((&arguments,))?;
     Ok(arguments)
@@ -49,7 +50,7 @@ fn inherit_credentials(
         .iter()
         .map(|credential| CredentialEntry(credential).name())
         .collect::<PyResult<Vec<_>>>()?;
-    let Some(index) = credential_index(&requested, &names) else {
+    let Some(index) = names.iter().position(|name| *name == requested) else {
         py.import("litellm._logging")?.getattr("verbose_logger")?.call_method1(
             "warning",
             ("litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", requested, names.len()),
@@ -60,9 +61,9 @@ fn inherit_credentials(
     let values = selected.values()?;
     let supplied: Vec<String> = arguments.keys().extract()?;
     let fields: Vec<String> = values.keys().extract()?;
-    for name in credential_default_fields(&supplied, &fields) {
-        if let Some(value) = values.get_item(name)? {
-            arguments.set_item(name, value)?;
+    for name in fields.iter().filter(|name| !supplied.contains(name)) {
+        if let Some(value) = values.get_item(name.as_str())? {
+            arguments.set_item(name.as_str(), value)?;
         }
     }
     Ok(())

litellm-rust/crates/callbacks-legacy/tests/deferred.rs

@@ -0,0 +1,162 @@
+use std::ffi::CStr;
+
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use rstest::rstest;
+
+use super::{PendingLogging, PendingSuccess};
+use crate::PythonLogger;
+use crate::test_support::{local, namespace, run};
+
+/// A deferred success for the namespace's `logger` and `response`, bound as `pending`.
+fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> {
+    let locals = namespace(py, c"response = object()");
+    run(py, &locals, script);
+    let pending = Py::new(
+        py,
+        PendingLogging {
+            pending: Some(PendingSuccess {
+                logger: PythonLogger::new(local(&locals, "logger").unbind(), true),
+                response: Some(local(&locals, "response").unbind()),
+                start: py.None(),
+                end: Some(py.None()),
+            }),
+        },
+    )
+    .unwrap();
+    locals.set_item("pending", pending).unwrap();
+    locals
+}
+
+#[test]
+fn release_enqueues_the_success_once_in_the_releasing_context() {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = defer(
+            py,
+            c"
+from contextvars import ContextVar
+
+marker = ContextVar('marker', default='unset')
+observed = []
+
+def on_enqueue(coroutine):
+    observed.append(marker.get())
+    pending.release(True)
+
+logger.on_enqueue = on_enqueue
+",
+        );
+        run(
+            py,
+            &locals,
+            c"
+marker.set('release')
+pending.release(True)
+pending.release(True)
+assert observed == ['release'], observed
+assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls
+assert logger.calls[0][1] is response
+",
+        );
+    });
+}
+
+#[test]
+fn a_blocked_release_drops_the_success_for_good() {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = defer(py, c"");
+        run(
+            py,
+            &locals,
+            c"
+pending.release(False)
+pending.release(True)
+assert logger.calls == [], logger.calls
+",
+        );
+    });
+}
+
+#[test]
+fn a_release_after_the_async_callbacks_went_away_only_keeps_the_books() {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = defer(py, c"logger.needed = {'async_success': False}");
+        run(
+            py,
+            &locals,
+            c"
+pending.release(True)
+assert logger.calls == [('success_bookkeeping', True)], logger.calls
+",
+        );
+    });
+}
+
+#[rstest]
+#[case::ordinary_error(c"RuntimeError('queue full')", false)]
+#[case::cancellation(c"asyncio.CancelledError()", true)]
+fn a_failed_enqueue_closes_the_coroutine_and_is_never_replayed(
+    #[case] failure: &CStr,
+    #[case] propagates: bool,
+) {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = defer(
+            py,
+            c"
+import asyncio
+
+def on_enqueue(coroutine):
+    raise failure
+
+logger.on_enqueue = on_enqueue
+",
+        );
+        locals
+            .set_item("failure", py.eval(failure, None, Some(&locals)).unwrap())
+            .unwrap();
+        let released = local(&locals, "pending").call_method1("release", (true,));
+        match released {
+            Ok(_) => assert!(!propagates),
+            Err(error) => {
+                assert!(propagates);
+                assert!(error.value(py).is(local(&locals, "failure")));
+            }
+        }
+        locals.set_item("propagates", propagates).unwrap();
+        run(
+            py,
+            &locals,
+            c"
+pending.release(True)
+assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls
+assert unraisable_from(logger) == ([] if propagates else [failure])
+",
+        );
+    });
+}
+
+#[test]
+fn an_unreleased_success_does_not_keep_its_logger_alive() {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = defer(py, c"");
+        run(
+            py,
+            &locals,
+            c"
+import gc
+import weakref
+
+logger.pending = pending
+reference = weakref.ref(logger)
+del logger, pending
+gc.collect()
+assert reference() is None
+",
+        );
+    });
+}

litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs

@@ -0,0 +1,246 @@
+use std::ffi::CStr;
+
+use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing};
+use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue};
+use pyo3::exceptions::asyncio::CancelledError;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use rstest::rstest;
+
+use super::LegacyLogging;
+use crate::test_support::{legacy_call, local, namespace, run};
+
+const CALL: &CStr = c"
+document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
+kwargs = {'logger': logger, 'document': document}
+";
+
+const TIMING: Timing = Timing {
+    start_time: 0.0,
+    end_time: 1.0,
+};
+
+fn begin<'py>(
+    py: Python<'py>,
+    locals: &Bound<'py, PyDict>,
+    asynchronous: bool,
+) -> (LegacyLogging, AdapterStep) {
+    let mut logging = legacy_call(py, locals, asynchronous);
+    let kwargs = local(locals, "kwargs")
+        .cast_into::<PyDict>()
+        .unwrap()
+        .unbind();
+    let step = logging.begin(py, kwargs, 0.0).unwrap();
+    (logging, step)
+}
+
+fn arguments<'py>(py: Python<'py>, step: AdapterStep) -> Bound<'py, PyDict> {
+    let AdapterStep::Arguments(arguments) = step else {
+        panic!("expected the prepared arguments");
+    };
+    arguments.into_bound(py)
+}
+
+fn awaits_deployment_hook(step: &AdapterStep) -> bool {
+    matches!(step, AdapterStep::Await(_))
+}
+
+#[rstest]
+#[case::synchronous(false)]
+#[case::asynchronous(true)]
+fn deployment_pre_call_hook_runs_only_for_asynchronous_calls(#[case] asynchronous: bool) {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(py, CALL);
+        let (_, step) = begin(py, &locals, asynchronous);
+        assert_eq!(awaits_deployment_hook(&step), asynchronous);
+        let names: Vec<String> = local(&locals, "logger")
+            .call_method0("names")
+            .unwrap()
+            .extract()
+            .unwrap();
+        assert_eq!(names.contains(&"pre_hook".to_string()), asynchronous);
+    });
+}
+
+#[test]
+fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(
+            py,
+            c"
+document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
+replacement = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk'}
+kwargs = {'logger': logger, 'document': document}
+replaced_kwargs = {'logger': logger, 'document': replacement, 'pages': [0]}
+",
+        );
+        let (mut logging, step) = begin(py, &locals, true);
+        assert!(awaits_deployment_hook(&step));
+        let step = logging
+            .resume(py, Ok(local(&locals, "replaced_kwargs").unbind()))
+            .unwrap();
+        locals.set_item("prepared", arguments(py, step)).unwrap();
+        run(
+            py,
+            &locals,
+            c"
+assert prepared['document'] is replacement
+assert prepared['pages'] is replaced_kwargs['pages']
+assert prepared['litellm_logging_obj'] is logger
+assert 'litellm_logging_obj' not in replaced_kwargs
+[checked] = [value for name, value in logger.calls if name == 'check_limits']
+assert checked is prepared
+",
+        );
+    });
+}
+
+#[test]
+fn response_returned_by_the_post_call_hook_is_finalized_and_returned() {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(
+            py,
+            c"
+kwargs = {'logger': logger}
+response = object()
+replacement = object()
+logger.hooks = {'pre': lambda kwargs: kwargs}
+",
+        );
+        let (mut logging, _) = begin(py, &locals, true);
+        logging
+            .resume(py, Ok(local(&locals, "kwargs").unbind()))
+            .unwrap();
+        let step = logging
+            .after_success(py, local(&locals, "response").unbind(), TIMING)
+            .unwrap();
+        assert!(awaits_deployment_hook(&step));
+        let step = logging
+            .resume(py, Ok(local(&locals, "replacement").unbind()))
+            .unwrap();
+        let AdapterStep::Response(returned) = step else {
+            panic!("expected the finalized response");
+        };
+        assert!(returned.bind(py).is(local(&locals, "replacement")));
+        run(
+            py,
+            &locals,
+            c"
+[finalized] = [value for name, value in logger.calls if name == 'finalize']
+assert finalized is replacement
+",
+        );
+    });
+}
+
+#[rstest]
+#[case::pre_call(false)]
+#[case::post_call(true)]
+fn cancelling_a_deployment_hook_ends_the_call_with_that_cancellation(#[case] post_call: bool) {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(py, c"kwargs = {'logger': logger}\nresponse = object()");
+        let (mut logging, _) = begin(py, &locals, true);
+        if post_call {
+            logging
+                .resume(py, Ok(local(&locals, "kwargs").unbind()))
+                .unwrap();
+            logging
+                .after_success(py, local(&locals, "response").unbind(), TIMING)
+                .unwrap();
+        }
+        let cancellation = CancelledError::new_err("cancelled");
+        let cancelled = cancellation.value(py).clone();
+        let error = logging.resume(py, Err(cancellation)).err().unwrap();
+        assert!(error.value(py).is(&cancelled));
+        let names: Vec<String> = local(&locals, "logger")
+            .call_method0("names")
+            .unwrap()
+            .extract()
+            .unwrap();
+        assert!(!names.iter().any(|name| name.contains("handler")));
+    });
+}
+
+#[rstest]
+#[case::hook_completed(false)]
+#[case::hook_cancelled(true)]
+fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelled: bool) {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(
+            py,
+            c"kwargs = {'logger': logger}\nfailure = ValueError('provider')",
+        );
+        let (mut logging, _) = begin(py, &locals, true);
+        logging
+            .resume(py, Ok(local(&locals, "kwargs").unbind()))
+            .unwrap();
+        let failure = PyErr::from_value(local(&locals, "failure"));
+        let failed = CallEvent::Failed {
+            timing: TIMING,
+            origin: FailureOrigin::Call,
+        };
+        let step = logging
+            .emit(py, &failed, Some(PublicValue::Error(&failure)))
+            .unwrap();
+        assert!(awaits_deployment_hook(&step));
+        let hook_result = if cancelled {
+            Err(CancelledError::new_err("cancelled"))
+        } else {
+            Ok(py.None())
+        };
+        assert!(matches!(
+            logging.resume(py, hook_result).unwrap(),
+            AdapterStep::Await(_)
+        ));
+        run(
+            py,
+            &locals,
+            c"
+assert logger.names()[-3:] == ['failure_hook', 'failure_handler', 'async_failure_handler'], logger.calls
+assert all(value is failure for name, value in logger.calls if name.endswith('_handler'))
+",
+        );
+    });
+}
+
+#[rstest]
+#[case::synchronous(false)]
+#[case::asynchronous(true)]
+fn a_limit_rejected_before_the_call_surfaces_as_the_callers_error(#[case] asynchronous: bool) {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(
+            py,
+            c"
+class BudgetExceeded(Exception):
+    pass
+
+rejection = BudgetExceeded('over budget')
+
+class LimitedLogger(StubLogger):
+    def check_limits(self, arguments):
+        raise rejection
+
+logger = LimitedLogger()
+logger.hooks = {'pre': lambda kwargs: kwargs}
+kwargs = {'logger': logger}
+",
+        );
+        let mut logging = legacy_call(py, &locals, asynchronous);
+        let kwargs = local(&locals, "kwargs")
+            .cast_into::<PyDict>()
+            .unwrap()
+            .unbind();
+        let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step {
+            AdapterStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())),
+            step => Ok(step),
+        });
+        let error = result.err().unwrap();
+        assert!(error.value(py).is(local(&locals, "rejection")));
+    });
+}

litellm-rust/crates/callbacks-legacy/tests/payload.rs

@@ -0,0 +1,365 @@
+use std::ffi::CStr;
+
+use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest};
+use litellm_host_python::{AdapterStep, CallbackAdapter};
+use pyo3::prelude::*;
+use rstest::rstest;
+use serde_json::{Value, json};
+
+use super::LegacyLogging;
+use crate::PythonLogger;
+use crate::test_support::{legacy_call, local, namespace, run};
+
+/// The payload phases of `Logging` on top of `StubLogger`, with `pre_call` handing the
+/// payload to the case's `on_pre_call`.
+const PAYLOAD_LOGGER: &CStr = c"
+class Request:
+    pass
+
+class PayloadLogger(StubLogger):
+    def update_from_kwargs(self, **update):
+        self.update = update
+
+    def pre_call(self, input, api_key, additional_args):
+        self.record('pre_call', None)
+        self.pre = additional_args
+        on_pre_call(additional_args)
+
+    def _pre_call(self, input, api_key, additional_args):
+        self.record('_pre_call', None)
+
+    def record_api_call_start_time(self):
+        self.record('record_api_call_start_time', None)
+
+    def post_call(self, original_response, additional_args):
+        self.record('post_call', None)
+        self.post = (original_response, additional_args)
+
+    def record_post_call(self, response, *rest):
+        self.record('record_post_call', response)
+
+request = Request()
+kwargs = {}
+logger = PayloadLogger()
+on_pre_call = lambda additional_args: None
+check = lambda: None
+";
+
+const DOCUMENT: &str = "data:application/pdf;base64,YWJj";
+const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk";
+
+fn document(source: &str) -> Value {
+    json!({"type": "document_url", "document_url": source})
+}
+
+fn before_send(script: &CStr, caller: Value, body: Value) -> WireRequest {
+    before_send_with_secrets(script, caller, body, &[])
+}
+
+/// Runs `before_send` over `body` for a caller whose route-side view is `caller`, with the
+/// Python objects `script` binds, then delivers the provider's raw response the way the
+/// driver does and runs the script's `check()`.
+fn before_send_with_secrets(
+    script: &CStr,
+    caller: Value,
+    body: Value,
+    secret_fields: &[&str],
+) -> WireRequest {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(py, PAYLOAD_LOGGER);
+        run(py, &locals, script);
+        let mut logging = LegacyLogging {
+            logger: Some(PythonLogger::new(local(&locals, "logger").unbind(), true)),
+            ..legacy_call(py, &locals, false)
+        };
+        let context = RequestContext {
+            model: "model".into(),
+            custom_llm_provider: "provider".into(),
+            optional_params: caller.clone(),
+            passthrough_fields: Passthrough::unchanged(caller.as_object().unwrap(), &body),
+            secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(),
+        };
+        let wire = WireRequest {
+            url: "https://provider.invalid/ocr".into(),
+            headers: vec![("x-route".into(), "route".into())],
+            body,
+        };
+        let step = logging.before_send(py, Box::new(wire), &context).unwrap();
+        let raw = CallEvent::ResponseReceived {
+            raw: RawResponse {
+                body: "raw response".into(),
+            },
+        };
+        assert!(matches!(
+            logging.emit(py, &raw, None).unwrap(),
+            AdapterStep::Done
+        ));
+        run(py, &locals, c"check()");
+        let AdapterStep::Wire(wire) = step else {
+            panic!("before_send did not hand back the wire request");
+        };
+        *wire
+    })
+}
+
+#[rstest]
+#[case::caller_keyword(c"
+document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
+pages = [0]
+kwargs = {'document': document, 'pages': pages}
+observed = []
+on_pre_call = lambda args: observed.append(
+    (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages)
+)
+def check():
+    assert observed == [(True, True)], observed
+")]
+#[case::request_attribute_behind_an_omitted_keyword(c"
+document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
+pages = [0]
+request.document = document
+kwargs = {'pages': pages}
+observed = []
+on_pre_call = lambda args: observed.append(
+    (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages)
+)
+def check():
+    assert observed == [(True, True)], observed
+")]
+fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) {
+    let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]});
+    let wire = before_send(
+        script,
+        json!({"document": document(DOCUMENT), "pages": [0]}),
+        body.clone(),
+    );
+    assert_eq!(wire.body, body);
+}
+
+#[test]
+fn pre_call_edit_of_a_passthrough_object_reaches_the_caller_and_the_wire() {
+    let wire = before_send(
+        c"
+document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
+kwargs = {'document': document}
+def on_pre_call(args):
+    args['complete_input_dict']['document']['document_url'] = 'data:application/pdf;base64,ZWRpdGVk'
+def check():
+    assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk'
+",
+        json!({"document": document(DOCUMENT)}),
+        json!({"document": document(DOCUMENT)}),
+    );
+    assert_eq!(wire.body["document"], document(EDITED));
+}
+
+#[test]
+fn a_body_key_the_route_rewrote_is_not_the_callers_object() {
+    let wire = before_send(
+        c"
+document = {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'}
+kwargs = {'document': document}
+observed = []
+def on_pre_call(args):
+    observed.append(args['complete_input_dict']['document'] is document)
+    args['complete_input_dict']['document']['document_name'] = 'edited.pdf'
+def check():
+    assert observed == [False], observed
+    assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'}
+",
+        json!({"document": document("https://example.invalid/scan.pdf")}),
+        json!({"document": document(DOCUMENT)}),
+    );
+    assert_eq!(
+        wire.body["document"],
+        json!({"type": "document_url", "document_url": DOCUMENT, "document_name": "edited.pdf"})
+    );
+}
+
+#[rstest]
+#[case::body(
+    c"
+def on_pre_call(args):
+    args['complete_input_dict'] = {'replacement': True}
+"
+)]
+#[case::headers(
+    c"
+def on_pre_call(args):
+    args['headers'] = {'x-replacement': 'yes'}
+"
+)]
+fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) {
+    let body = json!({"document": document(DOCUMENT)});
+    let wire = before_send(script, json!({}), body.clone());
+    assert_eq!(wire.body, body);
+    assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]);
+}
+
+#[test]
+fn pre_call_header_edit_reaches_the_wire() {
+    let wire = before_send(
+        c"
+def on_pre_call(args):
+    args['headers']['x-callback'] = 'edited'
+",
+        json!({}),
+        json!({}),
+    );
+    assert_eq!(
+        wire.headers,
+        [
+            ("x-route".to_string(), "route".to_string()),
+            ("x-callback".to_string(), "edited".to_string()),
+        ]
+    );
+}
+
+#[test]
+fn pre_call_receives_the_wire_request_and_the_logger_its_redacted_request() {
+    let body = json!({"model": "model", "document": document(DOCUMENT)});
+    before_send_with_secrets(
+        c"
+logger_fn = lambda *args: None
+kwargs = {
+    'litellm_call_id': 'call-1',
+    'client_secret': 'shh',
+    'proxy_server_request': {'body': {}},
+    'logger_fn': logger_fn,
+    'litellm_request_debug': True,
+    'ocr_cost_per_page': 0.05,
+}
+observed = []
+on_pre_call = observed.append
+def check():
+    [args] = observed
+    assert args['api_base'] == 'https://provider.invalid/ocr', args
+    assert args['complete_input_dict'] == {
+        'model': 'model',
+        'document': {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'},
+    }, args
+    update = logger.update
+    assert update['model'] == 'model' and update['custom_llm_provider'] == 'provider', update
+    assert update['litellm_params']['litellm_call_id'] == 'call-1', update
+    assert update['litellm_params']['api_base'] == 'https://provider.invalid/ocr', update
+    assert update['litellm_params']['logger_fn'] is logger_fn, update
+    assert update['litellm_params']['litellm_request_debug'] is True, update
+    assert update['litellm_params']['ocr_cost_per_page'] == 0.05, update
+    assert update['kwargs']['client_secret'] == '****', update
+    assert 'proxy_server_request' not in update['kwargs'], update
+    assert update['optional_params']['client_secret'] == '****', update
+",
+        json!({"client_secret": "shh"}),
+        body,
+        &["client_secret"],
+    );
+}
+
+#[rstest]
+#[case::added_key(
+    c"
+def on_pre_call(args):
+    args['complete_input_dict']['include_image_base64'] = True
+",
+    json!({"document": document(DOCUMENT), "include_image_base64": true})
+)]
+#[case::replaced_document(
+    c"
+document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}
+kwargs = {'document': document}
+def on_pre_call(args):
+    args['complete_input_dict']['document'] = {
+        'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk'
+    }
+def check():
+    assert document['document_url'] == 'data:application/pdf;base64,YWJj', document
+",
+    json!({"document": document(EDITED)})
+)]
+#[case::retained_body_edited_after_rebinding(
+    c"
+def on_pre_call(args):
+    retained = args['complete_input_dict']
+    args['complete_input_dict'] = {'rebound': True}
+    retained['include_image_base64'] = True
+",
+    json!({"document": document(DOCUMENT), "include_image_base64": true})
+)]
+fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) {
+    let body = json!({"document": document(DOCUMENT)});
+    let wire = before_send(script, json!({"document": document(DOCUMENT)}), body);
+    assert_eq!(wire.body, expected);
+}
+
+#[test]
+fn retained_headers_edited_after_rebinding_reach_the_wire() {
+    let wire = before_send(
+        c"
+def on_pre_call(args):
+    retained = args['headers']
+    args['headers'] = {'x-rebound': 'rebound'}
+    retained['x-retained'] = 'sent'
+",
+        json!({}),
+        json!({}),
+    );
+    assert_eq!(
+        wire.headers,
+        [
+            ("x-route".to_string(), "route".to_string()),
+            ("x-retained".to_string(), "sent".to_string()),
+        ]
+    );
+}
+
+#[test]
+fn post_call_receives_the_raw_response_and_the_payload_dicts_pre_call_saw() {
+    before_send(
+        c"
+def check():
+    original_response, additional_args = logger.post
+    assert original_response == 'raw response', original_response
+    assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict']
+    assert additional_args['headers'] is logger.pre['headers']
+",
+        json!({}),
+        json!({"document": document(DOCUMENT)}),
+    );
+}
+
+#[rstest]
+#[case::every_phase_listens(c"{}", &["pre_call", "post_call"])]
+#[case::no_input_callback(
+    c"{'input': False}",
+    &["_pre_call", "record_api_call_start_time", "record_post_call"]
+)]
+#[case::no_payload_consumer(c"{'payload': False}", &["record_api_call_start_time"])]
+fn payload_callbacks_run_only_for_the_phases_someone_listens_to(
+    #[case] needed: &CStr,
+    #[case] expected_calls: &[&str],
+) {
+    let script = std::ffi::CString::new(format!(
+        "
+logger.needed = {needed}
+def on_pre_call(args):
+    args['complete_input_dict']['include_image_base64'] = True
+def check():
+    assert logger.names() == {expected_calls:?}, logger.calls
+",
+        needed = needed.to_str().unwrap(),
+        expected_calls = expected_calls,
+    ))
+    .unwrap();
+    let body = json!({"document": document(DOCUMENT)});
+    let wire = before_send(&script, json!({}), body.clone());
+    let edited = json!({"document": document(DOCUMENT), "include_image_base64": true});
+    assert_eq!(
+        wire.body,
+        if expected_calls.contains(&"pre_call") {
+            edited
+        } else {
+            body
+        }
+    );
+}

litellm-rust/crates/callbacks-legacy/tests/support.rs

@@ -0,0 +1,188 @@
+use std::ffi::CStr;
+
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+
+use crate::{LegacyLogging, LegacySurface, PublicCall};
+
+/// Stand-ins for every litellm function the legacy contract calls. Tests share one
+/// interpreter and run concurrently, so each stub is installed idempotently and forwards to
+/// the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`).
+const STUBS: &CStr = c"
+import contextvars
+import sys
+import types
+
+for name in (
+    'litellm',
+    'litellm.utils',
+    'litellm.types',
+    'litellm.types.utils',
+    'litellm._internal_context',
+    'litellm.litellm_core_utils',
+    'litellm.litellm_core_utils.logging_worker',
+    'litellm.litellm_core_utils.litellm_logging',
+    'litellm.rust_bridge',
+    'litellm.rust_bridge.legacy_callbacks',
+):
+    sys.modules.setdefault(name, types.ModuleType(name))
+
+legacy = sys.modules['litellm.rust_bridge.legacy_callbacks']
+legacy.setup = lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace(
+    logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'],
+    kwargs=kwargs,
+    bridge_owned=True,
+)
+legacy.deployment_callbacks_needed = lambda: True
+legacy.check_limits = lambda arguments: arguments['logger'].check_limits(arguments)
+legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True)
+legacy.success_bookkeeping = lambda logger, response, start, end, asynchronous: logger.record(
+    'success_bookkeeping', asynchronous
+)
+legacy.failure_bookkeeping = lambda logger, error, start, end, asynchronous: logger.record(
+    'failure_bookkeeping', asynchronous
+)
+legacy.finalize = lambda response, logger, kwargs, start, end: logger.record('finalize', response)
+
+utils = sys.modules['litellm.utils']
+utils.async_pre_call_deployment_hook = lambda kwargs, call_type: kwargs['logger'].hook(
+    'pre', kwargs, call_type
+)
+utils.async_post_call_success_deployment_hook = lambda kwargs, response, call_type: kwargs[
+    'logger'
+].hook('success', response, call_type)
+utils.async_post_call_failure_deployment_hook = lambda kwargs, error, call_type: kwargs[
+    'logger'
+].hook('failure', error, call_type)
+utils._restore_correlation_context_if_supported = lambda logger: logger.record('restore', None)
+
+internal = sys.modules['litellm._internal_context']
+if not hasattr(internal, 'is_internal_call'):
+    internal.is_internal_call = contextvars.ContextVar('is_internal_call', default=False)
+
+sys.modules['litellm.types.utils'].CustomPricingLiteLLMParams = type(
+    'CustomPricingLiteLLMParams', (), {'model_fields': {'ocr_cost_per_page': None}}
+)
+
+
+unraisable = sys.modules.setdefault(
+    'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable')
+)
+if not hasattr(unraisable, 'events'):
+    unraisable.events = []
+    sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value))
+
+
+def unraisable_from(owner):
+    return [error for source, error in unraisable.events if source is owner]
+
+
+class Worker:
+    def ensure_initialized_and_enqueue(self, coroutine):
+        return coroutine.enqueue()
+
+
+class Executor:
+    def submit(self, run, handler, *args):
+        handler.__self__.record('submit', args)
+
+
+sys.modules['litellm.litellm_core_utils.logging_worker'].GLOBAL_LOGGING_WORKER = Worker()
+sys.modules['litellm.litellm_core_utils.litellm_logging'].executor = Executor()
+
+
+class StubCoroutine:
+    def __init__(self, logger):
+        self.logger = logger
+
+    def enqueue(self):
+        self.logger.record('enqueued', None)
+        self.logger.on_enqueue(self)
+
+    def close(self):
+        self.logger.record('closed', None)
+
+
+class StubLogger:
+    def __init__(self):
+        self.calls = []
+        self.needed = {}
+        self.hooks = {}
+        self.on_enqueue = lambda coroutine: None
+
+    def record(self, name, value):
+        self.calls.append((name, value))
+
+    def names(self):
+        return [name for name, _ in self.calls]
+
+    def hook(self, phase, value, call_type):
+        self.record(phase + '_hook', call_type)
+        return self.hooks.get(phase, lambda value: 'awaitable')(value)
+
+    def check_limits(self, arguments):
+        self.record('check_limits', arguments)
+
+    def failure_handler(self, error, trace, start, end):
+        self.record('failure_handler', error)
+
+    def async_failure_handler(self, error, trace, start, end):
+        self.record('async_failure_handler', error)
+        return 'awaitable'
+
+    def success_handler(self, response, start, end):
+        self.record('success_handler', response)
+
+    def async_success_handler(self, response, start, end):
+        self.record('async_success_handler', response)
+        return StubCoroutine(self)
+
+    def handle_sync_success_callbacks_for_async_calls(self, response, start, end):
+        self.record('sync_success_for_async_call', response)
+
+
+logger = StubLogger()
+";
+
+/// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it.
+pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> {
+    let locals = PyDict::new(py);
+    py.run(STUBS, Some(&locals), Some(&locals)).unwrap();
+    py.run(script, Some(&locals), Some(&locals)).unwrap();
+    locals
+}
+
+pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) {
+    py.run(code, Some(locals), Some(locals)).unwrap();
+}
+
+pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> {
+    locals.get_item(name).unwrap().unwrap()
+}
+
+/// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`).
+pub(crate) fn legacy_call(
+    py: Python<'_>,
+    locals: &Bound<'_, PyDict>,
+    asynchronous: bool,
+) -> LegacyLogging {
+    let request = locals
+        .get_item("request")
+        .unwrap()
+        .unwrap_or_else(|| py.None().into_bound(py));
+    let kwargs = locals
+        .get_item("kwargs")
+        .unwrap()
+        .map(|kwargs| kwargs.cast_into::<PyDict>().unwrap())
+        .unwrap_or_else(|| PyDict::new(py));
+    let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap();
+    LegacyLogging::new(
+        py,
+        LegacySurface {
+            call_type: "test",
+            input_description: "test input",
+        },
+        call,
+        asynchronous,
+    )
+}

litellm-rust/crates/callbacks-legacy/tests/terminal.rs

@@ -0,0 +1,291 @@
+use std::ffi::CStr;
+
+use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing};
+use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue};
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::exceptions::asyncio::CancelledError;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use rstest::rstest;
+
+use super::LegacyLogging;
+use crate::PythonLogger;
+use crate::test_support::{legacy_call, local, namespace, run};
+
+const TIMING: Timing = Timing {
+    start_time: 0.0,
+    end_time: 1.0,
+};
+
+fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging {
+    LegacyLogging {
+        logger: Some(PythonLogger::new(local(locals, "logger").unbind(), true)),
+        ..legacy_call(py, locals, asynchronous)
+    }
+}
+
+fn succeed(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep {
+    let response = local(locals, "response").unbind();
+    logging
+        .emit(
+            py,
+            &CallEvent::Succeeded { timing: TIMING },
+            Some(PublicValue::Response(&response)),
+        )
+        .unwrap()
+}
+
+fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep {
+    let failure = PyErr::from_value(local(locals, "failure"));
+    logging
+        .emit(
+            py,
+            &CallEvent::Failed {
+                timing: TIMING,
+                origin: FailureOrigin::Host,
+            },
+            Some(PublicValue::Error(&failure)),
+        )
+        .unwrap()
+}
+
+#[rstest]
+#[case::sync_listened(false, c"", &["submit"])]
+#[case::sync_unlistened(false, c"logger.needed = {'sync_success': False}", &["success_bookkeeping"])]
+#[case::async_listened(
+    true,
+    c"",
+    &["async_success_handler", "enqueued", "sync_success_for_async_call"]
+)]
+#[case::async_unlistened(
+    true,
+    c"logger.needed = {'async_success': False, 'sync_success_async': False}",
+    &["success_bookkeeping"]
+)]
+#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])]
+#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
+fn success_reaches_only_the_callbacks_that_listen(
+    #[case] asynchronous: bool,
+    #[case] script: &CStr,
+    #[case] expected: &[&str],
+) {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(py, c"response = object()");
+        run(py, &locals, script);
+        let mut logging = logged(py, &locals, asynchronous);
+        assert!(matches!(
+            succeed(py, &locals, &mut logging),
+            AdapterStep::Done
+        ));
+        let names: Vec<String> = local(&locals, "logger")
+            .call_method0("names")
+            .unwrap()
+            .extract()
+            .unwrap();
+        assert_eq!(names, expected);
+        run(
+            py,
+            &locals,
+            c"
+assert all(value is response for name, value in logger.calls if name.endswith('_handler'))
+assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False)
+",
+        );
+    });
+}
+
+#[rstest]
+#[case::synchronous(false, &["failure_handler"])]
+#[case::asynchronous(true, &[])]
+fn internal_calls_skip_failure_callbacks_only_when_asynchronous(
+    #[case] asynchronous: bool,
+    #[case] expected: &[&str],
+) {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(py, c"failure = ValueError('provider')");
+        let mut logging = LegacyLogging {
+            internal: true,
+            ..logged(py, &locals, asynchronous)
+        };
+        assert!(matches!(fail(py, &locals, &mut logging), AdapterStep::Done));
+        let names: Vec<String> = local(&locals, "logger")
+            .call_method0("names")
+            .unwrap()
+            .extract()
+            .unwrap();
+        assert_eq!(names, expected);
+    });
+}
+
+#[test]
+fn internal_async_calls_skip_the_async_success_fan_out() {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(py, c"response = object()");
+        let mut logging = LegacyLogging {
+            internal: true,
+            ..logged(py, &locals, true)
+        };
+        succeed(py, &locals, &mut logging);
+        run(
+            py,
+            &locals,
+            c"assert logger.names() == ['sync_success_for_async_call'], logger.calls",
+        );
+    });
+}
+
+#[test]
+fn a_failing_success_callback_is_reported_without_replacing_the_response() {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(
+            py,
+            c"
+response = object()
+failure = ValueError('terminal diagnostic')
+
+class FailingLogger(StubLogger):
+    def handle_sync_success_callbacks_for_async_calls(self, *args):
+        raise failure
+
+logger = FailingLogger()
+",
+        );
+        let mut logging = logged(py, &locals, true);
+        assert!(matches!(
+            succeed(py, &locals, &mut logging),
+            AdapterStep::Done
+        ));
+        assert!(
+            logging
+                .response
+                .as_ref()
+                .unwrap()
+                .bind(py)
+                .is(local(&locals, "response"))
+        );
+        run(py, &locals, c"assert unraisable_from(logger) == [failure]");
+    });
+}
+
+#[rstest]
+#[case::sync_listened(false, c"", &["failure_handler"])]
+#[case::sync_unlistened(false, c"logger.needed = {'sync_failure': False}", &["failure_bookkeeping"])]
+#[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])]
+#[case::async_unlistened(
+    true,
+    c"logger.needed = {'sync_failure': False, 'async_failure': False}",
+    &["failure_bookkeeping", "failure_bookkeeping"]
+)]
+fn failure_reaches_only_the_callbacks_that_listen(
+    #[case] asynchronous: bool,
+    #[case] script: &CStr,
+    #[case] expected: &[&str],
+) {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(py, c"failure = ValueError('provider')");
+        run(py, &locals, script);
+        let mut logging = logged(py, &locals, asynchronous);
+        let step = fail(py, &locals, &mut logging);
+        let awaits_async_handler = expected.contains(&"async_failure_handler");
+        assert_eq!(matches!(step, AdapterStep::Await(_)), awaits_async_handler);
+        let names: Vec<String> = local(&locals, "logger")
+            .call_method0("names")
+            .unwrap()
+            .extract()
+            .unwrap();
+        assert_eq!(names, expected);
+        run(
+            py,
+            &locals,
+            c"assert all(value is failure for name, value in logger.calls if name.endswith('_handler'))",
+        );
+    });
+}
+
+#[test]
+fn a_failing_sync_failure_callback_keeps_the_error_and_still_runs_the_async_family() {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(
+            py,
+            c"
+failure = ValueError('selected')
+
+class FailingLogger(StubLogger):
+    def failure_handler(self, error, trace, start, end):
+        self.record('failure_handler', error)
+        raise RuntimeError('handler failed')
+
+logger = FailingLogger()
+",
+        );
+        let mut logging = logged(py, &locals, true);
+        assert!(matches!(
+            fail(py, &locals, &mut logging),
+            AdapterStep::Await(_)
+        ));
+        assert!(
+            logging
+                .error
+                .as_ref()
+                .unwrap()
+                .bind(py)
+                .is(local(&locals, "failure"))
+        );
+        run(
+            py,
+            &locals,
+            c"assert logger.names() == ['failure_handler', 'async_failure_handler'], logger.calls",
+        );
+    });
+}
+
+#[rstest]
+#[case::completed(None, true)]
+#[case::handler_error(Some(false), true)]
+#[case::cancelled(Some(true), false)]
+fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled(
+    #[case] error: Option<bool>,
+    #[case] done: bool,
+) {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(py, c"failure = ValueError('provider')");
+        let mut logging = logged(py, &locals, true);
+        fail(py, &locals, &mut logging);
+        let result = match error {
+            None => Ok(py.None()),
+            Some(false) => Err(PyRuntimeError::new_err("handler failed")),
+            Some(true) => Err(CancelledError::new_err("cancelled")),
+        };
+        let expected = result.as_ref().err().map(|error| error.value(py).clone());
+        match logging.resume(py, result) {
+            Ok(step) => assert!(done && matches!(step, AdapterStep::Done)),
+            Err(propagated) => {
+                assert!(!done);
+                assert!(propagated.value(py).is(expected.unwrap()));
+            }
+        }
+    });
+}
+
+#[test]
+fn closing_restores_the_correlation_context_once() {
+    Python::initialize();
+    Python::attach(|py| {
+        let locals = namespace(py, c"");
+        let mut logging = logged(py, &locals, true);
+        logging.close(py);
+        logging.close(py);
+        run(
+            py,
+            &locals,
+            c"assert logger.names() == ['restore'], logger.calls",
+        );
+    });
+}

litellm-rust/crates/callbacks/Cargo.toml

@@ -1,15 +1,13 @@
 [package]
-name = "litellm-python-interop"
+name = "litellm-callbacks"
 version = "0.1.0"
 edition.workspace = true
 license.workspace = true
 repository.workspace = true
 
 [dependencies]
-pyo3.workspace = true
-pythonize.workspace = true
-serde.workspace = true
+serde_json.workspace = true
 
 [dev-dependencies]
 rstest.workspace = true
-serde_json.workspace = true
+tokio = { workspace = true, features = ["macros"] }

litellm-rust/crates/callbacks/src/event.rs

@@ -0,0 +1,135 @@
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use serde_json::{Map, Value};
+
+/// Seconds since the Unix epoch, on one clock for every host.
+pub fn epoch_seconds() -> f64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|duration| duration.as_secs_f64())
+        .unwrap_or(0.0)
+}
+
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct Timing {
+    pub start_time: f64,
+    pub end_time: f64,
+}
+
+/// The provider request as it is about to leave, offered to the host for rewriting.
+#[derive(Clone, Debug, PartialEq)]
+pub struct WireRequest {
+    pub url: String,
+    pub headers: Vec<(String, String)>,
+    pub body: Value,
+}
+
+/// What the route knows about the request it is sending, for a host that logs it. The
+/// route owns these facts; a host reads them beside the wire request and never rewrites
+/// them.
+#[derive(Clone, Debug, PartialEq)]
+pub struct RequestContext {
+    pub model: String,
+    pub custom_llm_provider: String,
+    /// The route's parameters before the provider transformation.
+    pub optional_params: Value,
+    pub passthrough_fields: Passthrough,
+    /// Optional-param names that carry credentials and must be redacted when logged.
+    pub secret_fields: Vec<String>,
+}
+
+/// Body keys whose values are the caller's inputs, unchanged by the route. The only way to
+/// build one is to compare the two, so a route cannot name a key it rewrote.
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct Passthrough(Vec<String>);
+
+impl Passthrough {
+    pub fn unchanged(caller: &Map<String, Value>, body: &Value) -> Self {
+        Self(
+            caller
+                .iter()
+                .filter(|(name, value)| body.get(name.as_str()) == Some(*value))
+                .map(|(name, _)| name.clone())
+                .collect(),
+        )
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = &str> {
+        self.0.iter().map(String::as_str)
+    }
+
+    pub fn contains(&self, name: &str) -> bool {
+        self.0.iter().any(|field| field == name)
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct RawResponse {
+    pub body: String,
+}
+
+/// Whether a failure surfaced inside the call, including a host op the call asked for,
+/// or in a host step around it (preparing the arguments, finalizing the response).
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum FailureOrigin {
+    Call,
+    Host,
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum CallEvent {
+    ResponseReceived {
+        raw: RawResponse,
+    },
+    Succeeded {
+        timing: Timing,
+    },
+    Failed {
+        timing: Timing,
+        origin: FailureOrigin,
+    },
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+    use serde_json::json;
+
+    use super::*;
+
+    #[rstest]
+    #[case::unchanged_scalar(json!({"pages": [0]}), json!({"pages": [0]}), &["pages"])]
+    #[case::unchanged_explicit_null(json!({"pages": null}), json!({"pages": null}), &["pages"])]
+    #[case::unchanged_nested_object(
+        json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}),
+        json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}, "model": "m"}),
+        &["document"]
+    )]
+    #[case::rewritten_value(
+        json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}),
+        json!({"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}}),
+        &[]
+    )]
+    #[case::dropped_nested_field(
+        json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "document_name": "b.png"}}),
+        json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}),
+        &[]
+    )]
+    #[case::added_nested_field(
+        json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}),
+        json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "detail": "high"}}),
+        &[]
+    )]
+    #[case::reordered_array(json!({"pages": [0, 1]}), json!({"pages": [1, 0]}), &[])]
+    #[case::consumed_by_the_route(json!({"api_key": "k", "pages": [0]}), json!({"pages": [0]}), &["pages"])]
+    #[case::added_by_the_route(json!({}), json!({"model": "m"}), &[])]
+    #[case::non_object_body(json!({"pages": [0]}), json!([{"pages": [0]}]), &[])]
+    fn passthrough_is_exactly_the_callers_unchanged_keys(
+        #[case] caller: Value,
+        #[case] body: Value,
+        #[case] expected: &[&str],
+    ) {
+        let passthrough = Passthrough::unchanged(caller.as_object().unwrap(), &body);
+        assert_eq!(passthrough.iter().collect::<Vec<_>>(), expected);
+    }
+}

litellm-rust/crates/callbacks/src/host.rs

@@ -0,0 +1,45 @@
+use std::future::Future;
+
+use crate::event::{CallEvent, RequestContext, WireRequest};
+use crate::route::Route;
+
+/// One suspension point of a native call, performed by the host.
+pub enum HostOp<R: Route> {
+    Route(R::Op),
+    BeforeSend {
+        wire: Box<WireRequest>,
+        context: Box<RequestContext>,
+    },
+    Emit(CallEvent),
+}
+
+pub enum HostResult<R: Route> {
+    Route(R::OpResult),
+    BeforeSend(Box<WireRequest>),
+    Emitted,
+}
+
+/// A host answer that is either available now or arrives once the host's own
+/// suspension (a Python awaitable, for example) resolves.
+pub enum HostStep<V, S> {
+    Ready(V),
+    Suspend(S),
+}
+
+/// An in-process host: answers route operations and observes the call without leaving
+/// the Rust runtime. Language hosts implement their own driver instead.
+pub trait Host<R: Route>: Send + Sync {
+    fn route(&self, op: R::Op) -> impl Future<Output = Result<R::OpResult, R::Error>> + Send;
+
+    fn before_send(
+        &self,
+        wire: WireRequest,
+        _context: &RequestContext,
+    ) -> impl Future<Output = Result<WireRequest, R::Error>> + Send {
+        async move { Ok(wire) }
+    }
+
+    fn emit(&self, _event: &CallEvent) -> impl Future<Output = Result<(), R::Error>> + Send {
+        async { Ok(()) }
+    }
+}

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

@@ -0,0 +1,12 @@
+//! The contract between a native call and the host runtime that drives it.
+//!
+//! A host is whatever sits on the far side of the language boundary: CPython today,
+//! another runtime later. Core implements [`machine::Machine`] per route and never learns
+//! which host is on the other end. The machine yields [`host::HostOp`]s; a driver answers
+//! them, observes [`event::CallEvent`]s and may rewrite the wire request before it is sent.
+
+pub mod event;
+pub mod host;
+pub mod machine;
+pub mod route;
+pub mod run;

litellm-rust/crates/callbacks/src/machine.rs

@@ -0,0 +1,63 @@
+use std::future::Future;
+use std::pin::Pin;
+
+use crate::host::{HostOp, HostResult};
+use crate::route::Route;
+
+pub enum MachineStep<R: Route, C> {
+    Host(HostOp<R>),
+    Complete(C),
+}
+
+pub type Step<'a, M> = Pin<
+    Box<
+        dyn Future<
+                Output = Result<
+                    MachineStep<<M as Machine>::Route, <M as Machine>::Complete>,
+                    <<M as Machine>::Route as Route>::Error,
+                >,
+            > + Send
+            + 'a,
+    >,
+>;
+
+pub type Interrupted<'a, M> = Pin<
+    Box<
+        dyn Future<
+                Output = Result<<M as Machine>::Complete, <<M as Machine>::Route as Route>::Error>,
+            > + Send
+            + 'a,
+    >,
+>;
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum HostFailure<E> {
+    Error(E),
+    Cancelled(E),
+}
+
+impl<E> HostFailure<E> {
+    pub fn into_error(self) -> E {
+        match self {
+            Self::Error(error) | Self::Cancelled(error) => error,
+        }
+    }
+}
+
+/// A resumable call. Core implements it per route; a host drives it. Every suspension
+/// point is an op the host performs and answers with a result.
+pub trait Machine: Send {
+    type Route: Route;
+    type Complete: Send + 'static;
+
+    /// `None` on the first call and whenever the previous step completed without
+    /// yielding an op; otherwise the result of the op last yielded.
+    fn resume(&mut self, result: Option<HostResult<Self::Route>>) -> Step<'_, Self>;
+
+    /// The host failed to perform the pending op, or the caller cancelled. The call
+    /// yields no further ops.
+    fn interrupt(
+        &mut self,
+        failure: HostFailure<<Self::Route as Route>::Error>,
+    ) -> Interrupted<'_, Self>;
+}

litellm-rust/crates/callbacks/src/route.rs

@@ -0,0 +1,9 @@
+/// One public call surface: what a completed call produces, how it fails, and the
+/// route-specific operations only its host can perform (request projection, file reads,
+/// token acquisition).
+pub trait Route: Send + Sync + 'static {
+    type Response: Send + 'static;
+    type Error: Clone + Send + Sync + 'static;
+    type Op: Send + 'static;
+    type OpResult: Send + 'static;
+}

litellm-rust/crates/callbacks/src/run.rs

@@ -0,0 +1,149 @@
+use crate::event::{CallEvent, FailureOrigin, Timing, epoch_seconds};
+use crate::host::{Host, HostOp, HostResult};
+use crate::machine::{HostFailure, Machine, MachineStep};
+use crate::route::Route;
+
+/// Drives a machine to completion against an in-process host and emits exactly one
+/// terminal event.
+pub async fn run<M, H>(mut machine: M, host: &H) -> Result<M::Complete, <M::Route as Route>::Error>
+where
+    M: Machine,
+    H: Host<M::Route>,
+{
+    let start_time = epoch_seconds();
+    let mut result = None;
+    let outcome = loop {
+        let step = match machine.resume(result.take()).await {
+            Ok(MachineStep::Complete(complete)) => break Ok(complete),
+            Ok(MachineStep::Host(op)) => op,
+            Err(error) => break Err(error),
+        };
+        let answer = match step {
+            HostOp::Route(op) => host.route(op).await.map(HostResult::Route),
+            HostOp::BeforeSend { wire, context } => host
+                .before_send(*wire, &context)
+                .await
+                .map(|wire| HostResult::BeforeSend(Box::new(wire))),
+            HostOp::Emit(event) => host.emit(&event).await.map(|()| HostResult::Emitted),
+        };
+        match answer {
+            Ok(answer) => result = Some(answer),
+            Err(error) => break machine.interrupt(HostFailure::Error(error)).await,
+        }
+    };
+    let timing = Timing {
+        start_time,
+        end_time: epoch_seconds(),
+    };
+    let terminal = match &outcome {
+        Ok(_) => CallEvent::Succeeded { timing },
+        Err(_) => CallEvent::Failed {
+            timing,
+            origin: FailureOrigin::Call,
+        },
+    };
+    let _ = host.emit(&terminal).await;
+    outcome
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Mutex;
+
+    use super::*;
+    use crate::machine::{Interrupted, Step};
+
+    struct Unit;
+
+    impl Route for Unit {
+        type Response = ();
+        type Error = &'static str;
+        type Op = &'static str;
+        type OpResult = ();
+    }
+
+    struct Scripted {
+        ops: Vec<&'static str>,
+        outcome: Result<(), &'static str>,
+    }
+
+    impl Machine for Scripted {
+        type Route = Unit;
+        type Complete = ();
+
+        fn resume(&mut self, _: Option<HostResult<Unit>>) -> Step<'_, Self> {
+            Box::pin(async move {
+                if !self.ops.is_empty() {
+                    return Ok(MachineStep::Host(HostOp::Route(self.ops.remove(0))));
+                }
+                self.outcome.map(MachineStep::Complete)
+            })
+        }
+
+        fn interrupt(&mut self, failure: HostFailure<&'static str>) -> Interrupted<'_, Self> {
+            Box::pin(async move { Err(failure.into_error()) })
+        }
+    }
+
+    #[derive(Default)]
+    struct Recording {
+        seen: Mutex<Vec<String>>,
+        fail: Option<&'static str>,
+    }
+
+    impl Host<Unit> for Recording {
+        async fn route(&self, op: &'static str) -> Result<(), &'static str> {
+            self.seen.lock().unwrap().push(format!("route:{op}"));
+            match self.fail {
+                Some(failing) if failing == op => Err("host failed"),
+                _ => Ok(()),
+            }
+        }
+
+        async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> {
+            self.seen.lock().unwrap().push(match event {
+                CallEvent::Succeeded { .. } => "succeeded".into(),
+                CallEvent::Failed { .. } => "failed".into(),
+                other => format!("{other:?}"),
+            });
+            Ok(())
+        }
+    }
+
+    fn scripted(ops: &[&'static str], outcome: Result<(), &'static str>) -> Scripted {
+        Scripted {
+            ops: ops.to_vec(),
+            outcome,
+        }
+    }
+
+    #[tokio::test]
+    async fn forwards_every_op_then_emits_one_succeeded() {
+        let host = Recording::default();
+        let outcome = run(scripted(&["project", "send"], Ok(())), &host).await;
+        assert_eq!(outcome, Ok(()));
+        assert_eq!(
+            *host.seen.lock().unwrap(),
+            ["route:project", "route:send", "succeeded"]
+        );
+    }
+
+    #[tokio::test]
+    async fn errors_and_host_failures_each_emit_failed_once() {
+        let host = Recording::default();
+        let outcome = run(scripted(&[], Err("boom")), &host).await;
+        assert_eq!(outcome, Err("boom"));
+        assert_eq!(*host.seen.lock().unwrap(), ["failed"]);
+
+        let host = Recording {
+            fail: Some("send"),
+            ..Recording::default()
+        };
+        let outcome = run(scripted(&["project", "send", "never"], Ok(())), &host).await;
+        assert_eq!(outcome, Err("host failed"));
+        assert_eq!(
+            *host.seen.lock().unwrap(),
+            ["route:project", "route:send", "failed"]
+        );
+    }
+}

litellm-rust/crates/core/AGENTS.md

@@ -2,7 +2,7 @@ litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-leve
 
 A route module owns the call entrypoint, runtime types, provider/auth/URL resolution, and the handler that performs the HTTP call. Provider code and base config traits live under `src/llms/`, mirroring their Python source paths. This applies to every API surface: shared orchestration stays in its route module (`ocr/`, `chat_completions/`, `messages/`, `audio_transcription/`, or `responses/`), while provider transformations live under the corresponding Python-mirrored `llms/<provider>/` path. Import implementations directly from their canonical paths; do not add a `src/providers/` layer or compatibility re-exports. Shared provider resolution lives under `src/litellm_core_utils/get_llm_provider_logic.rs`. Handlers belong in core, never in a host crate
 
-Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or host-specific callback execution. Core owns lifecycle sequencing and callback payload construction; hosts execute the selected integrations. Env reads are limited to credential fallback in a route's `prepare.rs`.
+Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Core runs each route as a machine that yields host operations and call events; which integrations consume those events is the host's business. Env reads are limited to credential fallback in a route's `prepare.rs`.
 
 Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates.
 

litellm-rust/crates/core/Cargo.toml

@@ -7,6 +7,7 @@ repository.workspace = true
 autotests = false
 
 [dependencies]
+litellm-callbacks.workspace = true
 bytes.workspace = true
 futures-util.workspace = true
 base64.workspace = true
@@ -41,3 +42,4 @@ veil.workspace = true
 aws-smithy-eventstream = "=0.61.1"
 aws-smithy-types = "1.6.1"
 rstest.workspace = true
+rstest_reuse.workspace = true

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

@@ -1,5 +1,4 @@
-use std::sync::OnceLock;
-use std::time::Duration;
+use std::{sync::OnceLock, time::Duration};
 
 use crate::constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS;
 

litellm-rust/crates/core/src/audio_transcription/handler.rs

@@ -1,8 +1,6 @@
 use serde_json::Value;
 
-use super::Error;
-use super::client::http_client;
-use super::types::ProviderAudioTranscriptionRequest;
+use super::{Error, client::http_client, types::ProviderAudioTranscriptionRequest};
 use crate::http_utils::{http_request, truncate_error_body};
 
 pub async fn execute_audio_transcription_provider_call(
@@ -44,8 +42,7 @@ async fn signed_headers(
     request: &ProviderAudioTranscriptionRequest,
     body: &[u8],
 ) -> Result<Vec<(String, String)>, Error> {
-    use std::collections::BTreeMap;
-    use std::time::SystemTime;
+    use std::{collections::BTreeMap, time::SystemTime};
 
     use litellm_auth_aws::{aws_auth_config, resolve_credentials, sign_bedrock_post};
     use litellm_providers::base_llm::audio_transcription::transformation::AudioTranscriptionAuth;

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

@@ -3,9 +3,8 @@ pub use error::Error;
 mod client;
 mod handler;
 mod prepare;
-pub use litellm_providers::audio_transcription::types;
-
 pub use handler::execute_audio_transcription_provider_call;
+pub use litellm_providers::audio_transcription::types;
 pub use prepare::prepare_audio_transcription_provider_call;
 use serde_json::Value;
 pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};

litellm-rust/crates/core/src/audio_transcription/prepare.rs

@@ -1,13 +1,18 @@
-use super::Error;
-use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
-use crate::http_utils::{has_header, string_headers};
-use crate::litellm_core_utils::get_llm_provider_logic::{
-    CustomLlmProvider, get_custom_llm_provider,
+use litellm_providers::{
+    base_llm::audio_transcription::transformation::{
+        AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
+    },
+    bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
 };
-use litellm_providers::base_llm::audio_transcription::transformation::{
-    AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
+
+use super::{
+    Error,
+    types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest},
+};
+use crate::{
+    http_utils::{has_header, string_headers},
+    litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider},
 };
-use litellm_providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
 
 fn provider_config(provider: &str) -> Option<&'static dyn BaseAudioTranscriptionConfig> {
     if provider == "bedrock" {

litellm-rust/crates/core/src/audio_transcription/tests.rs

@@ -1,11 +1,12 @@
-use std::io::{Read, Write};
-use std::net::TcpListener;
-use std::thread;
+use std::{
+    io::{Read, Write},
+    net::TcpListener,
+    thread,
+};
 
 use serde_json::{Map, json};
 
-use super::audio_transcription;
-use super::types::AudioTranscriptionRequest;
+use super::{audio_transcription, types::AudioTranscriptionRequest};
 
 #[tokio::test]
 async fn bedrock_request_is_signed_and_contains_audio() {

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

@@ -37,278 +37,6 @@ pub struct ArgumentSpec {
     pub secret: bool,
 }
 
-pub fn should_project(name: &str, consumed: &[ArgumentSpec], bound_fields: &[&str]) -> bool {
-    consumed.iter().any(|field| field.name == name)
-        || (!bound_fields.contains(&name) && !is_control(name))
-}
-
-pub fn is_control(name: &str) -> bool {
-    crate::params::is_control_param(name) || HOST_CONTROLS.contains(&name)
-}
-
-const HOST_CONTROLS: &[&str] = &[
-    "_agentic_loop_api_surface",
-    "_agentic_loop_depth",
-    "_agentic_loop_fingerprints",
-    "_code_interpreter_interception_active",
-    "_code_interpreter_interception_converted_stream",
-    "_code_interpreter_interception_sandbox_key",
-    "_code_interpreter_interception_session_scoped",
-    "_headroom_interception_converted_stream",
-    "_litellm_strip_stream_usage",
-    "_router_weights",
-    "_websearch_interception_converted_stream",
-    "_websearch_interception_emit_native_blocks",
-    "acompletion",
-    "adaptive_router_config",
-    "adaptive_router_default_model",
-    "aembedding",
-    "aimg_generation",
-    "allm_passthrough_route",
-    "allow_client_keepalive_override",
-    "allowed_model_region",
-    "allowed_openai_params",
-    "annotation_cost_per_page",
-    "api_version",
-    "arize_api_key",
-    "arize_space_id",
-    "arize_space_key",
-    "assistant_continue_message",
-    "async_call",
-    "atext_completion",
-    "attempted_targets",
-    "auto_router_config",
-    "auto_router_config_path",
-    "auto_router_default_model",
-    "auto_router_embedding_model",
-    "auto_router_max_input_chars",
-    "auto_router_model_compression",
-    "auto_router_routing_compression",
-    "aws_batch_role_arn",
-    "azure",
-    "azure_password",
-    "azure_username",
-    "base_model",
-    "bedrock_tags",
-    "bos_token",
-    "budget_duration",
-    "cache",
-    "cache_creation_input_audio_token_cost",
-    "cache_creation_input_token_cost",
-    "cache_creation_input_token_cost_above_1hr",
-    "cache_creation_input_token_cost_above_200k_tokens",
-    "cache_creation_input_token_cost_above_272k_tokens",
-    "cache_creation_input_token_cost_above_272k_tokens_flex",
-    "cache_creation_input_token_cost_above_272k_tokens_priority",
-    "cache_creation_input_token_cost_flex",
-    "cache_creation_input_token_cost_priority",
-    "cache_creation_input_token_cost_ultrafast",
-    "cache_key",
-    "cache_read_input_audio_token_cost",
-    "cache_read_input_token_cost",
-    "cache_read_input_token_cost_above_200k_tokens",
-    "cache_read_input_token_cost_above_200k_tokens_priority",
-    "cache_read_input_token_cost_above_272k_tokens",
-    "cache_read_input_token_cost_above_272k_tokens_flex",
-    "cache_read_input_token_cost_above_272k_tokens_priority",
-    "cache_read_input_token_cost_above_512k_tokens",
-    "cache_read_input_token_cost_flex",
-    "cache_read_input_token_cost_priority",
-    "cache_read_input_token_cost_ultrafast",
-    "caching",
-    "caching_groups",
-    "citation_cost_per_token",
-    "client",
-    "client_side_timeout",
-    "complete_response",
-    "completion_call_id",
-    "complexity_router_config",
-    "complexity_router_default_model",
-    "configurable_clientside_auth_params",
-    "context_window_fallback_dict",
-    "cooldown_time",
-    "cost_per_query",
-    "custom_prompt_dict",
-    "data_residency",
-    "dd_agent_host",
-    "dd_agent_port",
-    "dd_api_key",
-    "dd_site",
-    "default_api_key_rpm_limit",
-    "default_api_key_tpm_limit",
-    "disable_add_transform_inline_image_block",
-    "enable_json_schema_validation",
-    "enable_prompt_caching",
-    "enable_tag_filtering",
-    "ensure_alternating_roles",
-    "eos_token",
-    "fallback_depth",
-    "fallbacks",
-    "fastest_response",
-    "final_prompt_value",
-    "force_timeout",
-    "gcs_bucket_name",
-    "gcs_path_service_account",
-    "google_maps_grounding_cost_per_query",
-    "headers",
-    "hf_model_name",
-    "humanloop_api_key",
-    "id",
-    "input_cost_per_audio_per_second",
-    "input_cost_per_audio_per_second_above_128k_tokens",
-    "input_cost_per_audio_token",
-    "input_cost_per_audio_token_batches",
-    "input_cost_per_character",
-    "input_cost_per_character_above_128k_tokens",
-    "input_cost_per_image",
-    "input_cost_per_image_above_128k_tokens",
-    "input_cost_per_image_token",
-    "input_cost_per_image_token_batches",
-    "input_cost_per_pixel",
-    "input_cost_per_query",
-    "input_cost_per_second",
-    "input_cost_per_token",
-    "input_cost_per_token_above_128k_tokens",
-    "input_cost_per_token_above_200k_tokens",
-    "input_cost_per_token_above_200k_tokens_priority",
-    "input_cost_per_token_above_272k_tokens",
-    "input_cost_per_token_above_272k_tokens_flex",
-    "input_cost_per_token_above_272k_tokens_priority",
-    "input_cost_per_token_above_512k_tokens",
-    "input_cost_per_token_batches",
-    "input_cost_per_token_cache_hit",
-    "input_cost_per_token_flex",
-    "input_cost_per_token_priority",
-    "input_cost_per_token_ultrafast",
-    "input_cost_per_video_per_second",
-    "input_cost_per_video_per_second_above_128k_tokens",
-    "input_cost_per_video_per_second_above_15s_interval",
-    "input_cost_per_video_per_second_above_8s_interval",
-    "input_cost_per_video_token",
-    "input_cost_per_video_token_batches",
-    "itpm",
-    "keepalive_seconds",
-    "langfuse_environment",
-    "langfuse_host",
-    "langfuse_prompt_version",
-    "langfuse_public_key",
-    "langfuse_secret",
-    "langfuse_secret_key",
-    "langsmith_api_key",
-    "langsmith_base_url",
-    "langsmith_project",
-    "langsmith_sampling_rate",
-    "langsmith_tenant_id",
-    "litellm_credential_name",
-    "litellm_disabled_callbacks",
-    "litellm_request_debug",
-    "litellm_session_id",
-    "litellm_system_prompt",
-    "litellm_trace_id",
-    "litellm_trusted_callback_vars",
-    "logger_fn",
-    "max_agentic_loops",
-    "max_budget",
-    "max_fallbacks",
-    "max_parallel_requests",
-    "merge_reasoning_content_in_choices",
-    "metadata",
-    "mock_response",
-    "mock_timeout",
-    "model_alias_map",
-    "model_config",
-    "model_file_id_mapping",
-    "model_info",
-    "model_list",
-    "newrelic_api_key",
-    "newrelic_region",
-    "no-log",
-    "num_retries",
-    "ocr_cost_per_credit",
-    "ocr_cost_per_page",
-    "order",
-    "otpm",
-    "output_cost_per_audio_per_second",
-    "output_cost_per_audio_token",
-    "output_cost_per_character",
-    "output_cost_per_character_above_128k_tokens",
-    "output_cost_per_image",
-    "output_cost_per_image_token",
-    "output_cost_per_pixel",
-    "output_cost_per_reasoning_token",
-    "output_cost_per_reasoning_token_flex",
-    "output_cost_per_reasoning_token_priority",
-    "output_cost_per_second",
-    "output_cost_per_second_1080p",
-    "output_cost_per_second_480p",
-    "output_cost_per_second_4k",
-    "output_cost_per_second_720p",
-    "output_cost_per_token",
-    "output_cost_per_token_above_128k_tokens",
-    "output_cost_per_token_above_200k_tokens",
-    "output_cost_per_token_above_200k_tokens_priority",
-    "output_cost_per_token_above_272k_tokens",
-    "output_cost_per_token_above_272k_tokens_flex",
-    "output_cost_per_token_above_272k_tokens_priority",
-    "output_cost_per_token_above_512k_tokens",
-    "output_cost_per_token_batches",
-    "output_cost_per_token_flex",
-    "output_cost_per_token_priority",
-    "output_cost_per_token_ultrafast",
-    "output_cost_per_video_per_second",
-    "output_cost_per_video_token",
-    "output_vector_size",
-    "posthog_api_key",
-    "posthog_api_url",
-    "preset_cache_key",
-    "prompt_environment",
-    "prompt_id",
-    "prompt_label",
-    "prompt_variables",
-    "prompt_version",
-    "provider_specific_header",
-    "quality_router_config",
-    "quality_router_default_model",
-    "region_name",
-    "regional_endpoint_uplift_multiplier",
-    "regional_processing_uplift_multiplier_eu",
-    "regional_processing_uplift_multiplier_us",
-    "retry_policy",
-    "retry_strategy",
-    "roles",
-    "routing_strategy",
-    "rpm",
-    "rust",
-    "s3_bucket_name",
-    "s3_output_bucket_name",
-    "s3_region_name",
-    "search_context_cost_per_query",
-    "search_tool_name",
-    "secret_fields",
-    "self",
-    "shared_session",
-    "ssl_verify",
-    "stream_response",
-    "stream_timeout",
-    "supports_system_message",
-    "tags",
-    "text_completion",
-    "tiered_pricing",
-    "tpm",
-    "ttl",
-    "turn_off_message_logging",
-    "use_chat_completions_api",
-    "use_client",
-    "use_in_pass_through",
-    "use_litellm_proxy",
-    "use_xai_oauth",
-    "user_continue_message",
-    "verbose",
-    "wandb_api_key",
-    "weave_project_id",
-    "weight",
-];
-
 pub fn compose_body<B: Serialize>(
     arguments: &CallArguments,
     body: &B,
@@ -324,9 +52,9 @@ pub fn compose_body<B: Serialize>(
         Some(Value::Object(fields)) => Some(fields),
         Some(_) => return Err(crate::params::Error::ExtraBody),
     };
-    let extensions = arguments.iter().filter(|(name, _)| {
-        !consumed.contains(&name.as_str()) && name.as_str() != "extra_body" && !is_control(name)
-    });
+    let extensions = arguments
+        .iter()
+        .filter(|(name, _)| !consumed.contains(&name.as_str()));
     Ok(Value::Object(
         fields
             .into_iter()
@@ -389,7 +117,7 @@ mod tests {
     fn composition_preserves_extensions_and_applies_shallow_explicit_overrides() {
         let original = json!({
             "known": false, "future": {"old": 1}, "null": null, "zero": 0,
-            "metadata": {"host": true}, "shared_session": "host", "api_key": "secret",
+            "metadata": {"host": true}, "timeout": 30, "api_key": "secret",
             "extra_body": {
                 "known": null, "future": {"new": [false, 0, null]},
                 "metadata": {"provider": true}, "model": "ignored", "api_key": "ignored"
@@ -412,21 +140,6 @@ mod tests {
         assert_eq!(serde_json::to_value(arguments).unwrap(), original);
     }
 
-    #[test]
-    fn projection_prioritizes_consumed_fields_and_keeps_unknown_names() {
-        let fields = [ArgumentSpec {
-            name: "id",
-            secret: false,
-        }];
-        assert!(should_project("id", &fields, &[]));
-        assert!(!should_project("id", &[], &[]));
-        assert!(should_project("future_option", &[], &[]));
-        assert!(!should_project("document", &fields, &["document"]));
-        assert!(!should_project("metadata", &fields, &[]));
-        assert!(!should_project("callbacks", &fields, &[]));
-        assert!(!should_project("ocr_cost_per_page", &fields, &[]));
-    }
-
     #[test]
     fn invalid_extra_body_is_rejected_without_coercing_it_to_empty() {
         for value in [json!(false), json!(0), json!([]), json!("")] {

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

@@ -1,122 +0,0 @@
-use std::future::Future;
-use std::pin::Pin;
-
-pub enum HostCallStep<O, C> {
-    Host(O),
-    Complete(C),
-}
-
-pub type HostCallFuture<'a, O, C, E> =
-    Pin<Box<dyn Future<Output = Result<HostCallStep<O, C>, E>> + Send + 'a>>;
-
-pub trait HostCall: Send + Sync {
-    type Error: Send + Sync + 'static;
-    type Operation: Send + 'static;
-    type Result: Send + 'static;
-    type Complete: Send + 'static;
-
-    fn resume(
-        &mut self,
-        result: Option<Self::Result>,
-    ) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error>;
-
-    fn interrupt(
-        &mut self,
-        failure: HostFailure<Self::Error>,
-    ) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error>;
-}
-
-pub enum HostStep<V, S> {
-    Ready(V),
-    Suspend(S),
-}
-
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub enum HostPhase {
-    Setup,
-    DeploymentPreCall,
-    Prepare,
-    Execute,
-    ConstructResponse,
-    DeploymentPostCall,
-    Finalize,
-    Success,
-    MapFailure,
-    DeploymentFailure,
-    Failure,
-    AsyncFailure,
-    Complete,
-}
-
-#[derive(Clone, Debug)]
-pub enum HostFailure<E> {
-    Error(E),
-    Cancelled(E),
-}
-
-pub struct HostLifecycle {
-    phase: HostPhase,
-    asynchronous: bool,
-}
-
-impl HostLifecycle {
-    pub fn new(asynchronous: bool) -> Self {
-        Self {
-            phase: HostPhase::Setup,
-            asynchronous,
-        }
-    }
-
-    pub fn phase(&self) -> HostPhase {
-        self.phase
-    }
-
-    pub fn accept<E>(&mut self, result: Result<(), HostFailure<E>>) -> Option<E> {
-        if let Err(failure) = result {
-            if self.phase == HostPhase::DeploymentFailure {
-                self.phase = HostPhase::Failure;
-                return None;
-            }
-            let error = match failure {
-                HostFailure::Cancelled(error) => {
-                    self.phase = HostPhase::Complete;
-                    return Some(error);
-                }
-                HostFailure::Error(error) => error,
-            };
-            match self.phase {
-                HostPhase::Failure | HostPhase::AsyncFailure => {
-                    self.advance();
-                    return None;
-                }
-                HostPhase::Success => self.phase = HostPhase::Complete,
-                HostPhase::Execute | HostPhase::ConstructResponse => {
-                    self.phase = HostPhase::MapFailure;
-                }
-                _ => self.phase = HostPhase::Failure,
-            }
-            return Some(error);
-        }
-        self.advance();
-        None
-    }
-
-    fn advance(&mut self) {
-        self.phase = match self.phase {
-            HostPhase::Setup if self.asynchronous => HostPhase::DeploymentPreCall,
-            HostPhase::Setup | HostPhase::DeploymentPreCall => HostPhase::Prepare,
-            HostPhase::Prepare => HostPhase::Execute,
-            HostPhase::Execute => HostPhase::ConstructResponse,
-            HostPhase::ConstructResponse if self.asynchronous => HostPhase::DeploymentPostCall,
-            HostPhase::ConstructResponse | HostPhase::DeploymentPostCall => HostPhase::Finalize,
-            HostPhase::Finalize => HostPhase::Success,
-            HostPhase::MapFailure if self.asynchronous => HostPhase::DeploymentFailure,
-            HostPhase::MapFailure | HostPhase::DeploymentFailure => HostPhase::Failure,
-            HostPhase::Failure if self.asynchronous => HostPhase::AsyncFailure,
-            HostPhase::Failure
-            | HostPhase::AsyncFailure
-            | HostPhase::Success
-            | HostPhase::Complete => HostPhase::Complete,
-        };
-    }
-}

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

@@ -1,427 +0,0 @@
-use std::future::Future;
-use std::time::{Instant, SystemTime, UNIX_EPOCH};
-
-pub mod host;
-#[cfg(test)]
-#[path = "../../tests/host_lifecycle.rs"]
-mod host_tests;
-pub mod types;
-
-pub use types::{
-    CallLifecycleContext, CallLifecyclePhase, CallLifecyclePhaseTiming, CallLifecycleRequest,
-    CallLifecycleTiming,
-};
-
-pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
-    type Error: Send + Sync;
-    type PreCallFuture<'a>: Future<Output = Result<InitialReq, Self::Error>> + Send + 'a
-    where
-        Self: 'a,
-        InitialReq: 'a,
-        ProviderReq: 'a,
-        Resp: 'a;
-
-    type DuringCallFuture<'a>: Future<Output = Result<ProviderReq, Self::Error>> + 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 Self::Error,
-        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,
-    ) -> Result<Resp, Hooks::Error>
-    where
-        InitialReq: CallLifecycleRequest,
-        Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
-        ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
-        ProviderFuture: Future<Output = Result<Resp, Hooks::Error>>,
-    {
-        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,
-    ) -> Result<Resp, Hooks::Error>
-    where
-        Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
-        ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
-        ProviderFuture: Future<Output = Result<Resp, Hooks::Error>>,
-    {
-        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: &Hooks::Error,
-        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 std::pin::Pin;
-    use std::sync::Mutex;
-
-    use super::*;
-
-    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 Error = crate::messages::Error;
-        type PreCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
-        type DuringCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
-        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 crate::messages::Error,
-            _timing: &'a CallLifecycleTiming,
-        ) -> Self::FailureFuture<'a> {
-            Box::pin(async move {
-                self.events.lock().unwrap().push("failure");
-            })
-        }
-    }
-
-    impl CallLifecycleHooks<RecordingRequest, String, String> for RecordingHooks {
-        type Error = crate::messages::Error;
-        type PreCallFuture<'a> = BoxFuture<'a, Result<RecordingRequest, crate::messages::Error>>;
-        type DuringCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
-        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 crate::messages::Error,
-            _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, crate::messages::Error>(crate::messages::Error::Transport(
-                        crate::transport::Error::Network("provider down".to_string()),
-                    ))
-                },
-            )
-            .await
-            .expect_err("call fails");
-
-        assert_eq!(
-            error,
-            crate::messages::Error::Transport(crate::transport::Error::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

@@ -1,75 +0,0 @@
-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/chat_completions/client.rs

@@ -1,5 +1,4 @@
-use std::sync::OnceLock;
-use std::time::Duration;
+use std::{sync::OnceLock, time::Duration};
 
 use crate::constants::{CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_TIMEOUT_SECS};
 

litellm-rust/crates/core/src/chat_completions/common_utils.rs

@@ -1,9 +1,11 @@
+use litellm_providers::{
+    anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG,
+    base_llm::chat::transformation::BaseConfig,
+};
 use serde_json::{Map, Value};
 
 use super::Error;
 use crate::http_utils::string_headers as shared_string_headers;
-use litellm_providers::anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
-use litellm_providers::base_llm::chat::transformation::BaseConfig;
 
 const HEADER_CONTEXT: &str = "chat completions";
 

litellm-rust/crates/core/src/chat_completions/handler.rs

@@ -1,14 +1,16 @@
+use litellm_providers::base_llm::chat::transformation::ChatCompletionsAuth;
 use serde_json::Value;
 
-use super::Error;
-use super::client::http_client;
-use super::prepare::prepare_provider_request;
-use super::types::{
-    ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
-    ResolvedChatCompletionsRequest,
+use super::{
+    Error,
+    client::http_client,
+    prepare::prepare_provider_request,
+    types::{
+        ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
+        ResolvedChatCompletionsRequest,
+    },
 };
 use crate::http_utils::{http_request, truncate_error_body};
-use litellm_providers::base_llm::chat::transformation::ChatCompletionsAuth;
 
 pub(super) async fn execute_chat_completions_provider_call(
     request: ResolvedChatCompletionsRequest<'_>,
@@ -84,8 +86,7 @@ pub(super) async fn signed_headers(
     request: &ProviderChatCompletionsRequest,
     body: &[u8],
 ) -> Result<Vec<(String, String)>, Error> {
-    use std::collections::BTreeMap;
-    use std::time::SystemTime;
+    use std::{collections::BTreeMap, time::SystemTime};
 
     use litellm_auth_aws::{
         aws_auth_config, aws_signature_headers, host_supplied_credentials,

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

@@ -14,9 +14,8 @@ pub use litellm_providers::chat::{conversation, response_utils};
 pub(crate) mod handler;
 mod prepare;
 pub mod streaming;
-pub use litellm_providers::chat::types;
-
 use handler::execute_chat_completions_provider_call;
+pub use litellm_providers::chat::types;
 use prepare::{parse_messages, resolve_provider_config, resolve_request};
 use serde_json::{Map, Value};
 use types::{ChatCompletionsRequest, ChatCompletionsResponse};

litellm-rust/crates/core/src/chat_completions/prepare.rs

@@ -1,16 +1,18 @@
+use litellm_providers::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth};
 use serde_json::Value;
 
-use super::Error;
-use super::common_utils::{chat_completions_provider_config, string_headers};
-use super::types::{
-    ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest,
-    ResolvedChatCompletionsRequest,
+use super::{
+    Error,
+    common_utils::{chat_completions_provider_config, string_headers},
+    types::{
+        ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest,
+        ResolvedChatCompletionsRequest,
+    },
 };
-use crate::http_utils::has_header;
-use crate::litellm_core_utils::get_llm_provider_logic::{
-    CustomLlmProvider, get_custom_llm_provider,
+use crate::{
+    http_utils::has_header,
+    litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider},
 };
-use litellm_providers::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth};
 
 pub(super) fn resolve_provider_config<'a>(
     model: &'a str,

litellm-rust/crates/core/src/chat_completions/tests.rs

@@ -1,9 +1,11 @@
+use litellm_providers::base_llm::chat::transformation::ChatCompletionsAuth;
 use serde_json::{Map, Value, json};
 
-use super::Error;
-use super::prepare::{prepare_provider_request, resolve_request};
-use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest};
-use litellm_providers::base_llm::chat::transformation::ChatCompletionsAuth;
+use super::{
+    Error,
+    prepare::{prepare_provider_request, resolve_request},
+    types::{ChatCompletionsRequest, ProviderChatCompletionsRequest},
+};
 
 fn prepare_chat_completions_call(
     request: ChatCompletionsRequest<'_>,
@@ -587,8 +589,10 @@ fn the_gate_agrees_with_prepare_on_every_case_it_accepts() {
 }
 
 mod round_trip {
-    use tokio::io::{AsyncReadExt, AsyncWriteExt};
-    use tokio::net::{TcpListener, TcpStream};
+    use tokio::{
+        io::{AsyncReadExt, AsyncWriteExt},
+        net::{TcpListener, TcpStream},
+    };
 
     use super::*;
     use crate::chat_completions::chat_completions;

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

@@ -1,12 +1,12 @@
 pub mod audio_transcription;
 pub mod call_arguments;
-pub mod call_lifecycle;
 pub mod chat_completions;
 pub mod constants;
 pub mod error;
 pub mod http_utils;
 pub mod litellm_core_utils;
 pub mod llms;
+pub mod machine;
 mod media;
 pub mod messages;
 pub mod ocr;

litellm-rust/crates/core/src/llms/anthropic/chat/streaming.rs

@@ -6,11 +6,13 @@ use super::super::experimental_pass_through::messages::streaming::{
     AnthropicContentBlock, AnthropicContentBlockDelta, AnthropicMessagesStreamEvent,
     AnthropicStreamUsage,
 };
-use crate::chat_completions::Error;
-use crate::chat_completions::streaming::StreamTransformer;
-use crate::chat_completions::types::{
-    ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
-    ChatCompletionsUsage,
+use crate::chat_completions::{
+    Error,
+    streaming::StreamTransformer,
+    types::{
+        ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
+        ChatCompletionsUsage,
+    },
 };
 
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]

litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/batches.rs

@@ -1,11 +1,10 @@
+use litellm_providers::anthropic::experimental_pass_through::messages::transformation::resolve_anthropic_api_base;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use time::OffsetDateTime;
 use url::Url;
 
-use crate::messages::Error;
-use crate::messages::types::AnthropicMessagesResponse;
-use litellm_providers::anthropic::experimental_pass_through::messages::transformation::resolve_anthropic_api_base;
+use crate::messages::{Error, types::AnthropicMessagesResponse};
 
 const BATCHES_PATH_SUFFIX: &str = "/v1/messages/batches";
 

litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/count_tokens.rs

@@ -1,9 +1,13 @@
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
-use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
-use crate::messages::Error;
-use crate::messages::types::{AnthropicMessage, SystemPrompt};
+use crate::{
+    constants::ANTHROPIC_OAUTH_TOKEN_PREFIX,
+    messages::{
+        Error,
+        types::{AnthropicMessage, SystemPrompt},
+    },
+};
 
 const COUNT_TOKENS_ENDPOINT: &str = "https://api.anthropic.com/v1/messages/count_tokens";
 const TOKEN_COUNTING_BETA: &str = "token-counting-2024-11-01";

litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/streaming.rs

@@ -1,9 +1,11 @@
 use base64::Engine;
 use bytes::Buf;
 use futures_util::{Stream, StreamExt};
-use litellm_framing::Framer;
-use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
-use litellm_framing::sse::{SseFrame, SseFramer};
+use litellm_framing::{
+    Framer,
+    aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer},
+    sse::{SseFrame, SseFramer},
+};
 use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 

litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs

@@ -1,13 +1,22 @@
 use serde_json::Value;
 
-use crate::call_arguments::CallArguments;
-use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext};
-use crate::llms::cohere::ocr::transformation::{CohereParseConfig, CohereRequest};
-use crate::llms::cohere::ocr::{CohereOptions, validate_document};
-use crate::ocr::OcrClient;
-use crate::ocr::document::{inline_remote_document, validate_inline_document};
-use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument, PreparedOcrRequest};
-use crate::url_utils::ApiUrl;
+use crate::{
+    call_arguments::CallArguments,
+    llms::{
+        base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext},
+        cohere::ocr::{
+            CohereOptions,
+            transformation::{CohereParseConfig, CohereRequest},
+            validate_document,
+        },
+    },
+    ocr::{
+        OcrClient,
+        document::{inline_remote_document, validate_inline_document},
+        types::{LiteLLMOcrResponse, OcrDocument, PreparedOcrRequest},
+    },
+    url_utils::ApiUrl,
+};
 
 #[derive(Default)]
 pub(crate) struct AzureAICohereParseConfig;

litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs

@@ -1,6 +1,4 @@
-use std::collections::BTreeSet;
-use std::sync::Arc;
-use std::time::Duration;
+use std::{collections::BTreeSet, time::Duration};
 
 use base64::{Engine, engine::general_purpose::STANDARD};
 use litellm_auth::{InputSource, Sourced};
@@ -11,26 +9,31 @@ use serde_json::{Map, Value};
 use serde_with::serde_as;
 use tokio::time::Instant;
 
-use crate::call_arguments::CallArguments;
-use crate::constants::{
-    AZURE_DI_API_VERSION, AZURE_DI_DEFAULT_DPI, AZURE_DI_DEFAULT_HEIGHT, AZURE_DI_DEFAULT_WIDTH,
-    AZURE_DI_SUBSCRIPTION_HEADER, OCR_POLL_RETRY_SECS,
+use crate::{
+    call_arguments::CallArguments,
+    constants::{
+        AZURE_DI_API_VERSION, AZURE_DI_DEFAULT_DPI, AZURE_DI_DEFAULT_HEIGHT,
+        AZURE_DI_DEFAULT_WIDTH, AZURE_DI_SUBSCRIPTION_HEADER, OCR_POLL_RETRY_SECS,
+    },
+    llms::base_llm::ocr::transformation::{
+        BaseOcrConfig, OcrResponseContext, decode_and_normalize_response,
+    },
+    ocr::{
+        OcrClient,
+        client::read_json_response,
+        document::InlineDocument,
+        json::DecodedOcrResponse,
+        prepare::credential_env,
+        route::OcrHost,
+        types::{
+            LiteLLMOcrResponse, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage,
+            OcrPageDimensions, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest,
+            ResolvedOcrCredentials,
+        },
+    },
+    serde_compat::{FiniteF64, LaxI64},
+    url_utils::ApiUrl,
 };
-use crate::llms::base_llm::ocr::transformation::{
-    BaseOcrConfig, OcrResponseContext, decode_and_normalize_response,
-};
-use crate::ocr::OcrClient;
-use crate::ocr::client::read_json_response;
-use crate::ocr::document::InlineDocument;
-use crate::ocr::hooks::OcrHooks;
-use crate::ocr::json::DecodedOcrResponse;
-use crate::ocr::prepare::credential_env;
-use crate::ocr::types::{
-    LiteLLMOcrResponse, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage,
-    OcrPageDimensions, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, ResolvedOcrCredentials,
-};
-use crate::serde_compat::{FiniteF64, LaxI64};
-use crate::url_utils::ApiUrl;
 
 const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY";
 const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT";
@@ -235,7 +238,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
             context.headers,
             context.connection,
             context.request_format == OcrResponseFormat::Native,
-            context.hooks,
+            context.host,
         )
         .await?;
         Ok(LiteLLMOcrResponse {
@@ -439,13 +442,13 @@ async fn read_operation_response(
     headers: &[(String, String)],
     connection: &OcrConnection,
     native: bool,
-    hooks: &Arc<dyn OcrHooks>,
+    host: &OcrHost,
 ) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, crate::ocr::Error> {
     if response.status() != reqwest::StatusCode::ACCEPTED {
         let bytes =
             crate::ocr::client::read_response_bytes(response, connection.max_response_bytes)
                 .await?;
-        crate::ocr::handler::post_call(hooks, &bytes).await?;
+        crate::ocr::handler::emit_response_received(host, &bytes).await?;
         return crate::ocr::json::decode_response(&bytes, native);
     }
     let location = response
@@ -464,8 +467,8 @@ async fn read_operation_response(
     }
     let bytes =
         crate::ocr::client::read_response_bytes(response, connection.max_response_bytes).await?;
-    crate::ocr::handler::post_call(hooks, &bytes).await?;
-    poll_operation(http_client, operation, headers, connection, native, hooks).await
+    crate::ocr::handler::emit_response_received(host, &bytes).await?;
+    poll_operation(http_client, operation, headers, connection, native, host).await
 }
 
 async fn poll_operation(
@@ -474,7 +477,7 @@ async fn poll_operation(
     headers: &[(String, String)],
     connection: &OcrConnection,
     native: bool,
-    hooks: &Arc<dyn OcrHooks>,
+    host: &OcrHost,
 ) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, crate::ocr::Error> {
     let deadline = Instant::now()
         .checked_add(connection.poll_timeout)
@@ -516,7 +519,7 @@ async fn poll_operation(
         .map_err(|_| crate::ocr::Error::PollTimeout)??;
         match &decoded.data.status {
             Some(OperationStatus::Succeeded) => {
-                crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?;
+                crate::ocr::handler::emit_response_received(host, decoded.text.as_bytes()).await?;
                 return Ok(decoded);
             }
             Some(OperationStatus::Running | OperationStatus::NotStarted) => {
@@ -807,7 +810,12 @@ mod tests {
 
     use std::sync::{Arc, Mutex};
 
-    use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
+    use litellm_callbacks::event::CallEvent;
+
+    use crate::ocr::{
+        LocalOcrHost,
+        test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
+    };
 
     fn query_value(url: &str, key: &str) -> Option<String> {
         url::Url::parse(url)
@@ -981,28 +989,8 @@ mod tests {
         }
     }
 
-    struct SubmissionBoundary {
-        request_count: Arc<Mutex<Vec<String>>>,
-        post_calls: Arc<Mutex<Vec<(usize, Value)>>>,
-    }
-
-    impl crate::ocr::hooks::OcrHooks for SubmissionBoundary {
-        fn post_call(
-            &self,
-            request: crate::ocr::hooks::OcrPostCallRequest,
-        ) -> crate::ocr::hooks::OcrHookFuture<'_, crate::ocr::hooks::OcrPostCallRequest> {
-            Box::pin(async move {
-                self.post_calls.lock().unwrap().push((
-                    self.request_count.lock().unwrap().len(),
-                    request.original_response.clone(),
-                ));
-                Ok(request)
-            })
-        }
-    }
-
     #[tokio::test]
-    async fn accepted_response_runs_post_call_for_submission_and_completed_poll() {
+    async fn accepted_response_emits_response_received_for_submission_and_completed_poll() {
         let (base, seen, server) = mock_server(vec![
             MockResponse {
                 status: 202,
@@ -1012,23 +1000,31 @@ mod tests {
             MockResponse::json(json!({"status":"succeeded"})),
         ])
         .await;
-        let post_calls = Arc::new(Mutex::new(Vec::new()));
-        let request = crate::ocr::LiteLLMOcrRequest {
-            hooks: Arc::new(SubmissionBoundary {
-                request_count: seen.clone(),
-                post_calls: post_calls.clone(),
-            }),
-            ..wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}))
-        };
+        let responses_received = Arc::new(Mutex::new(Vec::new()));
+        let request_count = seen.clone();
+        let observed = responses_received.clone();
+        let host = LocalOcrHost::new(wire_request(
+            "azure_ai/doc-intelligence/prebuilt-read",
+            &base,
+            json!({}),
+        ))
+        .with_observer(move |event| {
+            if let CallEvent::ResponseReceived { raw } = event {
+                observed
+                    .lock()
+                    .unwrap()
+                    .push((request_count.lock().unwrap().len(), raw.body.clone()));
+            }
+        });
 
-        perform_ocr(request).await.unwrap();
+        perform_ocr_with(host).await.unwrap();
         server.await.unwrap();
         assert_eq!(seen.lock().unwrap().len(), 2);
         assert_eq!(
-            *post_calls.lock().unwrap(),
+            *responses_received.lock().unwrap(),
             [
-                (1, json!(r#"{"submitted":true}"#)),
-                (2, json!(r#"{"status":"succeeded"}"#)),
+                (1, r#"{"submitted":true}"#.to_string()),
+                (2, r#"{"status":"succeeded"}"#.to_string()),
             ]
         );
     }
@@ -1217,45 +1213,4 @@ mod tests {
             assert!(error.to_string().contains("dot segment"));
         }
     }
-
-    #[tokio::test]
-    async fn pre_call_guardrail_receives_caller_pages_before_mapping() {
-        use std::sync::Arc;
-
-        use crate::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreCallRequest};
-
-        struct RewritePages;
-        impl OcrHooks for RewritePages {
-            fn intercepts_requests(&self) -> bool {
-                true
-            }
-
-            fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
-                Box::pin(async move {
-                    assert_eq!(request.optional_params["pages"], json!([0, 2]));
-                    Ok(OcrPreCallRequest {
-                        optional_params: json!({"pages": [1]}),
-                        ..request
-                    })
-                })
-            }
-        }
-        let (base, seen, server) =
-            mock_server(vec![MockResponse::json(json!({"status": "succeeded"}))]).await;
-        let request = wire_request(
-            "azure_ai/doc-intelligence/prebuilt-read",
-            &base,
-            json!({"pages": [0, 2]}),
-        )
-        .with_host_hooks(Arc::new(RewritePages), None);
-        perform_ocr(request).await.unwrap();
-        server.await.unwrap();
-        let requests = seen.lock().unwrap();
-        let target = requests[0].split_whitespace().nth(1).unwrap();
-        assert_eq!(
-            query_value(&format!("{base}{target}"), "pages").as_deref(),
-            Some("2")
-        );
-        assert_eq!(requests.len(), 1);
-    }
 }

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

@@ -304,12 +304,12 @@ mod tests {
         );
     }
 
-    use std::sync::Arc;
-
     use serde_json::json;
 
-    use crate::ocr::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks};
-    use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
+    use crate::ocr::LocalOcrHost;
+    use crate::ocr::test_support::{
+        MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request,
+    };
 
     #[tokio::test]
     async fn facade_executes_azure_mistral_with_prepared_auth() {
@@ -374,82 +374,242 @@ mod tests {
         );
     }
 
-    struct ReplaceBodyDocument;
-
-    impl OcrHooks for ReplaceBodyDocument {
-        fn intercepts_requests(&self) -> bool {
-            true
-        }
-
-        fn during_call(
-            &self,
-            mut request: OcrDuringCallRequest,
-        ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
-            Box::pin(async move {
-                request.body["document"] = json!({
-                    "type":"document_url",
-                    "document_url":"https://example.com/not-inline.pdf"
-                });
-                Ok(request)
-            })
-        }
-    }
-
     #[tokio::test]
     async fn rejects_non_inline_body_after_guardrails() {
-        let mut request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({}));
-        request.hooks = Arc::new(ReplaceBodyDocument);
-        let error = perform_ocr(request).await.unwrap_err();
+        let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({}));
+        let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| {
+            wire.body["document"] = json!({
+                "type":"document_url",
+                "document_url":"https://example.com/not-inline.pdf"
+            });
+            Ok(wire)
+        });
+        let error = perform_ocr_with(host).await.unwrap_err();
         assert!(error.to_string().contains("data URI"));
     }
 
-    struct EchoCallerDocument(Value);
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicUsize, Ordering};
 
-    impl OcrHooks for EchoCallerDocument {
-        fn intercepts_requests(&self) -> bool {
-            true
+    use litellm_auth::{
+        ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle,
+    };
+
+    use crate::ocr::LiteLLMOcrRequest;
+    use crate::ocr::test_support::header;
+    use crate::ocr::wire::decode_request;
+
+    #[derive(Debug)]
+    struct CountingToken {
+        token: fn(usize) -> String,
+        calls: AtomicUsize,
+    }
+
+    impl CountingToken {
+        fn new(token: fn(usize) -> String) -> Arc<Self> {
+            Arc::new(Self {
+                token,
+                calls: AtomicUsize::new(0),
+            })
         }
 
-        fn during_call(
-            &self,
-            mut request: OcrDuringCallRequest,
-        ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
-            let document = self.0.clone();
+        fn calls(&self) -> usize {
+            self.calls.load(Ordering::SeqCst)
+        }
+    }
+
+    impl TokenProvider for CountingToken {
+        fn acquire(&self) -> TokenFuture<'_> {
+            let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
+            let token = SecretValue::new((self.token)(call));
             Box::pin(async move {
-                request.body["document"] = document;
-                Ok(request)
+                Ok(ResolvedCredential::AccessToken {
+                    token,
+                    expires_on: None,
+                })
             })
         }
     }
 
-    #[tokio::test]
-    async fn remote_document_stays_inlined_when_hook_echoes_caller_document() {
-        let (base, seen, server) = mock_server(vec![
-            MockResponse::json(json!("served document")),
-            MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}],"usage_info":{"pages_processed":1}})),
-        ])
-        .await;
-        let document_url = format!("{base}/document.pdf");
-        let mut request = crate::ocr::test_support::with_source(
-            wire_request("azure_ai/model", &base, json!({})),
-            &document_url,
-        );
-        request.hooks = Arc::new(EchoCallerDocument(
-            json!({"type":"document_url","document_url":document_url}),
-        ));
+    fn numbered_token(call: usize) -> String {
+        format!("callback-{call}")
+    }
 
-        let result = perform_ocr(request).await.unwrap();
+    fn azure_request(
+        provider: &Arc<CountingToken>,
+        api_base: Option<&str>,
+        api_key: Option<&str>,
+        extra_headers: Value,
+        optional_params: Value,
+    ) -> LiteLLMOcrRequest {
+        let wire = serde_json::from_value(json!({
+            "model": "azure_ai/mistral-ocr-latest",
+            "document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
+            "api_key": api_key,
+            "api_base": api_base,
+            "custom_llm_provider": null,
+            "extra_headers": extra_headers,
+            "optional_params": optional_params,
+            "timeout_seconds": 2.0
+        }))
+        .unwrap();
+        LiteLLMOcrRequest {
+            azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())),
+            ..decode_request(wire).unwrap()
+        }
+    }
+
+    fn ocr_page() -> MockResponse {
+        MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]}))
+    }
+
+    #[tokio::test]
+    async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() {
+        let provider = CountingToken::new(numbered_token);
+        let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await;
+
+        for _ in 0..2 {
+            perform_ocr(azure_request(
+                &provider,
+                Some(&base),
+                None,
+                Value::Null,
+                json!({}),
+            ))
+            .await
+            .unwrap();
+        }
         server.await.unwrap();
 
-        assert_eq!(result.pages[0].markdown, "hello");
+        assert_eq!(provider.calls(), 2);
         let requests = seen.lock().unwrap();
-        assert_eq!(requests.len(), 2);
-        assert!(requests[0].starts_with("GET /document.pdf "));
-        let body: Value =
-            serde_json::from_str(requests[1].split_once("\r\n\r\n").unwrap().1).unwrap();
         assert_eq!(
-            body["document"]["document_url"],
-            json!("data:application/json;base64,InNlcnZlZCBkb2N1bWVudCI=")
+            requests
+                .iter()
+                .map(|request| header(request, "authorization"))
+                .collect::<Vec<_>>(),
+            [Some("Bearer callback-1"), Some("Bearer callback-2")]
         );
     }
+
+    #[rstest]
+    #[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)]
+    #[case::provider_beats_static_token(
+        None,
+        Value::Null,
+        json!({"azure_ad_token":"static-token"}),
+        "Bearer callback-1",
+        1
+    )]
+    #[case::header_wins_on_the_wire_but_provider_still_runs(
+        None,
+        json!({"Authorization":"Bearer override"}),
+        json!({}),
+        "Bearer override",
+        1
+    )]
+    #[tokio::test]
+    async fn credential_precedence(
+        #[case] api_key: Option<&str>,
+        #[case] extra_headers: Value,
+        #[case] optional_params: Value,
+        #[case] expected_authorization: &str,
+        #[case] expected_calls: usize,
+    ) {
+        let provider = CountingToken::new(numbered_token);
+        let (base, seen, server) = mock_server(vec![ocr_page()]).await;
+
+        perform_ocr(azure_request(
+            &provider,
+            Some(&base),
+            api_key,
+            extra_headers,
+            optional_params,
+        ))
+        .await
+        .unwrap();
+        server.await.unwrap();
+
+        assert_eq!(provider.calls(), expected_calls);
+        let requests = seen.lock().unwrap();
+        assert_eq!(requests.len(), 1);
+        assert_eq!(
+            header(&requests[0], "authorization"),
+            Some(expected_authorization)
+        );
+    }
+
+    #[rstest]
+    #[case::missing_api_base(
+        false,
+        json!({}),
+        numbered_token,
+        |error: &crate::ocr::Error| matches!(error, crate::ocr::Error::Auth(litellm_auth::Error::MissingApiBase {
+            provider: "Azure AI",
+            environment_variable: AZURE_AI_API_BASE_ENV,
+        })),
+        0
+    )]
+    #[case::unsupported_oidc_reference(
+        true,
+        json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}),
+        numbered_token,
+        |error: &crate::ocr::Error| matches!(error, crate::ocr::Error::Auth(litellm_auth::Error::UnsupportedOidcReference)),
+        0
+    )]
+    #[case::empty_provider_token_ignores_static_token(
+        true,
+        json!({"azure_ad_token":"static-token"}),
+        |_| String::new(),
+        |error: &crate::ocr::Error| matches!(error, crate::ocr::Error::MissingAzureAiCredentials),
+        1
+    )]
+    #[tokio::test]
+    async fn credential_failures_send_no_provider_request(
+        #[case] with_api_base: bool,
+        #[case] optional_params: Value,
+        #[case] token: fn(usize) -> String,
+        #[case] expected: fn(&crate::ocr::Error) -> bool,
+        #[case] expected_calls: usize,
+    ) {
+        let provider = CountingToken::new(token);
+        let (base, seen, server) = mock_server(vec![ocr_page()]).await;
+
+        let error = perform_ocr(azure_request(
+            &provider,
+            with_api_base.then_some(base.as_str()),
+            None,
+            Value::Null,
+            optional_params,
+        ))
+        .await
+        .unwrap_err();
+        server.abort();
+
+        assert!(expected(&error), "unexpected error: {error:?}");
+        assert_eq!(provider.calls(), expected_calls);
+        assert!(seen.lock().unwrap().is_empty());
+    }
+
+    #[tokio::test]
+    async fn environment_supplies_api_base_and_bearer_key() {
+        let env = |name: &str| match name {
+            AZURE_AI_API_BASE_ENV => Some("https://env.example".to_string()),
+            AZURE_AI_API_KEY_ENV => Some("env-key".to_string()),
+            _ => None,
+        };
+        let connection = OcrConnection::default();
+
+        let headers = AzureAiOcrConfig
+            .resolve_headers(&connection, &Default::default(), &env)
+            .await
+            .unwrap();
+        let url = AzureAiOcrConfig.build_ocr_url(None, &env).unwrap();
+
+        assert_eq!(
+            headers,
+            [("Authorization".to_string(), "Bearer env-key".to_string())]
+        );
+        assert_eq!(url, "https://env.example/providers/mistral/azure/ocr");
+    }
 }

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

@@ -1,16 +1,18 @@
 use std::future::Future;
-use std::sync::Arc;
 
-use serde::Serialize;
-use serde::de::DeserializeOwned;
+use serde::{Serialize, de::DeserializeOwned};
 use serde_json::Value;
 
-use crate::call_arguments::CallArguments;
-use crate::ocr::OcrClient;
-use crate::ocr::hooks::OcrHooks;
-use crate::ocr::types::{
-    LiteLLMOcrResponse, OcrConnection, OcrCredentialInputs, OcrDocument, OcrResponseFormat,
-    PreparedOcrRequest, ResolvedOcrCredentials,
+use crate::{
+    call_arguments::CallArguments,
+    ocr::{
+        OcrClient,
+        route::OcrHost,
+        types::{
+            LiteLLMOcrResponse, OcrConnection, OcrCredentialInputs, OcrDocument, OcrResponseFormat,
+            PreparedOcrRequest, ResolvedOcrCredentials,
+        },
+    },
 };
 
 const HEALTH_CHECK_PDF_DATA_URI: &str = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y=";
@@ -37,7 +39,7 @@ pub(crate) struct OcrRequestContext<'a> {
 pub(crate) struct OcrResponseContext<'a> {
     pub client: &'a OcrClient,
     pub connection: &'a OcrConnection,
-    pub hooks: &'a Arc<dyn OcrHooks>,
+    pub host: &'a OcrHost,
     pub request_format: OcrResponseFormat,
     pub url: &'a str,
     pub headers: &'a [(String, String)],
@@ -133,7 +135,7 @@ pub(crate) trait BaseOcrConfig: Send + Sync + Sized + 'static {
                 context.connection.max_response_bytes,
             )
             .await?;
-            crate::ocr::handler::post_call(context.hooks, &bytes).await?;
+            crate::ocr::handler::emit_response_received(context.host, &bytes).await?;
             self.transform_ocr_response(model, &bytes, context.request_format)
         }
     }

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

@@ -2,18 +2,22 @@ use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 use serde_with::serde_as;
 
-use crate::call_arguments::{CallArguments, parse_options};
-use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE};
-use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response};
-use crate::ocr::OcrClient;
-use crate::ocr::document::InlineDocument;
-use crate::ocr::prepare::credential_env;
-use crate::ocr::types::{
-    LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrPageImage, OcrResponseFormat,
-    OcrUsageInfo, PreparedOcrRequest,
+use crate::{
+    call_arguments::{CallArguments, parse_options},
+    constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE},
+    llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response},
+    ocr::{
+        OcrClient,
+        document::InlineDocument,
+        prepare::credential_env,
+        types::{
+            LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrPageImage,
+            OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest,
+        },
+    },
+    serde_compat::LaxI64,
+    url_utils::ApiUrl,
 };
-use crate::serde_compat::LaxI64;
-use crate::url_utils::ApiUrl;
 
 const COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC";
 
@@ -339,7 +343,7 @@ mod tests {
             "cohere/parse",
             "https://example.com",
             json!({
-                "output_format":"markdown", "metadata":{"host":true},
+                "output_format":"markdown", "timeout":30,
                 "extra_body":{
                     "output_format": {"future":true},
                     "document":{"type":"image_url","image_url":"https://example.com/a.png",
@@ -353,7 +357,7 @@ mod tests {
             }))
             .unwrap(),
         );
-        let request = crate::ocr::prepare::prepare_request(request);
+        let request = crate::ocr::prepare::prepare_request_for_test(request);
         let http = CohereParseConfig
             .prepare_request(&request, &crate::ocr::test_support::ocr_client())
             .await
@@ -512,7 +516,7 @@ mod tests {
             request.response_format().unwrap(),
             crate::ocr::types::OcrResponseFormat::Litellm
         );
-        let request = crate::ocr::prepare::prepare_request(request);
+        let request = crate::ocr::prepare::prepare_request_for_test(request);
         let http = CohereParseConfig
             .prepare_request(&request, &crate::ocr::test_support::ocr_client())
             .await
@@ -747,15 +751,19 @@ mod tests {
     }
 
     #[rstest]
-    #[case::base("")]
-    #[case::version("/v2")]
-    #[case::complete("/v2/parse")]
-    fn completes_provider_urls_without_duplicate_paths_and_preserves_queries(#[case] suffix: &str) {
+    #[case::base("", "/v2/parse")]
+    #[case::version("/v2", "/v2/parse")]
+    #[case::complete("/v2/parse", "/v2/parse")]
+    #[case::proxy_prefix("/cohere/", "/cohere/v2/parse")]
+    fn completes_provider_urls_without_duplicate_paths_and_preserves_queries(
+        #[case] suffix: &str,
+        #[case] path: &str,
+    ) {
         assert_eq!(
             CohereParseConfig
                 .build_ocr_url(&format!("https://example.com{suffix}?tenant=a"))
                 .unwrap(),
-            "https://example.com/v2/parse?tenant=a"
+            format!("https://example.com{path}?tenant=a")
         );
     }
 
@@ -779,4 +787,84 @@ mod tests {
             Err(crate::ocr::Error::Auth(_))
         ));
     }
+
+    #[test]
+    fn environment_key_becomes_the_bearer() {
+        let headers = CohereParseConfig
+            .resolve_headers(&OcrConnection::default(), &|name| {
+                (name == COHERE_API_KEY_ENV).then(|| "env-key".to_string())
+            })
+            .unwrap();
+
+        assert_eq!(
+            headers,
+            [("Authorization".to_string(), "Bearer env-key".to_string())]
+        );
+    }
+
+    #[test]
+    fn missing_key_names_the_environment_variable() {
+        let error = CohereParseConfig
+            .resolve_headers(&OcrConnection::default(), &|_| None)
+            .unwrap_err();
+
+        assert!(error.to_string().contains(COHERE_API_KEY_ENV), "{error}");
+    }
+
+    #[rstest]
+    #[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")]
+    #[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")]
+    #[tokio::test]
+    async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key(
+        #[case] model: &str,
+        #[case] request_line: &str,
+    ) {
+        use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr};
+
+        let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
+        let request = crate::ocr::test_support::wire_request(model, &base, json!({}))
+            .with_document(
+                serde_json::from_value::<OcrDocument>(
+                    json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}),
+                )
+                .unwrap()
+                .into(),
+            );
+
+        perform_ocr(request).await.unwrap();
+        server.await.unwrap();
+
+        let requests = seen.lock().unwrap();
+        assert_eq!(requests.len(), 1);
+        assert!(requests[0].starts_with(request_line), "{}", requests[0]);
+        assert_eq!(
+            header(&requests[0], "authorization"),
+            Some("Bearer test-key")
+        );
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn route_rejects_non_image_document_without_a_request(
+        #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str,
+    ) {
+        use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr};
+
+        let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
+
+        let error = perform_ocr(crate::ocr::test_support::wire_request(
+            model,
+            &base,
+            json!({}),
+        ))
+        .await
+        .unwrap_err();
+        server.abort();
+
+        assert!(
+            matches!(error, crate::ocr::Error::CohereImageOnly),
+            "{error:?}"
+        );
+        assert!(seen.lock().unwrap().is_empty());
+    }
 }

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

@@ -1,17 +1,21 @@
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
-use crate::call_arguments::CallArguments;
-use crate::constants::MISTRAL_OCR_API_BASE;
-use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response};
-use crate::ocr::OcrClient;
-use crate::ocr::prepare::credential_env;
-use crate::ocr::types::{
-    LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat, OcrUsageInfo,
-    PreparedOcrRequest,
+use crate::{
+    call_arguments::CallArguments,
+    constants::MISTRAL_OCR_API_BASE,
+    llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response},
+    ocr::{
+        OcrClient,
+        prepare::credential_env,
+        types::{
+            LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat,
+            OcrUsageInfo, PreparedOcrRequest,
+        },
+    },
+    params::OpaqueParams,
+    url_utils::ApiUrl,
 };
-use crate::params::OpaqueParams;
-use crate::url_utils::ApiUrl;
 
 const MISTRAL_OCR_API_KEY_ENV_VAR: &str = "MISTRAL_API_KEY";
 
@@ -618,6 +622,22 @@ mod tests {
         );
     }
 
+    #[rstest]
+    fn environment_keeps_extra_headers_after_the_bearer_key(
+        #[with(Some("explicit"), vec![("X-Trace".into(), "trace-1".into())])]
+        connection: OcrConnection,
+    ) {
+        assert_eq!(
+            MistralOcrConfig
+                .resolve_headers(&connection, &|_| None)
+                .unwrap(),
+            [
+                ("Authorization".to_string(), "Bearer explicit".to_string()),
+                ("X-Trace".to_string(), "trace-1".to_string()),
+            ]
+        );
+    }
+
     #[rstest]
     fn environment_rejects_missing_key(connection: OcrConnection) {
         assert!(matches!(

litellm-rust/crates/core/src/llms/openai/responses/transformation.rs

@@ -1,6 +1,8 @@
-use crate::responses::Error;
-use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult};
-use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model};
+use crate::responses::{
+    Error,
+    types::{ResponsesWsEvent, ResponsesWsTransformResult},
+    websocket::{ResponsesWebSocketProviderConfig, enforce_model},
+};
 
 pub struct OpenAiResponsesApiConfig;
 

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

@@ -3,20 +3,24 @@ use std::collections::BTreeMap;
 use serde::{Deserialize, Deserializer, Serialize};
 use serde_json::{Map, Value, json};
 
-use crate::call_arguments::{CallArguments, compose_body};
-use crate::constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX};
-use crate::llms::base_llm::ocr::transformation::{
-    BaseOcrConfig, OcrRequestContext, decode_and_normalize_response,
+use crate::{
+    call_arguments::{CallArguments, compose_body},
+    constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX},
+    llms::base_llm::ocr::transformation::{
+        BaseOcrConfig, OcrRequestContext, decode_and_normalize_response,
+    },
+    ocr::{
+        OcrClient,
+        document::InlineDocument,
+        prepare::{build_http_request, credential_env, guardrail_document},
+        types::{
+            LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat,
+            OcrUsageInfo, PreparedOcrRequest,
+        },
+    },
+    params::OpaqueParams,
+    url_utils::ApiUrl,
 };
-use crate::ocr::OcrClient;
-use crate::ocr::document::InlineDocument;
-use crate::ocr::prepare::{build_http_request, credential_env, guardrail_document};
-use crate::ocr::types::{
-    LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat, OcrUsageInfo,
-    PreparedOcrRequest,
-};
-use crate::params::OpaqueParams;
-use crate::url_utils::ApiUrl;
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 #[serde(transparent)]
@@ -682,12 +686,13 @@ mod tests {
         );
     }
 
-    use std::sync::Arc;
-
+    use litellm_callbacks::event::{CallEvent, WireRequest};
     use rstest::rstest;
 
-    use crate::ocr::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrPostCallRequest};
-    use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
+    use crate::ocr::{
+        LocalOcrHost,
+        test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
+    };
 
     fn request_body(request: &str) -> Value {
         serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
@@ -783,38 +788,23 @@ mod tests {
         assert!(requests[1].starts_with("POST /parse "));
     }
 
-    struct ParseBoundary {
-        request_count: Arc<std::sync::Mutex<Vec<String>>>,
-    }
-
-    impl OcrHooks for ParseBoundary {
-        fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> {
-            Box::pin(async move {
-                assert_eq!(self.request_count.lock().unwrap().len(), 2);
-                assert_eq!(
-                    request.original_response,
-                    json!(r#"{"result":{"chunks":[]}}"#)
-                );
-                Ok(request)
-            })
-        }
-    }
-
     #[tokio::test]
-    async fn post_call_stays_after_reducto_upload_and_parse() {
+    async fn response_received_stays_after_reducto_upload_and_parse() {
         let (base, seen, server) = mock_server(vec![
             MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
             MockResponse::json(json!({"result":{"chunks":[]}})),
         ])
         .await;
-        let request = crate::ocr::LiteLLMOcrRequest {
-            hooks: Arc::new(ParseBoundary {
-                request_count: seen.clone(),
-            }),
-            ..wire_request("reducto/parse-v3", &base, json!({}))
-        };
+        let request_count = seen.clone();
+        let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({})))
+            .with_observer(move |event| {
+                if let CallEvent::ResponseReceived { raw } = event {
+                    assert_eq!(request_count.lock().unwrap().len(), 2);
+                    assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#);
+                }
+            });
 
-        perform_ocr(request).await.unwrap();
+        perform_ocr_with(host).await.unwrap();
         server.await.unwrap();
         assert_eq!(seen.lock().unwrap().len(), 2);
     }
@@ -932,28 +922,6 @@ mod tests {
         );
     }
 
-    struct RewriteDocument;
-
-    struct RewriteHeaders;
-
-    impl OcrHooks for RewriteHeaders {
-        fn intercepts_requests(&self) -> bool {
-            true
-        }
-
-        fn during_call(
-            &self,
-            request: OcrDuringCallRequest,
-        ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
-            Box::pin(async move {
-                Ok(OcrDuringCallRequest {
-                    headers: vec![("authorization".into(), "Bearer guarded".into())],
-                    ..request
-                })
-            })
-        }
-    }
-
     #[rstest]
     #[case("reducto/parse-v3")]
     #[case("reducto/parse-legacy")]
@@ -966,9 +934,14 @@ mod tests {
         .await;
         let mut request = wire_request(model, &base, json!({}));
         request.transport.extra_headers = vec![("authorization".into(), "Bearer original".into())];
-        request.hooks = Arc::new(RewriteHeaders);
+        let host = LocalOcrHost::new(request).with_before_send(|wire, _| {
+            Ok(WireRequest {
+                headers: vec![("authorization".into(), "Bearer guarded".into())],
+                ..wire
+            })
+        });
 
-        perform_ocr(request).await.unwrap();
+        perform_ocr_with(host).await.unwrap();
         server.await.unwrap();
         let requests = seen.lock().unwrap();
         assert_eq!(requests.len(), 2);
@@ -980,36 +953,23 @@ mod tests {
         }
     }
 
-    impl OcrHooks for RewriteDocument {
-        fn intercepts_requests(&self) -> bool {
-            true
-        }
-
-        fn during_call(
-            &self,
-            request: OcrDuringCallRequest,
-        ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
-            Box::pin(async move {
-                assert_eq!(
-                    request.body["document_url"],
-                    "data:application/pdf;base64,YWJj"
-                );
-                Ok(OcrDuringCallRequest {
-                    body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}),
-                    ..request
-                })
-            })
-        }
-    }
-
     #[tokio::test]
     async fn guardrail_rewrites_document_before_upload() {
         let (base, seen, server) =
             mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await;
-        let mut request = wire_request("reducto/parse-v3", &base, json!({}));
-        request.hooks = Arc::new(RewriteDocument);
+        let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({})))
+            .with_before_send(|wire, _| {
+                assert_eq!(
+                    wire.body["document_url"],
+                    "data:application/pdf;base64,YWJj"
+                );
+                Ok(WireRequest {
+                    body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}),
+                    ..wire
+                })
+            });
 
-        perform_ocr(request).await.unwrap();
+        perform_ocr_with(host).await.unwrap();
         server.await.unwrap();
         let requests = seen.lock().unwrap();
         assert_eq!(requests.len(), 1);

litellm-rust/crates/core/src/llms/vertex_ai/ocr/deepseek_transformation.rs

@@ -3,16 +3,20 @@ use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 
 use super::transformation::VertexAiOcrConfig;
-use crate::call_arguments::CallArguments;
-use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext};
-use crate::ocr::OcrClient;
-use crate::ocr::prepare::credential_env;
-use crate::ocr::types::{
-    LiteLLMOcrResponse, OcrDocument, OcrPage, OcrPageDimensions, OcrPageImage, OcrUsageInfo,
-    PreparedOcrRequest,
+use crate::{
+    call_arguments::CallArguments,
+    llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext},
+    ocr::{
+        OcrClient,
+        prepare::credential_env,
+        types::{
+            LiteLLMOcrResponse, OcrDocument, OcrPage, OcrPageDimensions, OcrPageImage,
+            OcrUsageInfo, PreparedOcrRequest,
+        },
+    },
+    params::OpaqueParams,
+    url_utils::ApiUrl,
 };
-use crate::params::OpaqueParams;
-use crate::url_utils::ApiUrl;
 
 const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com";
 const MODEL_PREFIX: &str = "deepseek-ai/";
@@ -456,8 +460,7 @@ mod tests {
 
     use rstest::rstest;
 
-    use crate::llms::base_llm::ocr::transformation::BaseOcrConfig;
-    use crate::ocr::types::OcrDocument;
+    use crate::{llms::base_llm::ocr::transformation::BaseOcrConfig, ocr::types::OcrDocument};
 
     fn document() -> OcrDocument {
         serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap()

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

@@ -2,17 +2,21 @@ use litellm_auth_gcp::{self as vertex, VertexConfig};
 use serde_json::Value;
 
 use super::common_utils::validate_destination;
-use crate::call_arguments::CallArguments;
-use crate::llms::base_llm::ocr::transformation::{
-    BaseOcrConfig, OcrEnvironment, OcrRequestContext,
+use crate::{
+    call_arguments::CallArguments,
+    llms::{
+        base_llm::ocr::transformation::{BaseOcrConfig, OcrEnvironment, OcrRequestContext},
+        mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest},
+    },
+    ocr::{
+        OcrClient,
+        document::{inline_remote_document, validate_inline_document},
+        prepare::credential_env,
+        types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, PreparedOcrRequest},
+    },
+    params::OpaqueParams,
+    url_utils::ApiUrl,
 };
-use crate::llms::mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest};
-use crate::ocr::OcrClient;
-use crate::ocr::document::{inline_remote_document, validate_inline_document};
-use crate::ocr::prepare::credential_env;
-use crate::ocr::types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, PreparedOcrRequest};
-use crate::params::OpaqueParams;
-use crate::url_utils::ApiUrl;
 
 const DEFAULT_LOCATION: &str = "us-central1";
 
@@ -198,9 +202,10 @@ fn validate_location(location: &str) -> Result<(), crate::ocr::Error> {
 
 #[cfg(test)]
 mod tests {
-    use super::VertexAiOcrConfig;
     use rstest::rstest;
 
+    use super::VertexAiOcrConfig;
+
     #[test]
     fn endpoint_uses_location_project_and_model() {
         assert_eq!(
@@ -329,10 +334,14 @@ mod tests {
     ) {
         use std::time::Duration;
 
-        use crate::llms::base_llm::ocr::transformation::BaseOcrConfig;
-        use crate::llms::mistral::ocr::transformation::MistralOcrConfig;
-        use crate::llms::vertex_ai::ocr::transformation::VertexAiOcrConfig;
-        use crate::ocr::test_support::ocr_client;
+        use crate::{
+            llms::{
+                base_llm::ocr::transformation::BaseOcrConfig,
+                mistral::ocr::transformation::MistralOcrConfig,
+                vertex_ai::ocr::transformation::VertexAiOcrConfig,
+            },
+            ocr::test_support::ocr_client,
+        };
 
         let client = ocr_client();
         let options = json!({
@@ -348,10 +357,10 @@ mod tests {
             options.clone(),
         );
         let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options);
-        let direct = crate::ocr::prepare::prepare_request(
+        let direct = crate::ocr::prepare::prepare_request_for_test(
             crate::ocr::test_support::resolved_request(direct),
         );
-        let vertex = crate::ocr::prepare::prepare_request(
+        let vertex = crate::ocr::prepare::prepare_request_for_test(
             crate::ocr::test_support::resolved_request(vertex),
         );
         let direct_http = MistralOcrConfig

litellm-rust/crates/core/src/machine/auth.rs

@@ -0,0 +1,53 @@
+use std::sync::Arc;
+
+use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
+use litellm_callbacks::route::Route;
+
+use super::{HostChannel, MachineFault};
+
+/// A route whose host can mint credentials on the call's behalf.
+pub trait TokenRoute: Route {
+    fn acquire_token_op() -> Self::Op;
+    fn token_credential(result: Self::OpResult) -> Option<ResolvedCredential>;
+}
+
+/// A [`TokenProvider`] that asks the host for each credential through the call's own
+/// operation channel, so the host answers it on the caller's thread and context.
+pub struct HostTokenProvider<R: Route> {
+    channel: HostChannel<R>,
+}
+
+impl<R: Route> std::fmt::Debug for HostTokenProvider<R> {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter.write_str("HostTokenProvider")
+    }
+}
+
+impl<R> HostTokenProvider<R>
+where
+    R: TokenRoute,
+    R::Error: From<MachineFault> + std::fmt::Display,
+{
+    pub fn handle(channel: HostChannel<R>) -> TokenProviderHandle {
+        TokenProviderHandle::new(Arc::new(Self { channel }))
+    }
+}
+
+impl<R> TokenProvider for HostTokenProvider<R>
+where
+    R: TokenRoute,
+    R::Error: From<MachineFault> + std::fmt::Display,
+{
+    fn acquire(&self) -> TokenFuture<'_> {
+        Box::pin(async move {
+            let result = self
+                .channel
+                .route(R::acquire_token_op())
+                .await
+                .map_err(|error| Error::AzureTokenAcquisition(error.to_string()))?;
+            R::token_credential(result).ok_or_else(|| {
+                Error::AzureTokenAcquisition("invalid token provider host result".into())
+            })
+        })
+    }
+}

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

@@ -0,0 +1,202 @@
+//! The one machine every route runs on: it owns the route's provider future, polls it in
+//! place, and turns the host operations that future requests into [`Machine`] steps. No
+//! task is spawned; dropping the machine drops the in-flight call.
+
+mod auth;
+
+use std::{future::Future, pin::Pin};
+
+pub use auth::{HostTokenProvider, TokenRoute};
+use litellm_callbacks::{
+    event::{CallEvent, RequestContext, WireRequest},
+    host::{HostOp, HostResult},
+    machine::{HostFailure, Interrupted, Machine, MachineStep, Step},
+    route::Route,
+};
+use tokio::sync::{mpsc, oneshot};
+
+/// The machine's own failures, distinct from anything the provider call reports.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum MachineFault {
+    /// The host driver went away while the call was waiting on it.
+    Abandoned,
+    /// The host answered out of turn: a result with nothing pending, or nothing when a
+    /// result was pending.
+    Protocol(&'static str),
+    /// The host answered a route operation with the wrong result variant.
+    Mismatch,
+}
+
+pub type ExecuteFuture<R> =
+    Pin<Box<dyn Future<Output = Result<<R as Route>::Response, <R as Route>::Error>> + Send>>;
+
+struct PendingOp<R: Route> {
+    op: HostOp<R>,
+    reply: oneshot::Sender<HostResult<R>>,
+}
+
+/// The provider side of the machine: how the in-flight call reaches its host.
+pub struct HostChannel<R: Route> {
+    ops: Option<mpsc::UnboundedSender<PendingOp<R>>>,
+}
+
+impl<R: Route> Clone for HostChannel<R> {
+    fn clone(&self) -> Self {
+        Self {
+            ops: self.ops.clone(),
+        }
+    }
+}
+
+impl<R: Route> HostChannel<R> {
+    /// A channel with no host behind it: the wire request goes out unchanged, events go
+    /// nowhere, and route operations fail. For tests that prepare a request without
+    /// driving it.
+    #[cfg(test)]
+    pub(crate) fn detached() -> Self {
+        Self { ops: None }
+    }
+}
+
+impl<R: Route> HostChannel<R>
+where
+    R::Error: From<MachineFault>,
+{
+    async fn invoke(&self, op: HostOp<R>) -> Result<HostResult<R>, R::Error> {
+        let ops = self.ops.as_ref().ok_or(MachineFault::Abandoned)?;
+        let (reply, answer) = oneshot::channel();
+        ops.send(PendingOp { op, reply })
+            .map_err(|_| MachineFault::Abandoned)?;
+        answer.await.map_err(|_| MachineFault::Abandoned.into())
+    }
+
+    pub async fn route(&self, op: R::Op) -> Result<R::OpResult, R::Error> {
+        match self.invoke(HostOp::Route(op)).await? {
+            HostResult::Route(result) => Ok(result),
+            _ => Err(MachineFault::Mismatch.into()),
+        }
+    }
+
+    pub async fn before_send(
+        &self,
+        wire: WireRequest,
+        context: RequestContext,
+    ) -> Result<WireRequest, R::Error> {
+        if self.ops.is_none() {
+            return Ok(wire);
+        }
+        let op = HostOp::BeforeSend {
+            wire: Box::new(wire),
+            context: Box::new(context),
+        };
+        match self.invoke(op).await? {
+            HostResult::BeforeSend(wire) => Ok(*wire),
+            _ => Err(MachineFault::Mismatch.into()),
+        }
+    }
+
+    pub async fn emit(&self, event: CallEvent) -> Result<(), R::Error> {
+        if self.ops.is_none() {
+            return Ok(());
+        }
+        match self.invoke(HostOp::Emit(event)).await? {
+            HostResult::Emitted => Ok(()),
+            _ => Err(MachineFault::Mismatch.into()),
+        }
+    }
+}
+
+enum Execution<R: Route> {
+    Unstarted(Box<dyn FnOnce(HostChannel<R>) -> ExecuteFuture<R> + Send>),
+    Running(ExecuteFuture<R>),
+    Done,
+}
+
+pub struct RouteMachine<R: Route> {
+    execution: Execution<R>,
+    ops: mpsc::UnboundedReceiver<PendingOp<R>>,
+    channel: HostChannel<R>,
+    reply: Option<oneshot::Sender<HostResult<R>>>,
+}
+
+impl<R: Route> RouteMachine<R>
+where
+    R::Error: From<MachineFault>,
+{
+    pub fn new(execute: impl FnOnce(HostChannel<R>) -> ExecuteFuture<R> + Send + 'static) -> Self {
+        let (ops_tx, ops) = mpsc::unbounded_channel();
+        Self {
+            execution: Execution::Unstarted(Box::new(execute)),
+            ops,
+            channel: HostChannel { ops: Some(ops_tx) },
+            reply: None,
+        }
+    }
+
+    async fn step(
+        &mut self,
+        result: Option<HostResult<R>>,
+    ) -> Result<MachineStep<R, R::Response>, R::Error> {
+        match (self.reply.take(), result) {
+            (Some(reply), Some(result)) => {
+                reply
+                    .send(result)
+                    .map_err(|_| MachineFault::Protocol("the call stopped waiting on the host"))?;
+            }
+            (None, None) if matches!(self.execution, Execution::Unstarted(_)) => {}
+            (Some(reply), None) => {
+                self.reply = Some(reply);
+                return Err(MachineFault::Protocol("host operation result is required").into());
+            }
+            (None, Some(_)) => {
+                return Err(MachineFault::Protocol("unexpected host operation result").into());
+            }
+            (None, None) => {
+                return Err(
+                    MachineFault::Protocol("call cannot be resumed after completion").into(),
+                );
+            }
+        }
+        if let Execution::Unstarted(_) = self.execution {
+            let Execution::Unstarted(start) =
+                std::mem::replace(&mut self.execution, Execution::Done)
+            else {
+                unreachable!()
+            };
+            self.execution = Execution::Running(start(self.channel.clone()));
+        }
+        let Execution::Running(future) = &mut self.execution else {
+            return Err(MachineFault::Protocol("call cannot be resumed after completion").into());
+        };
+        tokio::select! {
+            biased;
+            pending = self.ops.recv() => {
+                let pending = pending.ok_or(MachineFault::Abandoned)?;
+                self.reply = Some(pending.reply);
+                Ok(MachineStep::Host(pending.op))
+            }
+            outcome = future => {
+                self.execution = Execution::Done;
+                outcome.map(MachineStep::Complete)
+            }
+        }
+    }
+}
+
+impl<R: Route> Machine for RouteMachine<R>
+where
+    R::Error: From<MachineFault>,
+{
+    type Route = R;
+    type Complete = R::Response;
+
+    fn resume(&mut self, result: Option<HostResult<R>>) -> Step<'_, Self> {
+        Box::pin(self.step(result))
+    }
+
+    fn interrupt(&mut self, failure: HostFailure<R::Error>) -> Interrupted<'_, Self> {
+        self.reply = None;
+        self.execution = Execution::Done;
+        Box::pin(async move { Err(failure.into_error()) })
+    }
+}

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

@@ -1,12 +1,16 @@
-use std::future::Future;
-use std::io;
-use std::net::{IpAddr, SocketAddr};
-use std::pin::Pin;
-use std::sync::Arc;
-use std::time::Duration;
+use std::{
+    future::Future,
+    io,
+    net::{IpAddr, SocketAddr},
+    pin::Pin,
+    sync::Arc,
+    time::Duration,
+};
 
-use reqwest::Url;
-use reqwest::dns::{Addrs, Name, Resolve, Resolving};
+use reqwest::{
+    Url,
+    dns::{Addrs, Name, Resolve, Resolving},
+};
 
 use crate::constants::MEDIA_CONNECT_TIMEOUT_SECS;
 
@@ -281,8 +285,10 @@ impl Resolve for PublicDnsResolver {
 mod tests {
     use std::collections::HashSet;
 
-    use tokio::io::{AsyncReadExt, AsyncWriteExt};
-    use tokio::net::TcpListener;
+    use tokio::{
+        io::{AsyncReadExt, AsyncWriteExt},
+        net::TcpListener,
+    };
 
     use super::*;
 

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

@@ -1,5 +1,4 @@
-use std::sync::OnceLock;
-use std::time::Duration;
+use std::{sync::OnceLock, time::Duration};
 
 use crate::constants::{MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS};
 

litellm-rust/crates/core/src/messages/common_utils.rs

@@ -1,11 +1,13 @@
+use litellm_providers::{
+    anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
+    azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
+    base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig,
+};
 use serde_json::{Map, Value};
 
 use super::Error;
 use crate::http_utils::string_headers as shared_string_headers;
 pub(super) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_body};
-use litellm_providers::anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
-use litellm_providers::azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
-use litellm_providers::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig;
 
 const HEADER_CONTEXT: &str = "messages";
 

litellm-rust/crates/core/src/messages/handler.rs

@@ -1,10 +1,11 @@
-use super::Error;
-use super::client::http_client;
-use super::common_utils::truncate_error_body;
-use super::prepare::prepare_provider_request;
-use super::types::{AnthropicMessagesResponse, MessagesRequest};
-use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
-use crate::http_utils::http_request;
+use super::{
+    Error,
+    client::http_client,
+    common_utils::truncate_error_body,
+    prepare::prepare_provider_request,
+    types::{AnthropicMessagesResponse, MessagesRequest},
+};
+use crate::{constants::ANTHROPIC_MESSAGES_PROVIDER, http_utils::http_request};
 
 pub(super) async fn execute_messages_provider_call(
     request: MessagesRequest<'_>,

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

@@ -13,9 +13,8 @@ mod client;
 mod common_utils;
 mod handler;
 mod prepare;
-pub use litellm_providers::messages::types;
-
 use handler::{execute_messages_provider_call, execute_messages_provider_stream};
+pub use litellm_providers::messages::types;
 use types::{AnthropicMessagesResponse, MessagesRequest};
 
 pub async fn messages(request: MessagesRequest<'_>) -> Result<AnthropicMessagesResponse, Error> {

litellm-rust/crates/core/src/messages/prepare.rs

@@ -1,14 +1,16 @@
-use serde_json::{Map, Value};
-
-use super::Error;
-use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
-use super::types::{MessagesRequest, ProviderMessagesRequest};
-use crate::litellm_core_utils::get_llm_provider_logic::{
-    CustomLlmProvider, get_custom_llm_provider,
-};
 use litellm_providers::base_llm::anthropic_messages::transformation::{
     BaseAnthropicMessagesConfig, MessagesAuthStrategy,
 };
+use serde_json::{Map, Value};
+
+use super::{
+    Error,
+    common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers},
+    types::{MessagesRequest, ProviderMessagesRequest},
+};
+use crate::litellm_core_utils::get_llm_provider_logic::{
+    CustomLlmProvider, get_custom_llm_provider,
+};
 
 pub(super) fn prepare_provider_request(
     request: MessagesRequest<'_>,

litellm-rust/crates/core/src/messages/tests.rs

@@ -1,15 +1,19 @@
 use std::time::Duration;
 
 use serde_json::{Map, Value, json};
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
-use tokio::net::{TcpListener, TcpStream};
-
-use super::Error;
-use super::common_utils::{
-    has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body,
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::{TcpListener, TcpStream},
+};
+
+use super::{
+    Error,
+    common_utils::{
+        has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body,
+    },
+    messages,
+    types::MessagesRequest,
 };
-use super::messages;
-use super::types::MessagesRequest;
 
 async fn read_http_request(socket: &mut TcpStream) -> String {
     let mut request = Vec::new();

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

@@ -47,6 +47,17 @@ pub fn consumed_optional_param_names(
         .collect())
 }
 
+pub(crate) fn is_secret_param(name: &str) -> bool {
+    matches!(
+        name,
+        "azure_ad_token"
+            | "client_secret"
+            | "azure_federated_token_file"
+            | "vertex_credentials"
+            | "vertex_ai_credentials"
+    )
+}
+
 pub fn consumed_optional_params(
     model: &str,
     custom_llm_provider: Option<&str>,
@@ -56,14 +67,7 @@ pub fn consumed_optional_params(
             .into_iter()
             .map(|name| ArgumentSpec {
                 name,
-                secret: matches!(
-                    name,
-                    "azure_ad_token"
-                        | "client_secret"
-                        | "azure_federated_token_file"
-                        | "vertex_credentials"
-                        | "vertex_ai_credentials"
-                ),
+                secret: is_secret_param(name),
             })
             .collect()
     })

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

@@ -1,14 +1,14 @@
-use std::sync::OnceLock;
-use std::time::Duration;
+use std::{sync::OnceLock, time::Duration};
 
 use bytes::{Bytes, BytesMut};
 use litellm_auth_gcp::VertexAuth;
 use serde::de::DeserializeOwned;
 
-use super::json::{DecodedOcrResponse, decode_response};
-use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
-use crate::constants::OCR_CONNECT_TIMEOUT_SECS;
-use crate::media::MediaFetcher;
+use super::{
+    json::{DecodedOcrResponse, decode_response},
+    types::{LiteLLMOcrRequest, LiteLLMOcrResponse},
+};
+use crate::{constants::OCR_CONNECT_TIMEOUT_SECS, media::MediaFetcher};
 
 #[derive(Clone)]
 pub struct OcrClient {
@@ -37,36 +37,11 @@ impl OcrClient {
         &self,
         request: LiteLLMOcrRequest,
     ) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
-        use super::{
-            NativeOutcome, OcrAdmission, OcrCall, OcrCallStep, OcrHookHost, OcrHost,
-            OcrHostOperation, OcrHostResult,
-        };
-
-        let host = OcrHookHost::new(request.hooks.clone());
-        let mut request = Some(request);
-        let NativeOutcome::Completed(mut call) = OcrCall::admit(self.clone(), OcrAdmission::all())
-        else {
-            return Err(crate::ocr::Error::InvalidRequest(
-                "native OCR host admission declined".into(),
-            ));
-        };
-        let mut result = None;
-        loop {
-            match call.resume(result.take()).await? {
-                OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
-                    result = Some(OcrHostResult::Request(Ok((
-                        Box::new(request.take().ok_or_else(|| {
-                            crate::ocr::Error::InvalidRequest(
-                                "OCR request was already projected".into(),
-                            )
-                        })?),
-                        false,
-                    ))))
-                }
-                OcrCallStep::Host(operation) => result = Some(host.invoke(operation).await),
-                OcrCallStep::Complete(response) => return Ok(response),
-            }
-        }
+        litellm_callbacks::run::run(
+            super::ocr_machine(self.clone()),
+            &super::LocalOcrHost::new(request),
+        )
+        .await
     }
 
     pub(crate) fn provider_http(&self) -> &reqwest::Client {

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

@@ -1,20 +1,18 @@
-use std::collections::BTreeMap as Map;
-use std::io::Read;
-use std::path::Path;
+use std::{collections::BTreeMap as Map, io::Read, path::Path};
 
 use base64::{Engine, engine::general_purpose::STANDARD};
-use data_url::mime::Mime;
-use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError};
+use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError, mime::Mime};
 use reqwest::Url;
 
-use super::Error as OcrError;
-use super::Error as OcrRequestError;
-use super::Error as OcrResponseError;
-use super::types::{OcrConnection, OcrDocument, OcrDocumentInput};
-use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS};
-use crate::media::Error as MediaError;
-use crate::media::{DownloadPolicy, MediaFetcher};
-use crate::transport::Error as TransportError;
+use super::{
+    Error as OcrError, Error as OcrRequestError, Error as OcrResponseError,
+    types::{OcrConnection, OcrDocument, OcrDocumentInput},
+};
+use crate::{
+    constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS},
+    media::{DownloadPolicy, Error as MediaError, MediaFetcher},
+    transport::Error as TransportError,
+};
 
 pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, super::Error> {
     match input {
@@ -396,8 +394,10 @@ mod tests {
 
     #[tokio::test]
     async fn remote_conversion_preserves_kind_and_isolates_provider_credentials() {
-        use tokio::io::{AsyncReadExt, AsyncWriteExt};
-        use tokio::net::TcpListener;
+        use tokio::{
+            io::{AsyncReadExt, AsyncWriteExt},
+            net::TcpListener,
+        };
 
         let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
         let address = listener.local_addr().unwrap();

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

@@ -1,36 +1,22 @@
-use std::sync::Arc;
+use litellm_callbacks::event::{CallEvent, RawResponse};
 
-use super::OcrClient;
-use super::hooks::{OcrHooks, OcrLifecycleHooks, OcrPostCallRequest};
-use super::types::{LiteLLMOcrResponse, PreparedOcrRequest, ResolvedOcrRequest};
-use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext};
+use super::{
+    OcrClient,
+    route::OcrHost,
+    types::{LiteLLMOcrResponse, PreparedOcrRequest, ResolvedOcrRequest},
+};
 use crate::llms::base_llm::ocr::transformation::OcrResponseContext;
 
 pub(crate) async fn perform_ocr_request(
     client: &OcrClient,
     request: ResolvedOcrRequest,
+    host: &OcrHost,
+    caller_document: bool,
 ) -> Result<LiteLLMOcrResponse, super::Error> {
     request.response_format()?;
-    let context = CallLifecycleContext::new(
-        "ocr",
-        request.model.clone(),
-        request.provider_name(),
-        request
-            .litellm_call_id
-            .clone()
-            .unwrap_or_else(|| format!("ocr-{:032x}", rand::random::<u128>())),
-    );
-    let hooks = OcrLifecycleHooks {
-        hooks: request.hooks.clone(),
-        provider_name: context.custom_llm_provider.clone(),
-    };
-    CallLifecycle::default()
-        .run(context, request, &hooks, |request| async move {
-            PreparedOcrCall::prepare(client.clone(), request)
-                .await?
-                .execute()
-                .await
-        })
+    PreparedOcrCall::prepare(client.clone(), request, host, caller_document)
+        .await?
+        .execute()
         .await
 }
 
@@ -44,8 +30,10 @@ impl PreparedOcrCall {
     pub(crate) async fn prepare(
         client: OcrClient,
         request: ResolvedOcrRequest,
+        host: &OcrHost,
+        caller_document: bool,
     ) -> Result<Self, super::Error> {
-        let request = super::prepare::prepare_request(request);
+        let request = super::prepare::prepare_request(request, host.clone(), caller_document);
         let http = request.config.prepare_request(&request, &client).await?;
         Ok(Self {
             client,
@@ -89,7 +77,7 @@ impl PreparedOcrCall {
         let context = OcrResponseContext {
             client: &self.client,
             connection: &self.request.connection,
-            hooks: &self.request.hooks,
+            host: &self.request.host,
             request_format: self.request.response_format()?,
             url: &url,
             headers: &headers,
@@ -116,10 +104,14 @@ fn request_headers(request: &reqwest::Request) -> Result<Vec<(String, String)>,
         .collect()
 }
 
-pub(crate) async fn post_call(hooks: &Arc<dyn OcrHooks>, bytes: &[u8]) -> Result<(), super::Error> {
-    let original_response = serde_json::Value::String(String::from_utf8_lossy(bytes).into_owned());
-    hooks
-        .post_call(OcrPostCallRequest { original_response })
-        .await?;
-    Ok(())
+pub(crate) async fn emit_response_received(
+    host: &OcrHost,
+    bytes: &[u8],
+) -> Result<(), super::Error> {
+    host.emit(CallEvent::ResponseReceived {
+        raw: RawResponse {
+            body: String::from_utf8_lossy(bytes).into_owned(),
+        },
+    })
+    .await
 }

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

@@ -1,147 +0,0 @@
-use std::future::Future;
-use std::pin::Pin;
-use std::sync::Arc;
-
-use serde::Serialize;
-use serde_json::Value;
-
-use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument, ResolvedOcrRequest};
-use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
-use crate::ocr::Error;
-
-pub type OcrHookFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
-pub type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
-
-#[derive(Clone, Debug, Serialize)]
-pub struct OcrPreCallRequest {
-    pub model: String,
-    pub custom_llm_provider: String,
-    pub document: OcrDocument,
-    pub optional_params: Value,
-}
-
-#[derive(Clone, Debug, Serialize)]
-pub struct OcrDuringCallRequest {
-    pub model: String,
-    pub custom_llm_provider: String,
-    pub api_key: Option<String>,
-    pub url: String,
-    pub headers: Vec<(String, String)>,
-    pub body: Value,
-    #[serde(skip)]
-    pub retained_fields: Vec<String>,
-}
-
-#[derive(Clone, Debug, Serialize)]
-pub struct OcrPostCallRequest {
-    pub original_response: Value,
-}
-
-pub trait OcrHooks: Send + Sync {
-    fn intercepts_requests(&self) -> bool {
-        false
-    }
-    fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
-        Box::pin(async move { Ok(request) })
-    }
-    fn during_call(
-        &self,
-        request: OcrDuringCallRequest,
-    ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
-        Box::pin(async move { Ok(request) })
-    }
-    fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> {
-        Box::pin(async move { Ok(request) })
-    }
-    fn success<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        _response: &'a LiteLLMOcrResponse,
-        _timing: &'a CallLifecycleTiming,
-    ) -> OcrLogFuture<'a> {
-        Box::pin(async {})
-    }
-    fn failure<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        _error: &'a Error,
-        _timing: &'a CallLifecycleTiming,
-    ) -> OcrLogFuture<'a> {
-        Box::pin(async {})
-    }
-}
-
-pub struct NoopOcrHooks;
-impl OcrHooks for NoopOcrHooks {}
-
-pub(crate) struct OcrLifecycleHooks {
-    pub hooks: Arc<dyn OcrHooks>,
-    pub provider_name: String,
-}
-
-impl CallLifecycleHooks<ResolvedOcrRequest, ResolvedOcrRequest, LiteLLMOcrResponse>
-    for OcrLifecycleHooks
-{
-    type Error = crate::ocr::Error;
-    type PreCallFuture<'a> = OcrHookFuture<'a, ResolvedOcrRequest>;
-    type DuringCallFuture<'a> = OcrHookFuture<'a, ResolvedOcrRequest>;
-    type SuccessFuture<'a> = OcrLogFuture<'a>;
-    type FailureFuture<'a> = OcrLogFuture<'a>;
-
-    fn async_pre_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: ResolvedOcrRequest,
-    ) -> Self::PreCallFuture<'a> {
-        Box::pin(async move {
-            if !self.hooks.intercepts_requests() {
-                return Ok(request);
-            }
-            let changed = self
-                .hooks
-                .pre_call(OcrPreCallRequest {
-                    model: request.model.clone(),
-                    custom_llm_provider: self.provider_name.clone(),
-                    document: request.document,
-                    optional_params: Value::Object(request.optional_params.into()),
-                })
-                .await?;
-            let Value::Object(optional_params) = changed.optional_params else {
-                return Err(super::Error::RequestField {
-                    path: "guardrail.optional_params".into(),
-                });
-            };
-            Ok(LiteLLMOcrRequest {
-                document: changed.document,
-                optional_params: optional_params.into(),
-                ..request
-            })
-        })
-    }
-
-    fn async_during_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: ResolvedOcrRequest,
-    ) -> Self::DuringCallFuture<'a> {
-        Box::pin(async move { Ok(request) })
-    }
-
-    fn async_log_success_event<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        response: &'a LiteLLMOcrResponse,
-        timing: &'a CallLifecycleTiming,
-    ) -> Self::SuccessFuture<'a> {
-        self.hooks.success(context, response, timing)
-    }
-
-    fn async_log_failure_event<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        error: &'a Error,
-        timing: &'a CallLifecycleTiming,
-    ) -> Self::FailureFuture<'a> {
-        self.hooks.failure(context, error, timing)
-    }
-}

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

@@ -1,727 +0,0 @@
-use std::future::Future;
-use std::pin::Pin;
-use std::sync::Arc;
-use std::sync::atomic::{AtomicBool, Ordering};
-
-use litellm_auth::Error as AuthError;
-use litellm_auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
-use tokio::sync::{Notify, mpsc, oneshot};
-
-use super::handler::perform_ocr_request;
-use super::hooks::{
-    OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest,
-    OcrPreCallRequest,
-};
-use super::types::{OcrDocumentInput, OcrFileContent};
-use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
-use crate::call_lifecycle::host::{
-    HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase,
-};
-use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming};
-use crate::ocr::Error;
-
-pub type NativeResult<T> = Result<NativeOutcome<T>, Error>;
-
-#[derive(Debug, PartialEq, Eq)]
-pub enum NativeOutcome<T> {
-    Completed(T),
-    Declined(OcrDecline),
-}
-
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub enum OcrDecline {
-    ProviderWorkflow,
-    HostOperations,
-}
-
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub struct OcrAdmission {
-    pub provider_workflow: bool,
-    pub host_operations: bool,
-    pub asynchronous: bool,
-}
-
-impl OcrAdmission {
-    pub const fn all() -> Self {
-        Self {
-            provider_workflow: true,
-            host_operations: true,
-            asynchronous: false,
-        }
-    }
-}
-
-#[derive(Clone, Debug)]
-pub enum OcrHostOperation {
-    ProjectRequest,
-    ReadDocument,
-    Lifecycle(HostPhase),
-    ConstructResponse(Arc<LiteLLMOcrResponse>),
-    MapFailure(Error),
-    Success {
-        context: CallLifecycleContext,
-        response: Arc<LiteLLMOcrResponse>,
-        timing: CallLifecycleTiming,
-    },
-    Failure {
-        context: CallLifecycleContext,
-        error: Error,
-        timing: CallLifecycleTiming,
-    },
-    AcquireAzureAdToken,
-    PreCall(OcrPreCallRequest),
-    DuringCall(OcrDuringCallRequest),
-    PostCall(OcrPostCallRequest),
-}
-
-impl OcrHostOperation {
-    pub const fn phase(&self) -> Option<HostPhase> {
-        match self {
-            Self::Lifecycle(phase) => Some(*phase),
-            Self::Success { .. } => Some(HostPhase::Success),
-            Self::Failure { .. } => Some(HostPhase::Failure),
-            _ => None,
-        }
-    }
-}
-
-pub enum OcrHostResult {
-    Request(Result<(Box<LiteLLMOcrRequest<OcrDocumentInput>>, bool), Error>),
-    Document(Result<OcrFileContent, Error>),
-    Lifecycle(Result<(), HostFailure<Error>>),
-    AzureAdToken(Result<ResolvedCredential, AuthError>),
-    PreCall(Result<OcrPreCallRequest, Error>),
-    DuringCall(Result<OcrDuringCallRequest, Error>),
-    PostCall(Result<OcrPostCallRequest, Error>),
-}
-
-pub type OcrCallStep = HostCallStep<OcrHostOperation, LiteLLMOcrResponse>;
-
-pub struct OcrCall {
-    lifecycle: HostLifecycle,
-    execution: OcrExecution,
-    response: Option<Arc<LiteLLMOcrResponse>>,
-    error: Option<Error>,
-    pending: bool,
-    completed: bool,
-    projecting: bool,
-}
-
-impl OcrCall {
-    pub fn admit(client: OcrClient, admission: OcrAdmission) -> NativeOutcome<Self> {
-        if !admission.provider_workflow {
-            return NativeOutcome::Declined(OcrDecline::ProviderWorkflow);
-        }
-        if !admission.host_operations {
-            return NativeOutcome::Declined(OcrDecline::HostOperations);
-        }
-        NativeOutcome::Completed(Self {
-            lifecycle: HostLifecycle::new(admission.asynchronous),
-            execution: OcrExecution::new(client),
-            response: None,
-            error: None,
-            pending: false,
-            completed: false,
-            projecting: false,
-        })
-    }
-
-    pub async fn resume(&mut self, result: Option<OcrHostResult>) -> Result<OcrCallStep, Error> {
-        if self.completed {
-            return Err(Error::InvalidRequest(
-                "OCR call cannot be resumed after completion".into(),
-            ));
-        }
-        if self.pending != result.is_some() {
-            return Err(Error::InvalidRequest(
-                "OCR host operation result does not match pending state".into(),
-            ));
-        }
-        match &result {
-            Some(OcrHostResult::Lifecycle(Ok(())))
-                if self.lifecycle.phase() == HostPhase::Execute =>
-            {
-                return Err(Error::InvalidRequest(
-                    "OCR provider operation requires a typed result".into(),
-                ));
-            }
-            Some(result)
-                if !matches!(result, OcrHostResult::Lifecycle(_))
-                    && self.lifecycle.phase() != HostPhase::Execute =>
-            {
-                return Err(Error::InvalidRequest(
-                    "unexpected OCR provider operation result".into(),
-                ));
-            }
-            _ => {}
-        }
-        self.pending = false;
-        let provider_result = match result {
-            Some(OcrHostResult::Request(result)) if self.projecting => {
-                self.projecting = false;
-                match result {
-                    Ok((request, azure_ad_token_provider)) => {
-                        self.execution.request = Some(*request);
-                        self.execution.azure_ad_token_provider = azure_ad_token_provider;
-                    }
-                    Err(error) => self.accept(Err(HostFailure::Error(error))),
-                }
-                None
-            }
-            Some(OcrHostResult::Request(_)) => {
-                return Err(Error::InvalidRequest(
-                    "unexpected OCR request projection".into(),
-                ));
-            }
-            Some(OcrHostResult::Lifecycle(result)) => {
-                self.accept(result);
-                None
-            }
-            result => result,
-        };
-        if self.lifecycle.phase() == HostPhase::Execute {
-            if self.execution.request.is_none()
-                && self.execution.execution.is_none()
-                && !self.execution.completed
-            {
-                self.projecting = true;
-                return Ok(self.host_step(OcrHostOperation::ProjectRequest));
-            }
-            match self.execution.resume(provider_result).await {
-                Ok(OcrCallStep::Host(operation)) => return Ok(self.host_step(operation)),
-                Ok(OcrCallStep::Complete(response)) => {
-                    self.response = Some(Arc::new(response));
-                    self.accept(Ok(()));
-                }
-                Err(error) => self.accept(Err(HostFailure::Error(error))),
-            }
-        }
-        if self.error.is_some() {
-            self.execution.stop().await;
-        }
-        let operation = match self.lifecycle.phase() {
-            HostPhase::Complete => {
-                self.completed = true;
-                return match self.error.take() {
-                    Some(error) => Err(error),
-                    None => self
-                        .response
-                        .take()
-                        .map(Arc::unwrap_or_clone)
-                        .map(OcrCallStep::Complete)
-                        .ok_or_else(|| {
-                            Error::InvalidRequest("OCR completed without a response".into())
-                        }),
-                };
-            }
-            HostPhase::ConstructResponse => OcrHostOperation::ConstructResponse(
-                self.response
-                    .as_ref()
-                    .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))?
-                    .clone(),
-            ),
-            HostPhase::MapFailure => OcrHostOperation::MapFailure(
-                self.error
-                    .as_ref()
-                    .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))?
-                    .clone(),
-            ),
-            HostPhase::Success | HostPhase::Failure => {
-                let snapshot = self
-                    .execution
-                    .terminal
-                    .lock()
-                    .unwrap_or_else(|error| error.into_inner())
-                    .clone();
-                match (self.lifecycle.phase(), snapshot) {
-                    (HostPhase::Success, Some((context, timing))) => OcrHostOperation::Success {
-                        context,
-                        response: self
-                            .response
-                            .as_ref()
-                            .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))?
-                            .clone(),
-                        timing,
-                    },
-                    (HostPhase::Failure, Some((context, timing))) => OcrHostOperation::Failure {
-                        context,
-                        error: self
-                            .error
-                            .as_ref()
-                            .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))?
-                            .clone(),
-                        timing,
-                    },
-                    (phase, _) => OcrHostOperation::Lifecycle(phase),
-                }
-            }
-            phase => OcrHostOperation::Lifecycle(phase),
-        };
-        Ok(self.host_step(operation))
-    }
-
-    fn accept(&mut self, result: Result<(), HostFailure<Error>>) {
-        let cancelled = matches!(&result, Err(HostFailure::Cancelled(_)));
-        if let Some(error) = self.lifecycle.accept(result) {
-            if cancelled {
-                self.error = Some(error);
-            } else {
-                self.error.get_or_insert(error);
-            }
-            self.execution.cancel();
-        }
-    }
-
-    pub async fn interrupt(&mut self, failure: HostFailure<Error>) -> Result<OcrCallStep, Error> {
-        if self.completed {
-            return Err(Error::InvalidRequest(
-                "OCR call cannot be interrupted after completion".into(),
-            ));
-        }
-        self.pending = false;
-        self.accept(Err(failure));
-        self.resume(None).await
-    }
-
-    fn host_step(&mut self, operation: OcrHostOperation) -> OcrCallStep {
-        self.pending = true;
-        OcrCallStep::Host(operation)
-    }
-}
-
-impl HostCall for OcrCall {
-    type Error = crate::ocr::Error;
-    type Operation = OcrHostOperation;
-    type Result = OcrHostResult;
-    type Complete = LiteLLMOcrResponse;
-
-    fn resume(
-        &mut self,
-        result: Option<Self::Result>,
-    ) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error> {
-        Box::pin(OcrCall::resume(self, result))
-    }
-
-    fn interrupt(
-        &mut self,
-        failure: HostFailure<Self::Error>,
-    ) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error> {
-        Box::pin(OcrCall::interrupt(self, failure))
-    }
-}
-
-struct PendingOperation {
-    operation: OcrHostOperation,
-    result: oneshot::Sender<OcrHostResult>,
-}
-
-struct OcrExecution {
-    client: Option<OcrClient>,
-    request: Option<LiteLLMOcrRequest<OcrDocumentInput>>,
-    operations_tx: mpsc::UnboundedSender<PendingOperation>,
-    operations_rx: mpsc::UnboundedReceiver<PendingOperation>,
-    pending_result: Option<oneshot::Sender<OcrHostResult>>,
-    execution: Option<tokio::task::JoinHandle<Result<LiteLLMOcrResponse, Error>>>,
-    blocking_preparation: Arc<BlockingPreparation>,
-    completed: bool,
-    azure_ad_token_provider: bool,
-    terminal: Arc<std::sync::Mutex<Option<(CallLifecycleContext, CallLifecycleTiming)>>>,
-}
-
-impl OcrExecution {
-    fn new(client: OcrClient) -> Self {
-        let (operations_tx, operations_rx) = mpsc::unbounded_channel();
-        Self {
-            client: Some(client),
-            request: None,
-            operations_tx,
-            operations_rx,
-            pending_result: None,
-            execution: None,
-            blocking_preparation: Arc::new(BlockingPreparation::default()),
-            completed: false,
-            azure_ad_token_provider: false,
-            terminal: Arc::default(),
-        }
-    }
-
-    pub async fn resume(&mut self, result: Option<OcrHostResult>) -> Result<OcrCallStep, Error> {
-        if self.completed {
-            return Err(Error::InvalidRequest(
-                "OCR call cannot be resumed after completion".into(),
-            ));
-        }
-        match (self.pending_result.take(), result) {
-            (Some(sender), Some(result)) => sender
-                .send(result)
-                .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into()))?,
-            (None, None) if self.execution.is_none() => self.start(),
-            (Some(sender), None) => {
-                self.pending_result = Some(sender);
-                return Err(Error::InvalidRequest(
-                    "OCR host operation result is required".into(),
-                ));
-            }
-            (None, Some(_)) => {
-                return Err(Error::InvalidRequest(
-                    "unexpected OCR host operation result".into(),
-                ));
-            }
-            (None, None) => {}
-        }
-
-        let execution = self.execution.as_mut().ok_or_else(|| {
-            Error::InvalidRequest("OCR call cannot be resumed after completion".into())
-        })?;
-        tokio::select! {
-            operation = self.operations_rx.recv() => {
-                let operation = operation.ok_or_else(|| Error::InvalidRequest("OCR operation channel closed".into()))?;
-                self.pending_result = Some(operation.result);
-                Ok(OcrCallStep::Host(operation.operation))
-            }
-            result = execution => {
-                self.execution = None;
-                self.completed = true;
-                result
-                    .map_err(|error| Error::Transport(crate::transport::Error::Network(format!("OCR execution task failed: {error}"))))?
-                    .map(OcrCallStep::Complete)
-            }
-        }
-    }
-
-    fn start(&mut self) {
-        let client = self.client.take().expect("admitted OCR call has a client");
-        let mut request = self
-            .request
-            .take()
-            .expect("admitted OCR call has a request");
-        let intercepts_requests = request.hooks.intercepts_requests();
-        if self.azure_ad_token_provider {
-            request.azure_ad_token_provider = Some(TokenProviderHandle::new(Arc::new(
-                OcrAzureAdTokenProvider {
-                    operations: self.operations_tx.clone(),
-                },
-            )));
-        }
-        let hooks = Arc::new(ProtocolHooks {
-            operations: self.operations_tx.clone(),
-            intercepts_requests,
-            terminal: self.terminal.clone(),
-        });
-        request.hooks = hooks.clone();
-        let blocking_preparation = self.blocking_preparation.clone();
-        self.execution = Some(tokio::spawn(async move {
-            let request = prepare_request_document(request, &hooks, blocking_preparation).await?;
-            perform_ocr_request(&client, request).await
-        }));
-    }
-
-    fn cancel(&mut self) {
-        self.pending_result = None;
-        if let Some(execution) = &self.execution {
-            execution.abort();
-        }
-    }
-
-    async fn stop(&mut self) {
-        self.cancel();
-        if let Some(execution) = self.execution.as_mut() {
-            let _ = execution.await;
-        }
-        self.blocking_preparation.wait().await;
-        self.execution = None;
-    }
-}
-
-#[derive(Default)]
-struct BlockingPreparation {
-    running: AtomicBool,
-    finished: Notify,
-}
-
-impl BlockingPreparation {
-    fn start(self: &Arc<Self>) -> BlockingPreparationGuard {
-        self.running.store(true, Ordering::Release);
-        BlockingPreparationGuard(self.clone())
-    }
-
-    async fn wait(&self) {
-        loop {
-            let finished = self.finished.notified();
-            if !self.running.load(Ordering::Acquire) {
-                return;
-            }
-            finished.await;
-        }
-    }
-}
-
-struct BlockingPreparationGuard(Arc<BlockingPreparation>);
-
-impl Drop for BlockingPreparationGuard {
-    fn drop(&mut self) {
-        self.0.running.store(false, Ordering::Release);
-        self.0.finished.notify_waiters();
-    }
-}
-
-async fn prepare_request_document(
-    request: LiteLLMOcrRequest<OcrDocumentInput>,
-    hooks: &ProtocolHooks,
-    blocking_preparation: Arc<BlockingPreparation>,
-) -> Result<super::types::ResolvedOcrRequest, Error> {
-    let request = match &request.document {
-        OcrDocumentInput::HostReader { mime_type } => {
-            let mime_type = mime_type.clone();
-            let content = match hooks.invoke(OcrHostOperation::ReadDocument).await? {
-                OcrHostResult::Document(result) => result?,
-                _ => {
-                    return Err(Error::InvalidRequest(
-                        "invalid OCR document read host result".into(),
-                    ));
-                }
-            };
-            request.with_document(OcrDocumentInput::Bytes {
-                bytes: content.bytes,
-                file_name: content.file_name,
-                mime_type,
-            })
-        }
-        _ => request,
-    };
-    if let OcrDocumentInput::Document(_) = &request.document {
-        return request.map_document(super::document::prepare_document);
-    }
-    let guard = blocking_preparation.start();
-    tokio::task::spawn_blocking(move || {
-        let _guard = guard;
-        request.map_document(super::document::prepare_document)
-    })
-    .await
-    .map_err(|error| {
-        Error::InvalidRequest(format!("OCR document preparation task failed: {error}"))
-    })?
-}
-
-impl Drop for OcrExecution {
-    fn drop(&mut self) {
-        if let Some(execution) = &self.execution {
-            execution.abort();
-        }
-    }
-}
-
-struct ProtocolHooks {
-    operations: mpsc::UnboundedSender<PendingOperation>,
-    intercepts_requests: bool,
-    terminal: Arc<std::sync::Mutex<Option<(CallLifecycleContext, CallLifecycleTiming)>>>,
-}
-
-#[derive(Debug)]
-struct OcrAzureAdTokenProvider {
-    operations: mpsc::UnboundedSender<PendingOperation>,
-}
-
-impl TokenProvider for OcrAzureAdTokenProvider {
-    fn acquire(&self) -> TokenFuture<'_> {
-        Box::pin(async move {
-            let (result, receiver) = oneshot::channel();
-            self.operations
-                .send(PendingOperation {
-                    operation: OcrHostOperation::AcquireAzureAdToken,
-                    result,
-                })
-                .map_err(|_| {
-                    AuthError::AzureTokenAcquisition("OCR host driver was abandoned".into())
-                })?;
-            match receiver.await.map_err(|_| {
-                AuthError::AzureTokenAcquisition(
-                    "OCR token provider operation was abandoned".into(),
-                )
-            })? {
-                OcrHostResult::AzureAdToken(result) => result,
-                _ => Err(AuthError::AzureTokenAcquisition(
-                    "invalid OCR token provider host result".into(),
-                )),
-            }
-        })
-    }
-}
-
-impl ProtocolHooks {
-    async fn invoke(&self, operation: OcrHostOperation) -> Result<OcrHostResult, Error> {
-        let (result, receiver) = oneshot::channel();
-        self.operations
-            .send(PendingOperation { operation, result })
-            .map_err(|_| Error::InvalidRequest("OCR host driver was abandoned".into()))?;
-        receiver
-            .await
-            .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into()))
-    }
-}
-
-impl OcrHooks for ProtocolHooks {
-    fn intercepts_requests(&self) -> bool {
-        self.intercepts_requests
-    }
-
-    fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
-        Box::pin(async move {
-            match self.invoke(OcrHostOperation::PreCall(request)).await? {
-                OcrHostResult::PreCall(result) => result,
-                _ => Err(Error::InvalidRequest(
-                    "invalid OCR pre-call host result".into(),
-                )),
-            }
-        })
-    }
-
-    fn during_call(
-        &self,
-        request: OcrDuringCallRequest,
-    ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
-        Box::pin(async move {
-            match self.invoke(OcrHostOperation::DuringCall(request)).await? {
-                OcrHostResult::DuringCall(result) => result,
-                _ => Err(Error::InvalidRequest(
-                    "invalid OCR during-call host result".into(),
-                )),
-            }
-        })
-    }
-
-    fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> {
-        Box::pin(async move {
-            match self.invoke(OcrHostOperation::PostCall(request)).await? {
-                OcrHostResult::PostCall(result) => result,
-                _ => Err(Error::InvalidRequest(
-                    "invalid OCR post-call host result".into(),
-                )),
-            }
-        })
-    }
-
-    fn success<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        _response: &'a LiteLLMOcrResponse,
-        timing: &'a CallLifecycleTiming,
-    ) -> OcrLogFuture<'a> {
-        Box::pin(async move {
-            *self
-                .terminal
-                .lock()
-                .unwrap_or_else(|error| error.into_inner()) =
-                Some((context.clone(), timing.clone()));
-        })
-    }
-
-    fn failure<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        _error: &'a Error,
-        timing: &'a CallLifecycleTiming,
-    ) -> OcrLogFuture<'a> {
-        Box::pin(async move {
-            *self
-                .terminal
-                .lock()
-                .unwrap_or_else(|error| error.into_inner()) =
-                Some((context.clone(), timing.clone()));
-        })
-    }
-}
-
-pub type OcrHostFuture<'a> = Pin<Box<dyn Future<Output = OcrHostResult> + Send + 'a>>;
-
-pub trait OcrHost: Send + Sync {
-    fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_>;
-}
-
-pub struct NoopOcrHost;
-
-impl OcrHost for NoopOcrHost {
-    fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> {
-        Box::pin(async move {
-            match operation {
-                OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
-                    Error::InvalidRequest("OCR host has no request projection".into()),
-                )),
-                OcrHostOperation::ReadDocument => OcrHostResult::Document(Err(
-                    Error::InvalidRequest("OCR host has no document reader".into()),
-                )),
-                OcrHostOperation::Lifecycle(_)
-                | OcrHostOperation::ConstructResponse(_)
-                | OcrHostOperation::MapFailure(_)
-                | OcrHostOperation::Success { .. }
-                | OcrHostOperation::Failure { .. } => OcrHostResult::Lifecycle(Ok(())),
-                OcrHostOperation::AcquireAzureAdToken => {
-                    OcrHostResult::AzureAdToken(Err(AuthError::AzureTokenAcquisition(
-                        "OCR host has no Azure AD token provider".into(),
-                    )))
-                }
-                OcrHostOperation::PreCall(request) => OcrHostResult::PreCall(Ok(request)),
-                OcrHostOperation::DuringCall(request) => OcrHostResult::DuringCall(Ok(request)),
-                OcrHostOperation::PostCall(request) => OcrHostResult::PostCall(Ok(request)),
-            }
-        })
-    }
-}
-
-pub struct OcrHookHost {
-    hooks: Arc<dyn OcrHooks>,
-}
-
-impl OcrHookHost {
-    pub fn new(hooks: Arc<dyn OcrHooks>) -> Self {
-        Self { hooks }
-    }
-}
-
-impl OcrHost for OcrHookHost {
-    fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> {
-        Box::pin(async move {
-            match operation {
-                OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
-                    Error::InvalidRequest("OCR hook host has no request projection".into()),
-                )),
-                OcrHostOperation::ReadDocument => OcrHostResult::Document(Err(
-                    Error::InvalidRequest("OCR hook host has no document reader".into()),
-                )),
-                OcrHostOperation::Success {
-                    context,
-                    response,
-                    timing,
-                } => {
-                    self.hooks.success(&context, &response, &timing).await;
-                    OcrHostResult::Lifecycle(Ok(()))
-                }
-                OcrHostOperation::Failure {
-                    context,
-                    error,
-                    timing,
-                } => {
-                    self.hooks.failure(&context, &error, &timing).await;
-                    OcrHostResult::Lifecycle(Ok(()))
-                }
-                OcrHostOperation::Lifecycle(_)
-                | OcrHostOperation::ConstructResponse(_)
-                | OcrHostOperation::MapFailure(_) => OcrHostResult::Lifecycle(Ok(())),
-                OcrHostOperation::AcquireAzureAdToken => {
-                    OcrHostResult::AzureAdToken(Err(AuthError::AzureTokenAcquisition(
-                        "OCR hook host has no Azure AD token provider".into(),
-                    )))
-                }
-                OcrHostOperation::PreCall(request) => {
-                    OcrHostResult::PreCall(self.hooks.pre_call(request).await)
-                }
-                OcrHostOperation::DuringCall(request) => {
-                    OcrHostResult::DuringCall(self.hooks.during_call(request).await)
-                }
-                OcrHostOperation::PostCall(request) => {
-                    OcrHostResult::PostCall(self.hooks.post_call(request).await)
-                }
-            }
-        })
-    }
-}

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

@@ -4,11 +4,10 @@ pub(crate) mod document;
 pub mod error;
 pub use error::Error;
 pub(crate) mod handler;
-pub mod hooks;
 pub(crate) mod json;
-mod lifecycle;
 pub(crate) mod prepare;
 mod provider_config;
+pub mod route;
 pub mod types;
 pub mod wire;
 
@@ -17,11 +16,8 @@ pub use arguments::{
 };
 pub use client::{OcrClient, ocr};
 pub use document::{encode_file_document, mime_type_for_name, read_path_document};
-pub use lifecycle::{
-    NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline,
-    OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult,
-};
 pub use provider_config::{get_api_key_env_var, get_health_check_document};
+pub use route::{LocalOcrHost, Ocr, OcrHost, OcrMachine, OcrOp, OcrOpResult, ocr_machine};
 pub use types::{
     LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrConnectionInputs, OcrCredentialInputs,
     OcrDocument, OcrDocumentInput, OcrFileContent, OcrPage, OcrPageDimensions, OcrPageImage,
@@ -38,6 +34,9 @@ mod azure_document_intelligence_tests;
 #[path = "../../tests/deepseek_ocr.rs"]
 mod deepseek_tests;
 #[cfg(test)]
+#[path = "../../tests/ocr/passthrough.rs"]
+mod passthrough_tests;
+#[cfg(test)]
 #[path = "../../tests/reducto_ocr.rs"]
 mod reducto_tests;
 #[cfg(test)]

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

@@ -1,8 +1,9 @@
+use litellm_callbacks::event::{Passthrough, RequestContext, WireRequest};
 use serde::Serialize;
-use serde_json::Value;
+use serde_json::{Map, Value};
 
 use super::OcrClient;
-use super::hooks::OcrDuringCallRequest;
+use super::route::OcrHost;
 use super::types::{OcrConnection, OcrDocument, PreparedOcrRequest, ResolvedOcrRequest};
 
 pub(crate) async fn transform_request_body<B>(
@@ -22,56 +23,62 @@ where
         request.config.get_supported_ocr_params(&request.model),
     )?;
     validate(&composed)?;
-    let retained_fields = request
-        .optional_params
-        .keys()
-        .filter(|name| composed.get(*name).is_some())
-        .cloned()
-        .chain(
-            composed
-                .get("document")
-                .is_some()
-                .then(|| "document".to_string()),
+    let passthrough_fields = Passthrough::unchanged(&caller_inputs(request)?, &composed);
+    let changed = request
+        .host
+        .before_send(
+            wire_request(url, headers, composed),
+            request_context(request, passthrough_fields),
         )
-        .collect();
-    let original_document =
-        serde_json::to_value(&request.document).map_err(|_| super::Error::RequestField {
+        .await?;
+    if !changed.body.is_object() {
+        return Err(super::Error::RequestField {
+            path: "guardrail.body".into(),
+        });
+    }
+    validate(&changed.body)?;
+    build_http_request(client, request, url, &changed.headers, &changed.body)
+}
+
+fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireRequest {
+    WireRequest {
+        url: url.into(),
+        headers: headers.to_vec(),
+        body,
+    }
+}
+
+fn caller_inputs(request: &PreparedOcrRequest) -> Result<Map<String, Value>, super::Error> {
+    let document = request
+        .caller_document
+        .then(|| serde_json::to_value(&request.document))
+        .transpose()
+        .map_err(|_| super::Error::RequestField {
             path: "document".into(),
         })?;
-    let prepared_document = composed
-        .get("document")
-        .filter(|prepared| **prepared != original_document)
-        .cloned();
-    let (body, headers) = if request.hooks.intercepts_requests() {
-        let changed = request
-            .hooks
-            .during_call(OcrDuringCallRequest {
-                model: request.model.clone(),
-                custom_llm_provider: request.provider_name().into(),
-                api_key: request.connection.api_key.clone(),
-                url: url.into(),
-                headers: headers.to_vec(),
-                body: composed,
-                retained_fields,
-            })
-            .await?;
-        let Value::Object(mut fields) = changed.body else {
-            return Err(super::Error::RequestField {
-                path: "guardrail.body".into(),
-            });
-        };
-        if let Some(prepared) =
-            prepared_document.filter(|_| fields.get("document") == Some(&original_document))
-        {
-            fields.insert("document".into(), prepared);
-        }
-        let body = Value::Object(fields);
-        validate(&body)?;
-        (body, changed.headers)
-    } else {
-        (composed, headers.to_vec())
-    };
-    build_http_request(client, request, url, &headers, &body)
+    let params: Map<String, Value> = request.optional_params.clone().into();
+    Ok(params
+        .into_iter()
+        .chain(document.map(|document| ("document".to_string(), document)))
+        .collect())
+}
+
+fn request_context(
+    request: &PreparedOcrRequest,
+    passthrough_fields: Passthrough,
+) -> RequestContext {
+    RequestContext {
+        model: request.model.clone(),
+        custom_llm_provider: request.provider_name().into(),
+        optional_params: Value::Object(request.optional_params.clone().into()),
+        passthrough_fields,
+        secret_fields: request
+            .optional_params
+            .keys()
+            .filter(|name| super::arguments::is_secret_param(name))
+            .cloned()
+            .collect(),
+    }
 }
 
 pub(crate) fn build_http_request<B: Serialize>(
@@ -97,24 +104,15 @@ pub(crate) async fn guardrail_document(
     url: &str,
     headers: &[(String, String)],
 ) -> Result<(OcrDocument, Vec<(String, String)>), super::Error> {
-    if !request.hooks.intercepts_requests() {
-        return Ok((request.document.clone(), headers.to_vec()));
-    }
+    let body = serde_json::to_value(&request.document).map_err(|_| super::Error::RequestField {
+        path: "document".into(),
+    })?;
     let changed = request
-        .hooks
-        .during_call(OcrDuringCallRequest {
-            model: request.model.clone(),
-            custom_llm_provider: request.provider_name().into(),
-            api_key: request.connection.api_key.clone(),
-            url: url.into(),
-            headers: headers.to_vec(),
-            body: serde_json::to_value(&request.document).map_err(|_| {
-                super::Error::RequestField {
-                    path: "document".into(),
-                }
-            })?,
-            retained_fields: Vec::new(),
-        })
+        .host
+        .before_send(
+            wire_request(url, headers, body),
+            request_context(request, Passthrough::default()),
+        )
         .await?;
     let document = super::json::decode_request_value(changed.body, "guardrail.document")?;
     Ok((document, changed.headers))
@@ -139,7 +137,11 @@ pub(crate) fn credential_env(name: &str) -> Option<String> {
     std::env::var(name).ok()
 }
 
-pub(crate) fn prepare_request(request: ResolvedOcrRequest) -> PreparedOcrRequest {
+pub(crate) fn prepare_request(
+    request: ResolvedOcrRequest,
+    host: OcrHost,
+    caller_document: bool,
+) -> PreparedOcrRequest {
     use litellm_auth::{InputSource, Sourced};
 
     let credentials = request.credentials.clone();
@@ -174,7 +176,17 @@ pub(crate) fn prepare_request(request: ResolvedOcrRequest) -> PreparedOcrRequest
             ..credentials
         });
     let transport = request.transport.clone();
-    PreparedOcrRequest::new(request, OcrConnection::new(resolved, transport))
+    PreparedOcrRequest::new(
+        request,
+        OcrConnection::new(resolved, transport),
+        host,
+        caller_document,
+    )
+}
+
+#[cfg(test)]
+pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest {
+    prepare_request(request, OcrHost::detached(), true)
 }
 
 #[cfg(test)]

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

@@ -1,22 +1,29 @@
 use strum::{EnumString, IntoStaticStr};
 
-use super::OcrClient;
-use super::types::{
-    LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument, PreparedOcrRequest,
-    ResolvedOcrCredentials,
+use super::{
+    OcrClient,
+    types::{
+        LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument, PreparedOcrRequest,
+        ResolvedOcrCredentials,
+    },
 };
-use crate::litellm_core_utils::get_llm_provider_logic::{
-    CustomLlmProvider, get_custom_llm_provider,
+use crate::{
+    litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider},
+    llms::{
+        azure_ai::ocr::{
+            cohere_parse_transformation::AzureAICohereParseConfig,
+            document_intelligence::transformation::AzureDocumentIntelligenceOcrConfig,
+            transformation::AzureAiOcrConfig,
+        },
+        base_llm::ocr::transformation::{BaseOcrConfig, OcrResponseContext},
+        cohere::ocr::transformation::CohereParseConfig,
+        mistral::ocr::transformation::MistralOcrConfig,
+        reducto::ocr::transformation::{ReductoParseLegacyConfig, ReductoParseV3Config},
+        vertex_ai::ocr::{
+            deepseek_transformation::VertexAIDeepSeekOCRConfig, transformation::VertexAiOcrConfig,
+        },
+    },
 };
-use crate::llms::azure_ai::ocr::cohere_parse_transformation::AzureAICohereParseConfig;
-use crate::llms::azure_ai::ocr::document_intelligence::transformation::AzureDocumentIntelligenceOcrConfig;
-use crate::llms::azure_ai::ocr::transformation::AzureAiOcrConfig;
-use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrResponseContext};
-use crate::llms::cohere::ocr::transformation::CohereParseConfig;
-use crate::llms::mistral::ocr::transformation::MistralOcrConfig;
-use crate::llms::reducto::ocr::transformation::{ReductoParseLegacyConfig, ReductoParseV3Config};
-use crate::llms::vertex_ai::ocr::deepseek_transformation::VertexAIDeepSeekOCRConfig;
-use crate::llms::vertex_ai::ocr::transformation::VertexAiOcrConfig;
 
 macro_rules! dispatch_config {
     ($config:expr, $method:ident($($argument:expr),* $(,)?)) => {

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

@@ -0,0 +1,217 @@
+use std::sync::{Arc, Mutex};
+
+use litellm_auth::ResolvedCredential;
+use litellm_callbacks::{
+    event::{CallEvent, RequestContext, WireRequest},
+    route::Route,
+};
+
+use super::{
+    Error, LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient,
+    handler::perform_ocr_request,
+    types::{OcrDocumentInput, OcrFileContent, ResolvedOcrRequest},
+};
+use crate::machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute};
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum OcrOp {
+    ProjectRequest,
+    ReadDocument,
+    AcquireAzureAdToken,
+}
+
+pub enum OcrOpResult {
+    Request {
+        request: Box<LiteLLMOcrRequest<OcrDocumentInput>>,
+        caller_token: bool,
+    },
+    Document(OcrFileContent),
+    AzureAdToken(ResolvedCredential),
+}
+
+pub struct Ocr;
+
+impl Route for Ocr {
+    type Response = LiteLLMOcrResponse;
+    type Error = Error;
+    type Op = OcrOp;
+    type OpResult = OcrOpResult;
+}
+
+impl TokenRoute for Ocr {
+    fn acquire_token_op() -> OcrOp {
+        OcrOp::AcquireAzureAdToken
+    }
+
+    fn token_credential(result: OcrOpResult) -> Option<ResolvedCredential> {
+        match result {
+            OcrOpResult::AzureAdToken(credential) => Some(credential),
+            _ => None,
+        }
+    }
+}
+
+impl From<MachineFault> for Error {
+    fn from(fault: MachineFault) -> Self {
+        Self::InvalidRequest(match fault {
+            MachineFault::Abandoned => "OCR host driver was abandoned".into(),
+            MachineFault::Protocol(message) => format!("OCR {message}"),
+            MachineFault::Mismatch => "invalid OCR host operation result".into(),
+        })
+    }
+}
+
+pub type OcrHost = HostChannel<Ocr>;
+pub type OcrMachine = RouteMachine<Ocr>;
+
+/// The OCR call as a machine: projection, document reading and token acquisition are
+/// host operations; everything else runs in Rust.
+pub fn ocr_machine(client: OcrClient) -> OcrMachine {
+    RouteMachine::new(move |host| Box::pin(execute(client, host)))
+}
+
+async fn execute(client: OcrClient, host: OcrHost) -> Result<LiteLLMOcrResponse, Error> {
+    let OcrOpResult::Request {
+        request,
+        caller_token,
+    } = host.route(OcrOp::ProjectRequest).await?
+    else {
+        return Err(MachineFault::Mismatch.into());
+    };
+    let request = LiteLLMOcrRequest {
+        azure_ad_token_provider: caller_token
+            .then(|| HostTokenProvider::handle(host.clone()))
+            .or(request.azure_ad_token_provider),
+        ..*request
+    };
+    let caller_document = matches!(request.document, OcrDocumentInput::Document(_));
+    let request = prepare_request_document(request, &host).await?;
+    perform_ocr_request(&client, request, &host, caller_document).await
+}
+
+async fn prepare_request_document(
+    request: LiteLLMOcrRequest<OcrDocumentInput>,
+    host: &OcrHost,
+) -> Result<ResolvedOcrRequest, Error> {
+    let request = match &request.document {
+        OcrDocumentInput::HostReader { mime_type } => {
+            let mime_type = mime_type.clone();
+            let OcrOpResult::Document(content) = host.route(OcrOp::ReadDocument).await? else {
+                return Err(MachineFault::Mismatch.into());
+            };
+            request.with_document(OcrDocumentInput::Bytes {
+                bytes: content.bytes,
+                file_name: content.file_name,
+                mime_type,
+            })
+        }
+        _ => request,
+    };
+    if let OcrDocumentInput::Document(_) = &request.document {
+        return request.map_document(super::document::prepare_document);
+    }
+    tokio::task::spawn_blocking(move || request.map_document(super::document::prepare_document))
+        .await
+        .map_err(|error| Error::DocumentTask(Arc::new(error)))?
+}
+
+type Reader = Box<dyn Fn() -> Result<OcrFileContent, Error> + Send + Sync>;
+type BeforeSend =
+    Box<dyn Fn(WireRequest, &RequestContext) -> Result<WireRequest, Error> + Send + Sync>;
+type Observer = Box<dyn Fn(&CallEvent) + Send + Sync>;
+
+/// The in-process host for a request that is already in hand: the request answers
+/// projection, and the optional observer sees and may rewrite the wire request.
+pub struct LocalOcrHost {
+    request: Mutex<Option<LiteLLMOcrRequest<OcrDocumentInput>>>,
+    reader: Option<Reader>,
+    before_send: Option<BeforeSend>,
+    observer: Option<Observer>,
+}
+
+impl LocalOcrHost {
+    pub fn new(request: LiteLLMOcrRequest<OcrDocumentInput>) -> Self {
+        Self {
+            request: Mutex::new(Some(request)),
+            reader: None,
+            before_send: None,
+            observer: None,
+        }
+    }
+
+    pub fn with_reader(
+        self,
+        reader: impl Fn() -> Result<OcrFileContent, Error> + Send + Sync + 'static,
+    ) -> Self {
+        Self {
+            reader: Some(Box::new(reader)),
+            ..self
+        }
+    }
+
+    pub fn with_before_send(
+        self,
+        before_send: impl Fn(WireRequest, &RequestContext) -> Result<WireRequest, Error>
+        + Send
+        + Sync
+        + 'static,
+    ) -> Self {
+        Self {
+            before_send: Some(Box::new(before_send)),
+            ..self
+        }
+    }
+
+    pub fn with_observer(self, observer: impl Fn(&CallEvent) + Send + Sync + 'static) -> Self {
+        Self {
+            observer: Some(Box::new(observer)),
+            ..self
+        }
+    }
+}
+
+impl litellm_callbacks::host::Host<Ocr> for LocalOcrHost {
+    async fn route(&self, op: OcrOp) -> Result<OcrOpResult, Error> {
+        match op {
+            OcrOp::ProjectRequest => self
+                .request
+                .lock()
+                .unwrap_or_else(|error| error.into_inner())
+                .take()
+                .map(|request| OcrOpResult::Request {
+                    request: Box::new(request),
+                    caller_token: false,
+                })
+                .ok_or_else(|| Error::InvalidRequest("OCR request was already projected".into())),
+            OcrOp::ReadDocument => self
+                .reader
+                .as_ref()
+                .ok_or_else(|| Error::InvalidRequest("OCR host has no document reader".into()))
+                .and_then(|reader| reader())
+                .map(OcrOpResult::Document),
+            OcrOp::AcquireAzureAdToken => {
+                Err(Error::Auth(litellm_auth::Error::AzureTokenAcquisition(
+                    "OCR host has no Azure AD token provider".into(),
+                )))
+            }
+        }
+    }
+
+    async fn before_send(
+        &self,
+        wire: WireRequest,
+        context: &RequestContext,
+    ) -> Result<WireRequest, Error> {
+        match &self.before_send {
+            Some(before_send) => before_send(wire, context),
+            None => Ok(wire),
+        }
+    }
+
+    async fn emit(&self, event: &CallEvent) -> Result<(), Error> {
+        if let Some(observer) = &self.observer {
+            observer(event);
+        }
+        Ok(())
+    }
+}

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

@@ -1,7 +1,4 @@
-use std::collections::BTreeMap;
-use std::path::PathBuf;
-use std::sync::Arc;
-use std::time::Duration;
+use std::{collections::BTreeMap, path::PathBuf, time::Duration};
 
 use bytes::Bytes;
 use litellm_auth::{InputSource, Sourced, TokenProviderHandle};
@@ -9,11 +6,12 @@ use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 use serde_with::serde_as;
 
-use super::hooks::{NoopOcrHooks, OcrHooks};
 use super::provider_config::{OcrConfigKind, resolve_provider_config};
-use crate::call_arguments::CallArguments;
-use crate::constants::OCR_HTTP_TIMEOUT_SECS;
-use crate::serde_compat::{FiniteF64, LaxI64};
+use crate::{
+    call_arguments::CallArguments,
+    constants::OCR_HTTP_TIMEOUT_SECS,
+    serde_compat::{FiniteF64, LaxI64},
+};
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 #[serde(tag = "type")]
@@ -275,8 +273,6 @@ pub struct LiteLLMOcrRequest<D = OcrDocumentInput> {
     pub document: D,
     pub credentials: OcrCredentialInputs,
     pub transport: OcrTransportConfig,
-    pub hooks: Arc<dyn OcrHooks>,
-    pub litellm_call_id: Option<String>,
     pub optional_params: CallArguments,
     pub input_sources: BTreeMap<String, InputSource>,
     pub azure_ad_token_provider: Option<TokenProviderHandle>,
@@ -319,8 +315,6 @@ impl LiteLLMOcrRequest {
             document: document.into(),
             credentials: OcrCredentialInputs::default(),
             transport,
-            hooks: Arc::new(NoopOcrHooks),
-            litellm_call_id: None,
             optional_params,
             input_sources: BTreeMap::new(),
             azure_ad_token_provider: None,
@@ -339,8 +333,6 @@ impl<D> LiteLLMOcrRequest<D> {
             document: map(self.document)?,
             credentials: self.credentials,
             transport: self.transport,
-            hooks: self.hooks,
-            litellm_call_id: self.litellm_call_id,
             optional_params: self.optional_params,
             input_sources: self.input_sources,
             azure_ad_token_provider: self.azure_ad_token_provider,
@@ -354,8 +346,6 @@ impl<D> LiteLLMOcrRequest<D> {
             document,
             credentials: self.credentials,
             transport: self.transport,
-            hooks: self.hooks,
-            litellm_call_id: self.litellm_call_id,
             optional_params: self.optional_params,
             input_sources: self.input_sources,
             azure_ad_token_provider: self.azure_ad_token_provider,
@@ -378,18 +368,6 @@ impl<D> LiteLLMOcrRequest<D> {
         self.config.provider().into()
     }
 
-    pub fn with_host_hooks(
-        self,
-        hooks: Arc<dyn OcrHooks>,
-        litellm_call_id: Option<String>,
-    ) -> Self {
-        Self {
-            hooks,
-            litellm_call_id,
-            ..self
-        }
-    }
-
     pub fn with_connection_inputs(
         self,
         credentials: OcrCredentialInputs,
@@ -442,7 +420,10 @@ pub(crate) struct PreparedOcrRequest {
     pub model: String,
     pub document: OcrDocument,
     pub connection: OcrConnection,
-    pub hooks: Arc<dyn OcrHooks>,
+    pub host: super::route::OcrHost,
+    /// Whether the caller handed over the document as is, so the wire body's document
+    /// is the caller's own input rather than something the route prepared.
+    pub caller_document: bool,
     pub optional_params: CallArguments,
     pub input_sources: BTreeMap<String, InputSource>,
     pub azure_ad_token_provider: Option<TokenProviderHandle>,
@@ -450,14 +431,17 @@ pub(crate) struct PreparedOcrRequest {
 }
 
 impl PreparedOcrRequest {
-    pub(crate) fn new(request: ResolvedOcrRequest, connection: OcrConnection) -> Self {
+    pub(crate) fn new(
+        request: ResolvedOcrRequest,
+        connection: OcrConnection,
+        host: super::route::OcrHost,
+        caller_document: bool,
+    ) -> Self {
         let LiteLLMOcrRequest {
             model,
             document,
             credentials: _,
             transport: _,
-            hooks,
-            litellm_call_id: _,
             optional_params,
             input_sources,
             azure_ad_token_provider,
@@ -467,7 +451,8 @@ impl PreparedOcrRequest {
             model,
             document,
             connection,
-            hooks,
+            host,
+            caller_document,
             optional_params,
             input_sources,
             azure_ad_token_provider,

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

@@ -1,5 +1,4 @@
-use std::collections::BTreeMap;
-use std::time::Duration;
+use std::{collections::BTreeMap, time::Duration};
 
 use litellm_auth::InputSource;
 use serde::Deserialize;
@@ -105,10 +104,11 @@ pub fn decode_document(value: Value) -> Result<OcrDocument, Error> {
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use rstest::rstest;
     use serde_json::json;
 
+    use super::*;
+
     #[rstest]
     #[case::omitted(json!({"type":"document_url", "document_url":"https://example.com/a.pdf"}))]
     #[case::null(json!({"type":"document_url", "document_url":"https://example.com/a.pdf", "document_name":null}))]
@@ -120,7 +120,7 @@ mod tests {
     #[rstest]
     #[case::non_object(json!([]), "document")]
     #[case::missing_type(json!({"document_url":"https://example.com/a.pdf"}), "document")]
-    #[case::unsupported_type(json!({"type":"text"}), "document")]
+    #[case::unsupported_type(json!({"type":"text"}), "type")]
     #[case::missing_document_url(json!({"type":"document_url"}), "Document URL")]
     #[case::missing_image_url(json!({"type":"image_url"}), "Document URL")]
     fn ocr_contract_malformed_document_is_bad_request(
@@ -133,7 +133,7 @@ mod tests {
             Error::RequestField { .. } | Error::MissingDocumentUrl
         ));
         assert_eq!(error.http_status_code(), Some(400));
-        assert!(error.to_string().contains(field));
+        assert!(error.to_string().contains(field), "{error}");
     }
 
     #[test]

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

@@ -28,14 +28,6 @@ pub fn is_control_param(name: &str) -> bool {
             | "max_retries"
             | "req_format"
             | "max_response_bytes"
-            | "litellm_call_id"
-            | "litellm_logging_obj"
-            | "litellm_metadata"
-            | "proxy_server_request"
-            | "callbacks"
-            | "success_callback"
-            | "failure_callback"
-            | "guardrails"
             | "azure_ad_token"
             | "azure_ad_token_provider"
             | "tenant_id"

litellm-rust/crates/core/src/responses/instrumentation.rs

@@ -1,366 +0,0 @@
-use std::future::Future;
-use std::pin::Pin;
-use std::sync::Mutex;
-use std::time::{SystemTime, UNIX_EPOCH};
-
-use serde_json::Value;
-
-use super::Error;
-use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
-use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType};
-
-#[derive(Clone, Debug, Default, PartialEq, Eq)]
-pub struct ResponsesWsUsage {
-    pub prompt_tokens: u64,
-    pub completion_tokens: u64,
-    pub total_tokens: u64,
-}
-
-#[derive(Clone, Debug, Default, PartialEq, Eq)]
-pub struct ResponsesWsMetadata {
-    pub user_api_key_hash: Option<String>,
-    pub user_api_key_user_id: Option<String>,
-    pub user_api_key_team_id: Option<String>,
-}
-
-#[derive(Clone, Debug, PartialEq)]
-pub struct ResponsesWsLogPayload {
-    pub id: String,
-    pub litellm_call_id: String,
-    pub call_type: String,
-    pub model: String,
-    pub custom_llm_provider: String,
-    pub response_cost: f64,
-    pub usage: ResponsesWsUsage,
-    pub start_time: f64,
-    pub end_time: f64,
-    pub stream: bool,
-    pub metadata: ResponsesWsMetadata,
-}
-
-#[derive(Clone, Debug, PartialEq)]
-pub enum ResponsesWsLogOutcome {
-    Success {
-        payload: ResponsesWsLogPayload,
-        callback: ResponsesWsCallbackPayload,
-    },
-    Failure {
-        payload: ResponsesWsLogPayload,
-        callback: ResponsesWsCallbackPayload,
-        error_message: String,
-        error_kind: String,
-    },
-}
-
-#[derive(Clone, Debug, PartialEq)]
-pub struct ResponsesWsCallbackPayload {
-    pub object: String,
-    pub value: Value,
-}
-
-struct InstrumentationState {
-    litellm_call_id: String,
-    id: String,
-    model: String,
-    usage: ResponsesWsUsage,
-    start_time: f64,
-    end_time: f64,
-    metadata: ResponsesWsMetadata,
-    outcome: Option<ResponsesWsLogOutcome>,
-}
-
-pub struct ResponsesWsInstrumentation {
-    state: Mutex<InstrumentationState>,
-}
-
-impl ResponsesWsInstrumentation {
-    pub fn new(
-        litellm_call_id: impl Into<String>,
-        model: impl Into<String>,
-        metadata: ResponsesWsMetadata,
-    ) -> Self {
-        let litellm_call_id = litellm_call_id.into();
-        let now = epoch_seconds();
-        Self {
-            state: Mutex::new(InstrumentationState {
-                id: litellm_call_id.clone(),
-                litellm_call_id,
-                model: model.into(),
-                usage: ResponsesWsUsage::default(),
-                start_time: now,
-                end_time: now,
-                metadata,
-                outcome: None,
-            }),
-        }
-    }
-
-    pub fn observe(&self, event: &ResponsesWsEvent) {
-        if !matches!(
-            event.event_type,
-            ResponsesWsEventType::ResponseCreated
-                | ResponsesWsEventType::ResponseCompleted
-                | ResponsesWsEventType::ResponseFailed
-                | ResponsesWsEventType::ResponseIncomplete
-                | ResponsesWsEventType::Error
-        ) {
-            return;
-        }
-        let Ok(mut state) = self.state.lock() else {
-            return;
-        };
-        let Some(response) = event.data.get("response").and_then(Value::as_object) else {
-            return;
-        };
-        if let Some(id) = response
-            .get("id")
-            .and_then(Value::as_str)
-            .filter(|value| !value.is_empty())
-        {
-            state.id = id.to_string();
-            state.litellm_call_id = id.to_string();
-        }
-        if let Some(model) = response
-            .get("model")
-            .and_then(Value::as_str)
-            .filter(|value| !value.is_empty())
-        {
-            state.model = model.to_string();
-        }
-        let Some(usage) = response.get("usage").and_then(Value::as_object) else {
-            return;
-        };
-        if let Some(input) = usage.get("input_tokens").and_then(Value::as_u64) {
-            state.usage.prompt_tokens += input;
-        }
-        if let Some(output) = usage.get("output_tokens").and_then(Value::as_u64) {
-            state.usage.completion_tokens += output;
-        }
-        state.usage.total_tokens += usage
-            .get("total_tokens")
-            .and_then(Value::as_u64)
-            .unwrap_or_else(|| {
-                usage
-                    .get("input_tokens")
-                    .and_then(Value::as_u64)
-                    .unwrap_or(0)
-                    + usage
-                        .get("output_tokens")
-                        .and_then(Value::as_u64)
-                        .unwrap_or(0)
-            });
-    }
-
-    pub fn success_outcome(&self) -> ResponsesWsLogOutcome {
-        let mut state = self
-            .state
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner());
-        state.end_time = epoch_seconds();
-        ResponsesWsLogOutcome::Success {
-            payload: build_payload(&state),
-            callback: ResponsesWsCallbackPayload {
-                object: "responses_websocket".to_string(),
-                value: Value::Null,
-            },
-        }
-    }
-
-    pub fn failure_outcome(&self) -> ResponsesWsLogOutcome {
-        let mut state = self
-            .state
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner());
-        state.end_time = epoch_seconds();
-        ResponsesWsLogOutcome::Failure {
-            payload: build_payload(&state),
-            callback: ResponsesWsCallbackPayload {
-                object: "error".to_string(),
-                value: serde_json::json!({
-                    "message": "Responses WebSocket session ended in failure",
-                    "kind": "ResponsesWebSocketError",
-                }),
-            },
-            error_message: "Responses WebSocket session ended in failure".to_string(),
-            error_kind: "ResponsesWebSocketError".to_string(),
-        }
-    }
-
-    pub fn take_outcome(&self) -> Option<ResponsesWsLogOutcome> {
-        self.state
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
-            .outcome
-            .take()
-    }
-
-    pub fn take_or_build_outcome(&self, success: bool) -> ResponsesWsLogOutcome {
-        self.take_outcome().unwrap_or_else(|| {
-            if success {
-                self.success_outcome()
-            } else {
-                self.failure_outcome()
-            }
-        })
-    }
-}
-
-type LifecycleFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
-
-impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation {
-    type Error = Error;
-    type PreCallFuture<'a> = LifecycleFuture<'a, ()>;
-    type DuringCallFuture<'a> = LifecycleFuture<'a, ()>;
-    type SuccessFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
-    type FailureFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
-
-    fn async_pre_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: (),
-    ) -> Self::PreCallFuture<'a> {
-        Box::pin(async move { Ok(request) })
-    }
-
-    fn async_during_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: (),
-    ) -> Self::DuringCallFuture<'a> {
-        Box::pin(async move { Ok(request) })
-    }
-
-    fn async_log_success_event<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        _response: &'a (),
-        _timing: &'a CallLifecycleTiming,
-    ) -> Self::SuccessFuture<'a> {
-        Box::pin(async move {
-            let outcome = self.success_outcome();
-            if let Ok(mut state) = self.state.lock() {
-                state.outcome = Some(outcome);
-            }
-        })
-    }
-
-    fn async_log_failure_event<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        _error: &'a Error,
-        _timing: &'a CallLifecycleTiming,
-    ) -> Self::FailureFuture<'a> {
-        Box::pin(async move {
-            let outcome = self.failure_outcome();
-            if let Ok(mut state) = self.state.lock() {
-                state.outcome = Some(outcome);
-            }
-        })
-    }
-}
-
-fn build_payload(state: &InstrumentationState) -> ResponsesWsLogPayload {
-    ResponsesWsLogPayload {
-        id: state.id.clone(),
-        litellm_call_id: state.litellm_call_id.clone(),
-        call_type: "responses_websocket".to_string(),
-        model: state.model.clone(),
-        custom_llm_provider: "openai".to_string(),
-        response_cost: 0.0,
-        usage: state.usage.clone(),
-        start_time: state.start_time,
-        end_time: state.end_time,
-        stream: true,
-        metadata: state.metadata.clone(),
-    }
-}
-
-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::*;
-
-    fn event(value: Value) -> ResponsesWsEvent {
-        serde_json::from_value(value).expect("valid Responses WebSocket event")
-    }
-
-    #[test]
-    fn accumulates_upstream_usage_and_identity() {
-        let instrumentation =
-            ResponsesWsInstrumentation::new("call-1", "gpt-5", ResponsesWsMetadata::default());
-        instrumentation.observe(&event(serde_json::json!({
-            "type": "response.completed",
-            "response": {
-                "id": "resp-1",
-                "model": "gpt-5-mini",
-                "usage": {
-                    "input_tokens": 3,
-                    "output_tokens": 5,
-                    "total_tokens": 8
-                }
-            }
-        })));
-
-        let ResponsesWsLogOutcome::Success { payload, .. } = instrumentation.success_outcome()
-        else {
-            panic!("expected success outcome");
-        };
-        assert_eq!(payload.id, "resp-1");
-        assert_eq!(payload.model, "gpt-5-mini");
-        assert_eq!(payload.usage.prompt_tokens, 3);
-        assert_eq!(payload.usage.completion_tokens, 5);
-        assert_eq!(payload.usage.total_tokens, 8);
-        assert!(payload.end_time >= payload.start_time);
-    }
-
-    #[test]
-    fn builds_failure_payload_without_dispatching_callbacks() {
-        let instrumentation =
-            ResponsesWsInstrumentation::new("call-1", "gpt-5", ResponsesWsMetadata::default());
-        assert!(matches!(
-            instrumentation.failure_outcome(),
-            ResponsesWsLogOutcome::Failure { .. }
-        ));
-    }
-
-    #[tokio::test]
-    async fn lifecycle_records_success_outcome_for_provider_completion() {
-        let instrumentation =
-            ResponsesWsInstrumentation::new("call-1", "gpt-5", ResponsesWsMetadata::default());
-        let result = crate::call_lifecycle::CallLifecycle::default()
-            .run(
-                crate::call_lifecycle::CallLifecycleContext::new(
-                    "responses_websocket",
-                    "gpt-5",
-                    "openai",
-                    "call-1",
-                ),
-                (),
-                &instrumentation,
-                |_| async { Ok::<(), Error>(()) },
-            )
-            .await;
-
-        assert!(result.is_ok());
-        assert!(matches!(
-            instrumentation.take_outcome(),
-            Some(ResponsesWsLogOutcome::Success { .. })
-        ));
-    }
-
-    #[test]
-    fn builds_outcome_when_lifecycle_did_not_record_one() {
-        let instrumentation =
-            ResponsesWsInstrumentation::new("call-1", "gpt-5", ResponsesWsMetadata::default());
-        assert!(matches!(
-            instrumentation.take_or_build_outcome(true),
-            ResponsesWsLogOutcome::Success { .. }
-        ));
-    }
-}

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

@@ -1,5 +1,4 @@
 mod error;
 pub use error::Error;
-pub mod instrumentation;
 pub mod types;
 pub mod websocket;

litellm-rust/crates/core/src/responses/websocket.rs

@@ -1,24 +1,29 @@
-use std::collections::HashMap;
-use std::io;
-use std::sync::{Arc, OnceLock};
-use std::time::Duration;
+use std::{
+    collections::HashMap,
+    io,
+    sync::{Arc, OnceLock},
+    time::Duration,
+};
 
 use futures_util::{SinkExt, StreamExt};
 use rustls::{ClientConfig, RootCertStore};
-use tokio::net::TcpStream;
-use tokio::sync::Mutex;
-use tokio_tungstenite::tungstenite::Message;
-use tokio_tungstenite::tungstenite::client::IntoClientRequest;
-use tokio_tungstenite::tungstenite::error::TlsError;
-use tokio_tungstenite::tungstenite::handshake::client::Response;
-use tokio_tungstenite::tungstenite::http::{HeaderName, HeaderValue};
+use tokio::{net::TcpStream, sync::Mutex};
 use tokio_tungstenite::{
     Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
+    tungstenite::{
+        Message,
+        client::IntoClientRequest,
+        error::TlsError,
+        handshake::client::Response,
+        http::{HeaderName, HeaderValue},
+    },
 };
 
 use super::Error;
-use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH};
-use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult};
+use crate::{
+    constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH},
+    responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult},
+};
 
 pub trait ResponsesWebSocketProviderConfig: Sync {
     fn supports_native_websocket(&self) -> bool {

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

@@ -1,9 +1,9 @@
-use std::sync::Arc;
-
 use serde_json::{Value, json};
 
-use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks};
-use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
+use super::{
+    LocalOcrHost,
+    test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
+};
 
 #[tokio::test]
 async fn facade_executes_azure_mistral_with_prepared_auth() {
@@ -67,31 +67,16 @@ async fn facade_acquires_supplied_entra_token_for_final_request() {
     );
 }
 
-struct ReplaceBodyDocument;
-
-impl OcrHooks for ReplaceBodyDocument {
-    fn intercepts_requests(&self) -> bool {
-        true
-    }
-
-    fn during_call(
-        &self,
-        mut request: OcrDuringCallRequest,
-    ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
-        Box::pin(async move {
-            request.body["document"] = json!({
-                "type":"document_url",
-                "document_url":"https://example.com/not-inline.pdf"
-            });
-            Ok(request)
-        })
-    }
-}
-
 #[tokio::test]
 async fn rejects_non_inline_body_after_guardrails() {
-    let mut request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({}));
-    request.hooks = Arc::new(ReplaceBodyDocument);
-    let error = perform_ocr(request).await.unwrap_err();
+    let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({}));
+    let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| {
+        wire.body["document"] = json!({
+            "type":"document_url",
+            "document_url":"https://example.com/not-inline.pdf"
+        });
+        Ok(wire)
+    });
+    let error = perform_ocr_with(host).await.unwrap_err();
     assert!(error.to_string().contains("data URI"));
 }

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

@@ -1,10 +1,12 @@
-use std::sync::{Arc, Mutex};
-
+use litellm_callbacks::event::CallEvent;
 use rstest::rstest;
 use serde_json::{Value, json};
 
-use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
-use super::wire::{OcrWireRequest, decode_request};
+use super::{
+    LocalOcrHost,
+    test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
+    wire::{OcrWireRequest, decode_request},
+};
 
 fn query_value(url: &str, key: &str) -> Option<String> {
     url::Url::parse(url)
@@ -241,34 +243,8 @@ async fn accepted_response_polls_to_success_with_only_credentials() {
     }
 }
 
-struct SubmissionBoundary {
-    request_count: Arc<Mutex<Vec<String>>>,
-}
-
-impl super::hooks::OcrHooks for SubmissionBoundary {
-    fn post_call(
-        &self,
-        request: super::hooks::OcrPostCallRequest,
-    ) -> super::hooks::OcrHookFuture<'_, super::hooks::OcrPostCallRequest> {
-        Box::pin(async move {
-            match self.request_count.lock().unwrap().len() {
-                1 => assert_eq!(request.original_response, json!(r#"{"submitted":true}"#)),
-                2 => assert!(
-                    request
-                        .original_response
-                        .as_str()
-                        .unwrap()
-                        .contains("succeeded")
-                ),
-                count => panic!("unexpected callback after {count} requests"),
-            }
-            Ok(request)
-        })
-    }
-}
-
 #[tokio::test]
-async fn accepted_response_runs_post_call_before_polling() {
+async fn accepted_response_emits_response_received_before_polling() {
     let (base, seen, server) = mock_server(vec![
         MockResponse {
             status: 202,
@@ -278,14 +254,24 @@ async fn accepted_response_runs_post_call_before_polling() {
         MockResponse::json(json!({"status":"succeeded"})),
     ])
     .await;
-    let request = super::LiteLLMOcrRequest {
-        hooks: Arc::new(SubmissionBoundary {
-            request_count: seen.clone(),
-        }),
-        ..wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}))
-    };
+    let request_count = seen.clone();
+    let host = LocalOcrHost::new(wire_request(
+        "azure_ai/doc-intelligence/prebuilt-read",
+        &base,
+        json!({}),
+    ))
+    .with_observer(move |event| {
+        let CallEvent::ResponseReceived { raw } = event else {
+            return;
+        };
+        match request_count.lock().unwrap().len() {
+            1 => assert_eq!(raw.body, r#"{"submitted":true}"#),
+            2 => assert!(raw.body.contains("succeeded")),
+            count => panic!("unexpected callback after {count} requests"),
+        }
+    });
 
-    perform_ocr(request).await.unwrap();
+    perform_ocr_with(host).await.unwrap();
     server.await.unwrap();
     assert_eq!(seen.lock().unwrap().len(), 2);
 }
@@ -474,44 +460,3 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() {
         assert!(error.to_string().contains("dot segment"));
     }
 }
-
-#[tokio::test]
-async fn pre_call_guardrail_receives_caller_pages_before_mapping() {
-    use std::sync::Arc;
-
-    use crate::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreCallRequest};
-
-    struct RewritePages;
-    impl OcrHooks for RewritePages {
-        fn intercepts_requests(&self) -> bool {
-            true
-        }
-
-        fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
-            Box::pin(async move {
-                assert_eq!(request.optional_params["pages"], json!([0, 2]));
-                Ok(OcrPreCallRequest {
-                    optional_params: json!({"pages": [1]}),
-                    ..request
-                })
-            })
-        }
-    }
-    let (base, seen, server) =
-        mock_server(vec![MockResponse::json(json!({"status": "succeeded"}))]).await;
-    let request = wire_request(
-        "azure_ai/doc-intelligence/prebuilt-read",
-        &base,
-        json!({"pages": [0, 2]}),
-    )
-    .with_host_hooks(Arc::new(RewritePages), None);
-    perform_ocr(request).await.unwrap();
-    server.await.unwrap();
-    let requests = seen.lock().unwrap();
-    let target = requests[0].split_whitespace().nth(1).unwrap();
-    assert_eq!(
-        query_value(&format!("{base}{target}"), "pages").as_deref(),
-        Some("2")
-    );
-    assert_eq!(requests.len(), 1);
-}

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

@@ -1,12 +1,16 @@
 use rstest::rstest;
 use serde_json::{Value, json};
 
-use crate::llms::base_llm::ocr::transformation::BaseOcrConfig;
-use crate::llms::vertex_ai::ocr::deepseek_transformation::{
-    DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig,
-    normalize_response as transform_ocr_response,
+use crate::{
+    llms::{
+        base_llm::ocr::transformation::BaseOcrConfig,
+        vertex_ai::ocr::deepseek_transformation::{
+            DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig,
+            normalize_response as transform_ocr_response,
+        },
+    },
+    ocr::types::OcrDocument,
 };
-use crate::ocr::types::OcrDocument;
 
 fn document() -> OcrDocument {
     serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap()

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

@@ -1,117 +0,0 @@
-use crate::call_lifecycle::host::{HostFailure, HostLifecycle, HostPhase};
-use crate::ocr::Error;
-
-fn run(fail_at: Option<HostPhase>, asynchronous: bool) -> (Vec<HostPhase>, Vec<Error>) {
-    let mut lifecycle = HostLifecycle::new(asynchronous);
-    let mut events = Vec::new();
-    let mut failures = Vec::new();
-    while lifecycle.phase() != HostPhase::Complete {
-        let phase = lifecycle.phase();
-        events.push(phase);
-        let result = if Some(phase) == fail_at {
-            Err(HostFailure::Error(Error::InvalidRequest(
-                "selected failure".into(),
-            )))
-        } else {
-            Ok(())
-        };
-        if let Some(error) = lifecycle.accept(result) {
-            failures.push(error);
-        }
-    }
-    (events, failures)
-}
-
-#[test]
-fn public_outcome_is_finalized_before_a_single_terminal_dispatch() {
-    for asynchronous in [false, true] {
-        let (events, failures) = run(None, asynchronous);
-        assert!(failures.is_empty());
-        assert_eq!(
-            &events[events.len() - 2..],
-            &[HostPhase::Finalize, HostPhase::Success]
-        );
-        assert_eq!(
-            events
-                .iter()
-                .filter(|phase| **phase == HostPhase::Execute)
-                .count(),
-            1
-        );
-        assert_eq!(
-            events.contains(&HostPhase::DeploymentPostCall),
-            asynchronous
-        );
-    }
-}
-
-#[test]
-fn only_provider_and_response_construction_failures_use_provider_mapping() {
-    for phase in [
-        HostPhase::Setup,
-        HostPhase::DeploymentPreCall,
-        HostPhase::Prepare,
-        HostPhase::Execute,
-        HostPhase::ConstructResponse,
-        HostPhase::DeploymentPostCall,
-        HostPhase::Finalize,
-    ] {
-        let (events, failures) = run(Some(phase), true);
-        assert_eq!(failures.len(), 1);
-        assert!(!events.contains(&HostPhase::Success));
-        let mapped = matches!(phase, HostPhase::Execute | HostPhase::ConstructResponse);
-        assert_eq!(events.contains(&HostPhase::MapFailure), mapped);
-        assert_eq!(events.contains(&HostPhase::DeploymentFailure), mapped);
-        assert_eq!(
-            &events[events.len() - 2..],
-            &[HostPhase::Failure, HostPhase::AsyncFailure]
-        );
-        assert!(
-            events
-                .iter()
-                .filter(|phase| **phase == HostPhase::Execute)
-                .count()
-                <= 1
-        );
-    }
-}
-
-#[test]
-fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_dispatch() {
-    let mut lifecycle = HostLifecycle::new(true);
-    while lifecycle.phase() != HostPhase::Execute {
-        lifecycle.accept::<Error>(Ok(()));
-    }
-    let selected = Error::InvalidRequest("provider".into());
-    assert!(matches!(
-        lifecycle.accept(Err(HostFailure::Error(selected.clone()))),
-        Some(Error::InvalidRequest(message)) if message == "provider"
-    ));
-    lifecycle.accept::<Error>(Ok(()));
-    for phase in [
-        HostPhase::DeploymentFailure,
-        HostPhase::Failure,
-        HostPhase::AsyncFailure,
-    ] {
-        assert_eq!(lifecycle.phase(), phase);
-        assert!(
-            lifecycle
-                .accept(Err(HostFailure::Error(Error::InvalidRequest(
-                    "callback".into()
-                ))))
-                .is_none()
-        );
-    }
-    assert_eq!(lifecycle.phase(), HostPhase::Complete);
-}
-
-#[test]
-fn cancellation_skips_terminal_dispatch() {
-    let mut lifecycle = HostLifecycle::new(true);
-    let error = Error::InvalidRequest("cancelled".into());
-    assert!(matches!(
-        lifecycle.accept(Err(HostFailure::Cancelled(error.clone()))),
-        Some(Error::InvalidRequest(message)) if message == "cancelled"
-    ));
-    assert_eq!(lifecycle.phase(), HostPhase::Complete);
-}

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

@@ -0,0 +1,279 @@
+use std::collections::BTreeSet;
+use std::sync::{Arc, Mutex};
+
+use litellm_callbacks::event::{RequestContext, WireRequest};
+use rstest::rstest;
+use rstest_reuse::{self, apply, template};
+use serde_json::{Map, Value, json};
+
+use super::LocalOcrHost;
+use super::test_support::{
+    MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body,
+    wire_request_with_document,
+};
+
+#[derive(Clone, Copy, Debug)]
+enum Route {
+    Mistral,
+    AzureAi,
+    VertexMistral,
+    AzureCohereParse,
+    Cohere,
+}
+
+impl Route {
+    fn model(self) -> &'static str {
+        match self {
+            Self::Mistral => "mistral/model",
+            Self::AzureAi => "azure_ai/model",
+            Self::VertexMistral => "vertex_ai/mistral-ocr-maas",
+            Self::AzureCohereParse => "azure_ai/cohere-parse",
+            Self::Cohere => "cohere/model",
+        }
+    }
+
+    fn document_type(self) -> &'static str {
+        match self {
+            Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url",
+            Self::AzureCohereParse | Self::Cohere => "image_url",
+        }
+    }
+
+    fn options(self) -> Value {
+        match self {
+            Self::Mistral | Self::AzureAi => json!({"pages": [0]}),
+            Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}),
+            Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}),
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+enum Source {
+    Inline,
+    Remote,
+    RemoteWithExtraField,
+}
+
+/// What the host does to the wire request in `before_send`.
+#[derive(Clone, Copy, Debug)]
+enum Host {
+    Detached,
+    /// What `litellm-callbacks-legacy` does before `pre_call`: every passthrough body key
+    /// is replaced by the caller's own value.
+    Realiasing,
+    ReplacesDocument,
+}
+
+const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ=";
+
+impl Host {
+    fn before_send(
+        self,
+        caller: &Map<String, Value>,
+        wire: WireRequest,
+        context: &RequestContext,
+    ) -> WireRequest {
+        let Value::Object(fields) = wire.body else {
+            return wire;
+        };
+        let body = fields
+            .into_iter()
+            .map(|(name, value)| match self {
+                Self::Detached => (name, value),
+                Self::Realiasing => {
+                    let aliased = context
+                        .passthrough_fields
+                        .contains(&name)
+                        .then(|| caller.get(&name).cloned())
+                        .flatten()
+                        .unwrap_or(value);
+                    (name, aliased)
+                }
+                Self::ReplacesDocument if name == "document" => {
+                    let document_type = value["type"].clone();
+                    let key = document_type.as_str().unwrap_or_default().to_string();
+                    (name, json!({"type": document_type, key: REPLACED_DOCUMENT}))
+                }
+                Self::ReplacesDocument => (name, value),
+            })
+            .collect();
+        WireRequest {
+            body: Value::Object(body),
+            ..wire
+        }
+    }
+}
+
+struct Sent {
+    caller: Map<String, Value>,
+    result: Result<(), crate::ocr::Error>,
+    before_send: Option<(WireRequest, RequestContext)>,
+    provider_body: Option<Value>,
+}
+
+fn caller_document(route: Route, source: Source, document_base: &str) -> Value {
+    let document_type = route.document_type();
+    let remote = format!("{document_base}/scan.png");
+    match source {
+        Source::Inline => {
+            json!({"type": document_type, document_type: "data:image/png;base64,YWJj"})
+        }
+        Source::Remote => json!({"type": document_type, document_type: remote}),
+        Source::RemoteWithExtraField => {
+            json!({"type": document_type, document_type: remote, "document_name": "scan.png"})
+        }
+    }
+}
+
+async fn send(route: Route, source: Source, host: Host, document_base: &str) -> Sent {
+    let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await;
+    let document = caller_document(route, source, document_base);
+    let caller: Map<String, Value> = route
+        .options()
+        .as_object()
+        .unwrap()
+        .clone()
+        .into_iter()
+        .chain([("document".to_string(), document.clone())])
+        .collect();
+    let observed = Arc::new(Mutex::new(None));
+    let captured = observed.clone();
+    let host_caller = caller.clone();
+    let request = wire_request_with_document(route.model(), &base, document, route.options());
+    let local = LocalOcrHost::new(request).with_before_send(move |wire, context| {
+        *captured.lock().unwrap() = Some((wire.clone(), context.clone()));
+        Ok(host.before_send(&host_caller, wire, context))
+    });
+    let result = perform_ocr_with(local).await.map(|_| ());
+    match result {
+        Ok(()) => provider.await.unwrap(),
+        Err(_) => provider.abort(),
+    }
+    let provider_body = seen
+        .lock()
+        .unwrap()
+        .first()
+        .map(|request| request_body(request));
+    let before_send = observed.lock().unwrap().take();
+    Sent {
+        caller,
+        result,
+        before_send,
+        provider_body,
+    }
+}
+
+fn served_document_uri() -> String {
+    use base64::Engine;
+    format!(
+        "data:image/png;base64,{}",
+        base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT)
+    )
+}
+
+#[template]
+#[rstest]
+fn every_route_and_source(
+    #[values(
+        Route::Mistral,
+        Route::AzureAi,
+        Route::VertexMistral,
+        Route::AzureCohereParse,
+        Route::Cohere
+    )]
+    route: Route,
+    #[values(Source::Inline, Source::Remote, Source::RemoteWithExtraField)] source: Source,
+) {
+}
+
+#[template]
+#[rstest]
+fn every_route(
+    #[values(
+        Route::Mistral,
+        Route::AzureAi,
+        Route::VertexMistral,
+        Route::AzureCohereParse,
+        Route::Cohere
+    )]
+    route: Route,
+) {
+}
+
+#[template]
+#[rstest]
+#[case::azure_ai(Route::AzureAi)]
+#[case::vertex_mistral(Route::VertexMistral)]
+#[case::azure_cohere_parse(Route::AzureCohereParse)]
+fn inlining_routes(#[case] route: Route) {}
+
+#[apply(every_route_and_source)]
+#[tokio::test]
+async fn passthrough_fields_are_exactly_the_caller_values_sent_unchanged(
+    route: Route,
+    source: Source,
+) {
+    let (document_base, _documents) = document_server().await;
+    let sent = send(route, source, Host::Detached, &document_base).await;
+    sent.result.unwrap();
+    let (wire, context) = sent.before_send.unwrap();
+    let passthrough: BTreeSet<&str> = context.passthrough_fields.iter().collect();
+    let unchanged: BTreeSet<&str> = sent
+        .caller
+        .iter()
+        .filter(|(name, value)| wire.body.get(name.as_str()) == Some(*value))
+        .map(|(name, _)| name.as_str())
+        .collect();
+    assert_eq!(
+        passthrough,
+        unchanged,
+        "body: {:#}\ncaller: {:#}",
+        wire.body,
+        Value::Object(sent.caller.clone())
+    );
+}
+
+#[apply(every_route_and_source)]
+#[tokio::test]
+async fn realiasing_leaves_the_provider_request_unchanged(route: Route, source: Source) {
+    let (document_base, _documents) = document_server().await;
+    let detached = send(route, source, Host::Detached, &document_base).await;
+    let realiased = send(route, source, Host::Realiasing, &document_base).await;
+    detached.result.unwrap();
+    realiased.result.unwrap();
+    assert_eq!(realiased.provider_body, detached.provider_body);
+}
+
+#[apply(inlining_routes)]
+#[tokio::test]
+async fn inlining_routes_send_the_downloaded_document(
+    route: Route,
+    #[values(Host::Detached, Host::Realiasing)] host: Host,
+) {
+    let (document_base, _documents) = document_server().await;
+    let sent = send(route, Source::Remote, host, &document_base).await;
+    sent.result.unwrap();
+    assert_eq!(
+        sent.provider_body.unwrap()["document"][route.document_type()],
+        json!(served_document_uri())
+    );
+}
+
+#[apply(every_route)]
+#[tokio::test]
+async fn document_replaced_by_the_host_reaches_the_provider(route: Route) {
+    let (document_base, _documents) = document_server().await;
+    let sent = send(
+        route,
+        Source::Remote,
+        Host::ReplacesDocument,
+        &document_base,
+    )
+    .await;
+    sent.result.unwrap();
+    assert_eq!(
+        sent.provider_body.unwrap()["document"][route.document_type()],
+        json!(REPLACED_DOCUMENT)
+    );
+}

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

@@ -1,11 +1,15 @@
 use std::sync::{Arc, Mutex};
 
 use serde_json::{Value, json};
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
-use tokio::net::TcpListener;
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::TcpListener,
+};
 
-use crate::ocr::wire::{OcrWireRequest, decode_request};
-use crate::ocr::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
+use crate::ocr::{
+    LiteLLMOcrRequest, LiteLLMOcrResponse, LocalOcrHost, OcrClient, ocr_machine,
+    wire::{OcrWireRequest, decode_request},
+};
 
 pub(crate) fn ocr_client() -> OcrClient {
     let document_http = reqwest::Client::builder()
@@ -21,10 +25,30 @@ pub(crate) async fn perform_ocr(
     ocr_client().perform(request).await
 }
 
+pub(crate) async fn perform_ocr_with(
+    host: LocalOcrHost,
+) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
+    litellm_callbacks::run::run(ocr_machine(ocr_client()), &host).await
+}
+
 pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest {
+    wire_request_with_document(
+        model,
+        base,
+        json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}),
+        options,
+    )
+}
+
+pub(crate) fn wire_request_with_document(
+    model: &str,
+    base: &str,
+    document: Value,
+    options: Value,
+) -> LiteLLMOcrRequest {
     decode_request(OcrWireRequest {
         model: model.into(),
-        document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}),
+        document,
         api_key: Some("test-key".into()),
         api_base: Some(base.into()),
         custom_llm_provider: None,
@@ -50,6 +74,32 @@ pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOc
     request.with_document(document.into())
 }
 
+pub(crate) fn request_body(request: &str) -> Value {
+    serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
+}
+
+pub(crate) const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document";
+
+/// Serves [`SERVED_DOCUMENT`] as `image/png` to every connection until aborted.
+pub(crate) async fn document_server() -> (String, tokio::task::JoinHandle<()>) {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let base = format!("http://{}", listener.local_addr().unwrap());
+    let task = tokio::spawn(async move {
+        loop {
+            let (mut socket, _) = listener.accept().await.unwrap();
+            let mut buffer = [0u8; 4096];
+            let _ = socket.read(&mut buffer).await.unwrap();
+            let head = format!(
+                "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
+                SERVED_DOCUMENT.len()
+            );
+            socket.write_all(head.as_bytes()).await.unwrap();
+            socket.write_all(SERVED_DOCUMENT).await.unwrap();
+        }
+    });
+    (base, task)
+}
+
 pub(crate) struct MockResponse {
     pub status: u16,
     pub headers: Vec<(&'static str, String)>,
@@ -123,3 +173,13 @@ pub(crate) async fn mock_server(
     });
     (base, requests, task)
 }
+
+pub(crate) fn header<'a>(request: &'a str, name: &str) -> Option<&'a str> {
+    request
+        .lines()
+        .take_while(|line| !line.is_empty())
+        .find_map(|line| {
+            let (key, value) = line.split_once(':')?;
+            key.eq_ignore_ascii_case(name).then(|| value.trim())
+        })
+}

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

@@ -1,10 +1,11 @@
-use std::sync::Arc;
-
+use litellm_callbacks::event::{CallEvent, WireRequest};
 use rstest::rstest;
 use serde_json::{Value, json};
 
-use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrPostCallRequest};
-use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
+use super::{
+    LocalOcrHost,
+    test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
+};
 
 fn request_body(request: &str) -> Value {
     serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
@@ -129,38 +130,24 @@ async fn data_uri_upload_preserves_multipart_headers(
     }
 }
 
-struct ParseBoundary {
-    request_count: Arc<std::sync::Mutex<Vec<String>>>,
-}
-
-impl OcrHooks for ParseBoundary {
-    fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> {
-        Box::pin(async move {
-            assert_eq!(self.request_count.lock().unwrap().len(), 2);
-            assert_eq!(
-                request.original_response,
-                json!(r#"{"result":{"chunks":[]}}"#)
-            );
-            Ok(request)
-        })
-    }
-}
-
 #[tokio::test]
-async fn post_call_stays_after_reducto_upload_and_parse() {
+async fn response_received_stays_after_reducto_upload_and_parse() {
     let (base, seen, server) = mock_server(vec![
         MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
         MockResponse::json(json!({"result":{"chunks":[]}})),
     ])
     .await;
-    let request = super::LiteLLMOcrRequest {
-        hooks: Arc::new(ParseBoundary {
-            request_count: seen.clone(),
-        }),
-        ..wire_request("reducto/parse-v3", &base, json!({}))
-    };
+    let request_count = seen.clone();
+    let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))).with_observer(
+        move |event| {
+            if let CallEvent::ResponseReceived { raw } = event {
+                assert_eq!(request_count.lock().unwrap().len(), 2);
+                assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#);
+            }
+        },
+    );
 
-    perform_ocr(request).await.unwrap();
+    perform_ocr_with(host).await.unwrap();
     server.await.unwrap();
     assert_eq!(seen.lock().unwrap().len(), 2);
 }
@@ -300,38 +287,66 @@ async fn facade_omits_native_response_by_default_and_preserves_auth_priority() {
     );
 }
 
-struct RewriteDocument;
+#[tokio::test]
+async fn native_format_retains_the_provider_response() {
+    let raw = json!({
+        "result":{"chunks":[{"content":"native OCR response"}]},
+        "usage":{"num_pages":1}
+    });
+    let (base, _, server) = mock_server(vec![MockResponse::json(raw.clone())]).await;
+    let request = super::test_support::with_source(
+        wire_request("reducto/parse-v3", &base, json!({"req_format":"native"})),
+        "reducto://ready.pdf",
+    );
 
-impl OcrHooks for RewriteDocument {
-    fn intercepts_requests(&self) -> bool {
-        true
-    }
+    let response = perform_ocr(request).await.unwrap();
+    server.await.unwrap();
 
-    fn during_call(
-        &self,
-        request: OcrDuringCallRequest,
-    ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
-        Box::pin(async move {
-            assert_eq!(
-                request.body["document_url"],
-                "data:application/pdf;base64,YWJj"
-            );
-            Ok(OcrDuringCallRequest {
-                body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}),
-                ..request
-            })
-        })
-    }
+    assert_eq!(response.pages[0].markdown, "native OCR response");
+    assert_eq!(response.provider_native_response.as_ref(), raw.as_object());
+}
+
+#[tokio::test]
+async fn unknown_model_reaches_parse_and_keeps_its_name() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
+        "result":{"chunks":[{"content":"future model response"}]}
+    }))])
+    .await;
+    let request = super::test_support::with_source(
+        wire_request("reducto/future-parse-model", &base, json!({})),
+        "reducto://ready.pdf",
+    );
+
+    let response = perform_ocr(request).await.unwrap();
+    server.await.unwrap();
+
+    assert_eq!(response.model, "future-parse-model");
+    assert_eq!(response.pages[0].markdown, "future model response");
+    let requests = seen.lock().unwrap();
+    assert!(requests[0].starts_with("POST /parse "));
+    assert_eq!(
+        request_body(&requests[0]),
+        json!({"input":"reducto://ready.pdf"})
+    );
 }
 
 #[tokio::test]
 async fn guardrail_rewrites_document_before_upload() {
     let (base, seen, server) =
         mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await;
-    let mut request = wire_request("reducto/parse-v3", &base, json!({}));
-    request.hooks = Arc::new(RewriteDocument);
+    let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({})))
+        .with_before_send(|wire, _| {
+            assert_eq!(
+                wire.body["document_url"],
+                "data:application/pdf;base64,YWJj"
+            );
+            Ok(WireRequest {
+                body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}),
+                ..wire
+            })
+        });
 
-    perform_ocr(request).await.unwrap();
+    perform_ocr_with(host).await.unwrap();
     server.await.unwrap();
     let requests = seen.lock().unwrap();
     assert_eq!(requests.len(), 1);

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

@@ -102,10 +102,14 @@ async fn request_controlled_api_base_is_rejected_before_vertex_auth() {
 async fn adapters_build_complete_requests_and_share_mistral_normalization() {
     use std::time::Duration;
 
-    use crate::llms::base_llm::ocr::transformation::BaseOcrConfig;
-    use crate::llms::mistral::ocr::transformation::MistralOcrConfig;
-    use crate::llms::vertex_ai::ocr::transformation::VertexAiOcrConfig;
-    use crate::ocr::test_support::ocr_client;
+    use crate::{
+        llms::{
+            base_llm::ocr::transformation::BaseOcrConfig,
+            mistral::ocr::transformation::MistralOcrConfig,
+            vertex_ai::ocr::transformation::VertexAiOcrConfig,
+        },
+        ocr::test_support::ocr_client,
+    };
 
     let client = ocr_client();
     let options = json!({
@@ -121,10 +125,12 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() {
         options.clone(),
     );
     let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options);
-    let direct =
-        crate::ocr::prepare::prepare_request(super::test_support::resolved_request(direct));
-    let vertex =
-        crate::ocr::prepare::prepare_request(super::test_support::resolved_request(vertex));
+    let direct = crate::ocr::prepare::prepare_request_for_test(
+        super::test_support::resolved_request(direct),
+    );
+    let vertex = crate::ocr::prepare::prepare_request_for_test(
+        super::test_support::resolved_request(vertex),
+    );
     let direct_http = MistralOcrConfig
         .prepare_request(&direct, &client)
         .await

litellm-rust/crates/host-python/AGENTS.md

@@ -1,7 +1,9 @@
 - Target invariants; implementation and runtime validation may lag these rules
-- Keep this crate a small, domain-neutral foundation: Python/Serde conversion and interpreter-boundary utilities
-  - No LiteLLM domain dependencies, route types, callback policy, public API registration or cdylib build features
-  - Generic code alone does not justify extraction: runtime integration stays in `python-bridge/src/execution.rs`, host adaptation in its `lifecycle.rs`
+- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `CallbackAdapter`/`RouteHost` traits
+  - No LiteLLM domain dependencies beyond `litellm-callbacks`: no route types, no `Logging` policy, no public API registration, no cdylib build features
+  - The driver emits `Succeeded` or `Failed` exactly once and never dispatches after a cancellation; which Python objects consume those events is the adapter's business
+  - `RouteHost::invoke` receives the keyword view the adapter's `begin` returned, not the caller's dict; a route host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance)
+  - A failure that surfaces inside the call, including a host op the call asked for, is mapped through the route's `map_failure`; a failure in `begin` or `after_success` is raised as is
 - Use standard PyO3 ownership and conversion APIs
   - Prefer `Bound<'py, T>` for attached operations/results, `Py<T>` for retention; binding/unbinding does not copy payloads
   - Use `pythonize` for selected Serde data, never a JSON-text round trip; share conversion with `Pythonized<T>`
@@ -10,7 +12,8 @@
 - Use `Python::detach` for Rust-only work; Python operations require attachment
   - Keep diagnostic counters in the consumer; wrapper invocations do not measure every interpreter release
   - Release exclusive class borrows/locks before Python calls or decrements that can invoke finalizers; expose retained Python edges to GC without calling Python during traversal
-- Keep coroutine driving in the shared Python driver and native adapter
-  - Driver: `litellm/rust_bridge/lifecycle.py`; handle: `python-bridge/src/lifecycle.rs`; native-backed behavior tests: `python-bridge/tests/lifecycle.py`
+- Keep coroutine driving in the shared Python driver and the native handle
+  - Driver: `litellm/rust_bridge/lifecycle.py`; handle: `src/handle.rs`; call driver: `src/driver.rs`; native-backed behavior tests: `tests/lifecycle.py`
+  - Every adapter suspension is awaited inline in the caller's task; `into_future` creates a separate task and cannot satisfy this contract
 - References: [ownership](https://pyo3.rs/v0.29.2/types.html), [conversions](https://pyo3.rs/v0.29.2/conversions/traits.html), [pythonize errors](https://docs.rs/pythonize/0.29.0/src/pythonize/error.rs.html)
   - [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [re-entry](https://pyo3.rs/v0.29.2/class/call.html), [parallelism](https://pyo3.rs/v0.29.2/parallelism.html), [async delivery source](https://docs.rs/pyo3-async-runtimes/0.29.0/src/pyo3_async_runtimes/generic.rs.html)

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

@@ -0,0 +1,19 @@
+[package]
+name = "litellm-host-python"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+futures-util.workspace = true
+litellm-callbacks.workspace = true
+pyo3.workspace = true
+pyo3-async-runtimes.workspace = true
+pythonize.workspace = true
+serde.workspace = true
+tokio = { workspace = true, features = ["sync"] }
+
+[dev-dependencies]
+rstest.workspace = true
+serde_json.workspace = true

litellm-rust/crates/host-python/src/adapter.rs

@@ -0,0 +1,104 @@
+use litellm_callbacks::event::{CallEvent, RequestContext, Timing, WireRequest};
+use litellm_callbacks::route::Route;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::gc::{PyTraverseError, PyVisit};
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+
+pub fn missing_state() -> PyErr {
+    PyRuntimeError::new_err("missing native call state")
+}
+
+/// What an adapter step produced: either the value the driver asked for, or a Python
+/// awaitable the driver hands back to the caller's task before asking again.
+pub enum AdapterStep {
+    Await(Py<PyAny>),
+    Arguments(Py<PyDict>),
+    Wire(Box<WireRequest>),
+    Response(Py<PyAny>),
+    Done,
+}
+
+/// The host-typed value the driver attaches to a terminal event.
+pub enum PublicValue<'a> {
+    Response(&'a Py<PyAny>),
+    Error(&'a PyErr),
+}
+
+/// One consumer of a call's lifecycle on the Python side. The driver calls the steps in
+/// order: `begin` before the machine starts, `before_send` and `emit` while it runs,
+/// `after_success` and one terminal `emit` after it completes. Whenever a step returns
+/// [`AdapterStep::Await`], the driver awaits it in the caller's task and continues the
+/// same step through `resume`.
+///
+/// A step that fails with an ordinary exception fails the call with that exception,
+/// except on a terminal event, where the adapter is expected to report and swallow its
+/// own errors. An exception that is not a `PyException`, such as a cancellation, ends
+/// the call without further dispatch.
+pub trait CallbackAdapter: Send + Sync {
+    fn begin(
+        &mut self,
+        py: Python<'_>,
+        arguments: Py<PyDict>,
+        started_at: f64,
+    ) -> PyResult<AdapterStep>;
+
+    fn before_send(
+        &mut self,
+        py: Python<'_>,
+        wire: Box<WireRequest>,
+        context: &RequestContext,
+    ) -> PyResult<AdapterStep>;
+
+    fn after_success(
+        &mut self,
+        py: Python<'_>,
+        response: Py<PyAny>,
+        timing: Timing,
+    ) -> PyResult<AdapterStep>;
+
+    fn emit(
+        &mut self,
+        py: Python<'_>,
+        event: &CallEvent,
+        public: Option<PublicValue<'_>>,
+    ) -> PyResult<AdapterStep>;
+
+    fn resume(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<AdapterStep>;
+
+    fn close(&mut self, py: Python<'_>);
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>;
+}
+
+/// The Python side of one route: answers the route's own operations, builds the public
+/// response and maps failures to public exceptions.
+pub trait RouteHost: Send + Sync {
+    type Route: Route;
+
+    /// `arguments` is the keyword view the callback adapter's `begin` produced, not the
+    /// caller's own dict. A route host that projects from it inherits whatever that
+    /// adapter rewrote.
+    fn invoke(
+        &mut self,
+        py: Python<'_>,
+        arguments: &Bound<'_, PyDict>,
+        op: <Self::Route as Route>::Op,
+    ) -> PyResult<<Self::Route as Route>::OpResult>;
+
+    fn complete(
+        &mut self,
+        py: Python<'_>,
+        response: <Self::Route as Route>::Response,
+    ) -> PyResult<Py<PyAny>>;
+
+    fn native_error(error: <Self::Route as Route>::Error) -> PyErr;
+
+    fn host_error(error: &PyErr) -> <Self::Route as Route>::Error;
+
+    fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult<PyErr>;
+
+    fn close(&mut self, py: Python<'_>);
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>;
+}

litellm-rust/crates/host-python/src/callable.rs

@@ -0,0 +1,135 @@
+//! Failures raised by a caller-supplied Python callable.
+
+use pyo3::exceptions::{PyException, PyRuntimeError, PyTypeError};
+use pyo3::prelude::*;
+use pyo3::types::PyString;
+
+/// Reports a caller-supplied callable's failure under `template`, a Python format string
+/// with one field for the original exception, while leaving alone the failures a caller
+/// can already read: a `TypeError`, so a rejected return value is not reported twice, and
+/// anything that is not a `PyException`, a cancellation for example. Everything else
+/// becomes a `RuntimeError` carrying the original as both its `__cause__` and its
+/// `__context__`, with the message rendered by Python so the exception's own `__format__`
+/// is honored. A `__format__` that raises surfaces as that failure instead, with the
+/// original attached as its context.
+pub fn wrap_failure<T>(py: Python<'_>, template: &str, result: PyResult<T>) -> PyResult<T> {
+    result.map_err(|error| {
+        if error.is_instance_of::<PyTypeError>(py) || !error.is_instance_of::<PyException>(py) {
+            return error;
+        }
+        match PyString::new(py, template).call_method1("format", (error.value(py),)) {
+            Ok(message) => {
+                let wrapped = PyRuntimeError::new_err(message.unbind());
+                wrapped.set_context(py, Some(error.clone_ref(py)));
+                wrapped.set_cause(py, Some(error));
+                wrapped
+            }
+            Err(format_error) => {
+                format_error.set_context(py, Some(error));
+                format_error
+            }
+        }
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use pyo3::types::PyDict;
+
+    use super::*;
+
+    const TEMPLATE: &str = "Failed to reach the caller: {}";
+
+    fn raised<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> {
+        locals.get_item(name).unwrap().unwrap()
+    }
+
+    fn failure<'py>(error: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
+        Err(PyErr::from_value(error.clone()))
+    }
+
+    #[test]
+    fn only_ordinary_exceptions_are_reported_under_the_template() {
+        crate::initialize_python();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            py.run(
+                pyo3::ffi::c_str!(
+                    r#"
+class CallerError(Exception):
+    def __format__(self, specification):
+        return 'unavailable'
+ordinary = CallerError('must use __format__')
+type_error = TypeError('signature')
+abort = KeyboardInterrupt('cancelled')
+"#
+                ),
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+
+            let original = raised(&locals, "ordinary");
+            let wrapped = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err();
+            assert!(wrapped.is_instance_of::<PyRuntimeError>(py));
+            assert!(wrapped.cause(py).unwrap().value(py).is(&original));
+            assert!(
+                wrapped
+                    .value(py)
+                    .getattr("__context__")
+                    .unwrap()
+                    .is(&original)
+            );
+            assert_eq!(
+                wrapped.value(py).str().unwrap().to_str().unwrap(),
+                "Failed to reach the caller: unavailable"
+            );
+
+            for name in ["type_error", "abort"] {
+                let original = raised(&locals, name);
+                let error = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err();
+                assert!(error.value(py).is(&original));
+            }
+        });
+    }
+
+    #[test]
+    fn a_raising_format_surfaces_instead_of_the_report_and_keeps_the_original_as_context() {
+        crate::initialize_python();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            py.run(
+                pyo3::ffi::c_str!(
+                    r#"
+class Unformattable(Exception):
+    def __format__(self, specification):
+        raise ValueError('formatting failed')
+original = Unformattable('cannot render')
+"#
+                ),
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+
+            let original = raised(&locals, "original");
+            let error = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err();
+            assert!(error.is_instance_of::<pyo3::exceptions::PyValueError>(py));
+            assert!(
+                error
+                    .value(py)
+                    .getattr("__context__")
+                    .unwrap()
+                    .is(&original)
+            );
+        });
+    }
+
+    #[test]
+    fn successful_results_pass_through_untouched() {
+        crate::initialize_python();
+        Python::attach(|py| {
+            assert_eq!(wrap_failure(py, TEMPLATE, Ok(7)).unwrap(), 7);
+        });
+    }
+}

litellm-rust/crates/host-python/src/execution.rs

@@ -4,15 +4,15 @@ use std::pin::Pin;
 use std::task::{Context, Poll, Waker};
 use std::time::Duration;
 
+use crate::{Pythonized, panic_to_pyerr, release_gil};
 use futures_util::FutureExt;
-use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use serde::Serialize;
 use tokio::runtime::{Handle, Runtime};
 use tokio::time::{self, MissedTickBehavior};
 
-pub(crate) fn run_sync<T, E, F>(
+pub fn run_sync<T, E, F>(
     py: Python<'_>,
     future: F,
     map_error: fn(E) -> PyErr,
@@ -30,7 +30,7 @@ where
     )
 }
 
-pub(crate) fn run_sync_value<T, F>(py: Python<'_>, future: F) -> PyResult<T>
+pub fn run_sync_value<T, F>(py: Python<'_>, future: F) -> PyResult<T>
 where
     T: Send + 'static,
     F: Future<Output = PyResult<T>> + Send + 'static,
@@ -73,7 +73,7 @@ where
     Pythonized(result).into_pyobject(py).map(Bound::unbind)
 }
 
-pub(crate) fn run_async<T, E, F>(
+pub fn run_async<T, E, F>(
     py: Python<'_>,
     future: F,
     map_error: fn(E) -> PyErr,
@@ -90,7 +90,7 @@ where
     })
 }
 
-pub(crate) fn run_async_value<T, F>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
+pub fn run_async_value<T, F>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
 where
     T: for<'py> IntoPyObject<'py> + Send + 'static,
     F: Future<Output = PyResult<T>> + Send + 'static,
@@ -98,7 +98,7 @@ where
     pyo3_async_runtimes::tokio::future_into_py(py, async move { catch_future_panic(future).await? })
 }
 
-pub(crate) fn poll_async_value<T, F>(py: Python<'_>, future: Pin<&mut F>) -> PyResult<Poll<T>>
+pub fn poll_async_value<T, F>(py: Python<'_>, future: Pin<&mut F>) -> PyResult<Poll<T>>
 where
     T: Send,
     F: Future<Output = PyResult<T>> + Send,
@@ -158,14 +158,14 @@ where
 #[cfg(test)]
 mod tests {
     use std::ffi::CString;
-    use std::future::poll_fn;
-    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::future::{pending, poll_fn};
+    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
     use std::sync::{Arc, mpsc};
     use std::task::Poll;
     use std::thread;
     use std::time::Instant;
 
-    use litellm_core::messages::Error;
+    use pyo3::exceptions::PyLookupError;
     use pyo3::panic::PanicException;
     use pyo3::types::{PyDict, PyModule};
     use rstest::{fixture, rstest};
@@ -188,10 +188,19 @@ mod tests {
     #[fixture]
     #[once]
     fn initialized_python() -> InitializedPython {
-        Python::initialize();
+        crate::initialize_python();
         InitializedPython
     }
 
+    #[derive(Debug)]
+    struct Error(String);
+
+    impl std::fmt::Display for Error {
+        fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            formatter.write_str(&self.0)
+        }
+    }
+
     fn runtime_error(error: Error) -> PyErr {
         PyRuntimeError::new_err(error.to_string())
     }
@@ -200,6 +209,52 @@ mod tests {
         panic!("error mapper panicked")
     }
 
+    static ECHO_FUTURE_DROPPED: AtomicBool = AtomicBool::new(false);
+
+    struct EchoDropGuard;
+
+    impl Drop for EchoDropGuard {
+        fn drop(&mut self) {
+            ECHO_FUTURE_DROPPED.store(true, Ordering::SeqCst);
+        }
+    }
+
+    fn echo_error(error: Error) -> PyErr {
+        if error.0 == "panic in mapper" {
+            panic!("error mapper panicked")
+        }
+        PyLookupError::new_err(error.0)
+    }
+
+    #[pyfunction]
+    fn async_echo(py: Python<'_>, value: String) -> PyResult<Bound<'_, PyAny>> {
+        ECHO_FUTURE_DROPPED.store(false, Ordering::SeqCst);
+        let drop_guard = (value == "pending").then_some(EchoDropGuard);
+        run_async(
+            py,
+            async move {
+                let _drop_guard = drop_guard;
+                tokio::task::yield_now().await;
+                match value.as_str() {
+                    "error" => Err(Error("mapped error".into())),
+                    "map_panic" => Err(Error("panic in mapper".into())),
+                    "panic" => panic!("route future panicked"),
+                    "pending" => {
+                        pending::<()>().await;
+                        unreachable!()
+                    }
+                    _ => Ok(value),
+                }
+            },
+            echo_error,
+        )
+    }
+
+    #[pyfunction]
+    fn echo_future_dropped() -> bool {
+        ECHO_FUTURE_DROPPED.load(Ordering::SeqCst)
+    }
+
     struct PanickingOutput;
 
     static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0);
@@ -439,7 +494,7 @@ mod tests {
         python.attach(|py| {
             let error = run_sync::<bool, Error, _>(
                 py,
-                async { Err(Error::InvalidRequest("invalid".to_string())) },
+                async { Err(Error("invalid".to_string())) },
                 panicking_error_mapper,
             )
             .expect_err("panicked mapper should become a Python exception");
@@ -572,4 +627,77 @@ asyncio.run(exercise())
                 .expect("result delivery should leave Tokio workers responsive");
         });
     }
+
+    #[rstest]
+    fn async_runner_delivers_values_and_errors_and_drops_cancelled_futures(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
+        python.attach(|py| {
+            let module = PyModule::new(py, "runtime").expect("module should be created");
+            for function in [
+                wrap_pyfunction!(async_echo, &module).expect("function should wrap"),
+                wrap_pyfunction!(echo_future_dropped, &module).expect("function should wrap"),
+            ] {
+                module
+                    .add_function(function)
+                    .expect("function should register");
+            }
+            let locals = PyDict::new(py);
+            locals
+                .set_item("runtime", &module)
+                .expect("module should enter Python locals");
+            let code = CString::new(
+                r#"
+import asyncio
+
+async def exercise():
+    assert await runtime.async_echo("value") == "value"
+
+    try:
+        await runtime.async_echo("error")
+    except LookupError as error:
+        assert str(error) == "mapped error"
+    else:
+        raise AssertionError("mapped error was not raised")
+
+    try:
+        await runtime.async_echo("panic")
+    except BaseException as error:
+        assert type(error).__name__ == "PanicException"
+        assert str(error) == "route future panicked"
+    else:
+        raise AssertionError("panic was not raised")
+
+    try:
+        await runtime.async_echo("map_panic")
+    except BaseException as error:
+        assert type(error).__name__ == "PanicException"
+        assert str(error) == "error mapper panicked"
+    else:
+        raise AssertionError("mapper panic was not raised")
+
+    task = asyncio.ensure_future(runtime.async_echo("pending"))
+    await asyncio.sleep(0)
+    task.cancel()
+    try:
+        await task
+    except asyncio.CancelledError:
+        pass
+    else:
+        raise AssertionError("cancelled route completed")
+
+    for _ in range(100):
+        if runtime.echo_future_dropped():
+            break
+        await asyncio.sleep(0.001)
+    assert runtime.echo_future_dropped()
+
+asyncio.run(exercise())
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("async route contract should hold");
+        });
+    }
 }

litellm-rust/crates/host-python/src/handle.rs

@@ -1,16 +1,16 @@
 use std::panic::{AssertUnwindSafe, catch_unwind};
 
-use litellm_python_interop::panic_to_pyerr;
+use crate::panic_to_pyerr;
 use pyo3::exceptions::{PyBaseException, PyRuntimeError};
 use pyo3::gc::{PyTraverseError, PyVisit};
 use pyo3::prelude::*;
 
-pub(super) enum ExecutionStep {
+pub enum ExecutionStep {
     Return(Py<PyAny>),
     Await(Py<PyAny>),
 }
 
-pub(super) trait ExecutionBody: Send + Sync {
+pub trait ExecutionBody: Send + Sync {
     fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep>;
     fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>;
 }
@@ -23,12 +23,12 @@ enum ExecutionState {
 }
 
 #[pyclass]
-pub(super) struct Execution {
+pub struct Execution {
     state: ExecutionState,
 }
 
 impl Execution {
-    pub(super) fn new(body: impl ExecutionBody + 'static) -> Self {
+    pub fn new(body: impl ExecutionBody + 'static) -> Self {
         Self {
             state: ExecutionState::Created(Box::new(body)),
         }