Merge pull request #41752 from BerriAI/litellm_ocr_callbacks_legacy_contract

refactor(rust): isolate legacy callback contract

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)),
         }