wip

.github/workflows/test-rust.yml

@@ -4,6 +4,13 @@ on:
   push:
     paths:
       - "litellm-rust/**"
+      - "litellm/rust_bridge/**"
+      - "litellm/ocr/**"
+      - "litellm/llms/base_llm/ocr/**"
+      - "litellm/llms/custom_httpx/llm_http_handler.py"
+      - "tests/test_litellm/ocr/**"
+      - "tests/test_litellm/conftest.py"
+      - "Makefile"
       - ".cargo/**"
       - "pyproject.toml"
       - "uv.lock"
@@ -21,6 +28,13 @@ on:
       - "litellm_**"
     paths:
       - "litellm-rust/**"
+      - "litellm/rust_bridge/**"
+      - "litellm/ocr/**"
+      - "litellm/llms/base_llm/ocr/**"
+      - "litellm/llms/custom_httpx/llm_http_handler.py"
+      - "tests/test_litellm/ocr/**"
+      - "tests/test_litellm/conftest.py"
+      - "Makefile"
       - ".cargo/**"
       - "pyproject.toml"
       - "uv.lock"
@@ -128,6 +142,9 @@ jobs:
 
       - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl
 
+      - name: Require native OCR pytest acceptance
+        run: make test-rust-ocr RUST_OCR_WHEEL="$(realpath dist/*.whl)"
+
       - name: Check Python fixtures for Cargo tests
         run: make lint-rust-python-fixtures
 

Makefile

@@ -11,7 +11,7 @@
 	lint-test-quality lint-test-quality-budget-update \
 	install-dev install-proxy-dev install-test-deps install-hooks \
 	install-helm-unittest check-circular-imports check-import-safety check check-inner pre-commit \
-	lint-install lint-fetch-base bootstrap install-rust-python-test-deps test-rust-python lint-rust-python-fixtures
+	lint-install lint-fetch-base bootstrap install-rust-python-test-deps test-rust-python test-rust-ocr lint-rust-python-fixtures
 
 # Default target
 help:
@@ -57,6 +57,7 @@ help:
 	@echo "  make test-unit-helm     - Run helm unit tests"
 	@echo "  make test-rust-extension - Build the Rust extension and run its public Python tests"
 	@echo "  make test-rust-python   - Run ignored Python-integrated Cargo tests"
+	@echo "  make test-rust-ocr      - Build a wheel and require native OCR pytest acceptance"
 	@echo "  make lint-rust-python-fixtures - Check Rust test Python fixtures with Ruff"
 	@echo ""
 	@echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide"
@@ -307,6 +308,23 @@ test-rust-extension:
 	LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
 	"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
 
+test-rust-ocr:
+	@temporary=$$(mktemp -d) && \
+	trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \
+	if [ -n "$(RUST_OCR_WHEEL)" ]; then \
+		wheel="$(RUST_OCR_WHEEL)"; \
+	else \
+		$(UV) build --wheel --out-dir "$$temporary/wheels" || exit $$?; \
+		set -- "$$temporary"/wheels/*.whl; \
+		[ "$$#" -eq 1 ] || exit 1; \
+		wheel="$$1"; \
+	fi && \
+	UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --no-default-groups --group dev --extra proxy && \
+	$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$wheel" && \
+	LITELLM_LOCAL_MODEL_COST_MAP=True "$$temporary/venv/bin/python" -I -c 'from litellm.rust_bridge import _native; assert callable(_native.ocr) and callable(_native.aocr)' && \
+	LITELLM_REQUIRE_NATIVE_OCR=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
+	"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib tests/test_litellm/ocr/test_rust_bridge.py -v
+
 test-rust-python: install-rust-python-test-deps
 	@python=$$($(UV_RUN) python -c 'import sys; print(sys.executable)') && \
 	site_packages=$$("$$python" -c 'import os, sysconfig; print(os.pathsep.join(dict.fromkeys(sysconfig.get_path(key) for key in ("purelib", "platlib"))))') && \

litellm-rust/README.md

@@ -73,7 +73,28 @@ make lint-rust-python-fixtures
 `lint-rust-python-fixtures` runs pinned Ruff lint and formatting checks without
 syncing the project environment
 
-These tests validate retained callback identity, mutation, invocation context,
+Run the native OCR acceptance gate from the repository root:
+
+```bash
+make test-rust-ocr
+```
+
+This builds the current release wheel, installs locked SDK dependencies, the
+`dev` test group, and the `proxy` extra in a temporary Python 3.12 environment,
+then installs the wheel without resolving dependencies again. The proxy extra
+is needed by the shared pytest fixtures. Python isolated mode and pytest's
+importlib mode keep the checkout from shadowing the installed wheel
+
+The gate checks that native `ocr` and `aocr` are importable, then runs
+`tests/test_litellm/ocr/test_rust_bridge.py` with
+`LITELLM_REQUIRE_NATIVE_OCR=1`, so unavailable native OCR fails instead of
+skipping. CI uses `make test-rust-ocr RUST_OCR_WHEEL=/absolute/path/to/current.whl`
+to test the release wheel it just built. The stdlib-only
+`native_route_wheel_test.py` also exercises sync/async OCR through a small
+boundary, including 429 handling in `finish`/`afinish`, alongside the other
+native routes
+
+The Python-integrated Cargo tests validate retained callback identity, mutation, invocation context,
 and ownership against Python behavior, including existing LiteLLM components.
 They do not wire retained callbacks into production routes or change provider
 preparation, authentication, HTTP transport, or response transformation

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

@@ -32,6 +32,8 @@ pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10;
 
 pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600;
 
+pub(crate) const BUFFERED_POST_CONNECT_TIMEOUT_SECS: u64 = 10;
+
 /// `object` field every non-streaming chat completion response carries.
 pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
 

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

@@ -1,5 +1,7 @@
 //! Header and upstream-body helpers shared by every route module.
 
+pub mod buffered_post;
+
 use serde_json::{Map, Value};
 
 use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;

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

@@ -0,0 +1,74 @@
+use std::sync::OnceLock;
+use std::time::Duration;
+
+use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
+
+use crate::constants::BUFFERED_POST_CONNECT_TIMEOUT_SECS;
+use crate::error::Error;
+
+pub struct Request {
+    pub url: String,
+    pub headers: Vec<(Vec<u8>, Vec<u8>)>,
+    pub body: Vec<u8>,
+    pub timeout_seconds: f64,
+}
+
+pub struct Response {
+    pub status: u16,
+    pub headers: Vec<(Vec<u8>, Vec<u8>)>,
+    pub content: Vec<u8>,
+}
+
+pub async fn send(request: Request) -> Result<Response, Error> {
+    let timeout = Duration::try_from_secs_f64(request.timeout_seconds)
+        .ok()
+        .filter(|timeout| !timeout.is_zero())
+        .ok_or_else(|| Error::InvalidRequest("timeout must be positive and finite".into()))?;
+    let mut headers = HeaderMap::new();
+    for (name, value) in request.headers {
+        let name = HeaderName::from_bytes(&name)
+            .map_err(|_| Error::InvalidRequest("invalid header name".into()))?;
+        let value = HeaderValue::from_bytes(&value)
+            .map_err(|_| Error::InvalidRequest("invalid header value".into()))?;
+        headers.append(name, value);
+    }
+
+    static CLIENT: OnceLock<Result<reqwest::Client, reqwest::Error>> = OnceLock::new();
+    let client = CLIENT
+        .get_or_init(|| {
+            reqwest::Client::builder()
+                .connect_timeout(Duration::from_secs(BUFFERED_POST_CONNECT_TIMEOUT_SECS))
+                .redirect(reqwest::redirect::Policy::none())
+                .no_gzip()
+                .no_brotli()
+                .no_deflate()
+                .no_zstd()
+                .build()
+        })
+        .as_ref()
+        .map_err(|_| Error::Network("could not initialize HTTP client".into()))?;
+    let response = client
+        .post(request.url)
+        .headers(headers)
+        .body(request.body)
+        .timeout(timeout)
+        .send()
+        .await
+        .map_err(|_| Error::Network("transport failed".into()))?;
+    let status = response.status().as_u16();
+    let headers = response
+        .headers()
+        .iter()
+        .map(|(name, value)| (name.as_str().as_bytes().to_vec(), value.as_bytes().to_vec()))
+        .collect();
+    let content = response
+        .bytes()
+        .await
+        .map_err(|_| Error::Network("could not read response".into()))?
+        .to_vec();
+    Ok(Response {
+        status,
+        headers,
+        content,
+    })
+}

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

@@ -0,0 +1,45 @@
+use litellm_core::error::Error;
+use litellm_core::http_utils::buffered_post::{Request, send};
+
+fn request() -> Request {
+    Request {
+        url: "unknown://private-document?secret=credential".into(),
+        headers: vec![],
+        body: vec![0, 255],
+        timeout_seconds: 1.0,
+    }
+}
+
+#[tokio::test]
+async fn rejects_invalid_timeouts_and_headers_without_echoing_wire_values() {
+    for timeout_seconds in [0.0, -1.0, f64::NAN, f64::INFINITY, f64::MAX] {
+        let result = send(Request {
+            timeout_seconds,
+            ..request()
+        })
+        .await;
+        assert!(
+            matches!(result, Err(Error::InvalidRequest(message)) if message == "timeout must be positive and finite")
+        );
+    }
+    for (headers, expected) in [
+        (
+            vec![(b"private\nname".to_vec(), b"secret".to_vec())],
+            "invalid header name",
+        ),
+        (
+            vec![(b"x-proof".to_vec(), b"private\nvalue".to_vec())],
+            "invalid header value",
+        ),
+    ] {
+        let result = send(Request {
+            headers,
+            ..request()
+        })
+        .await;
+        assert!(matches!(result, Err(Error::InvalidRequest(message)) if message == expected));
+    }
+    assert!(
+        matches!(send(request()).await, Err(Error::Network(message)) if message == "transport failed")
+    );
+}

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

@@ -59,41 +59,3 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add("RustBridgeDeclined", py.get_type::<RustBridgeDeclined>())?;
     module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())
 }
-
-pub(crate) fn ocr_error_to_pyerr(err: Error) -> PyErr {
-    match err {
-        Error::MissingField("document_url" | "image_url") => {
-            PyValueError::new_err("Document URL is required")
-        }
-        Error::Http { status, body } => RustUpstreamError::new_err((status, body)),
-        other => core_error_to_pyerr(other),
-    }
-}
-
-#[cfg(test)]
-mod ocr_error_tests {
-    use super::*;
-
-    #[test]
-    fn ocr_errors_preserve_python_validation_and_provider_details() {
-        Python::initialize();
-        Python::attach(|py| {
-            for field in ["document_url", "image_url"] {
-                let mapped = ocr_error_to_pyerr(Error::MissingField(field));
-                assert!(mapped.is_instance_of::<PyValueError>(py));
-                assert_eq!(mapped.value(py).to_string(), "Document URL is required");
-            }
-            let mapped = ocr_error_to_pyerr(Error::Http {
-                status: 429,
-                body: r#"{"message":"rate limited"}"#.to_string(),
-            });
-            assert!(mapped.is_instance_of::<RustUpstreamError>(py));
-            let args: (u16, String) = mapped
-                .value(py)
-                .getattr("args")
-                .and_then(|args| args.extract())
-                .expect("OCR failures retain status and unprefixed provider message");
-            assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string()));
-        });
-    }
-}

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

@@ -20,7 +20,20 @@ where
     T: Serialize + Send + 'static,
     F: Future<Output = Result<T, Error>> + Send + 'static,
 {
-    run_sync_on(
+    let result = run_sync_value(py, future, map_error)?;
+    Pythonized(result).into_pyobject(py).map(Bound::unbind)
+}
+
+pub(crate) fn run_sync_value<T, F>(
+    py: Python<'_>,
+    future: F,
+    map_error: fn(Error) -> PyErr,
+) -> PyResult<T>
+where
+    T: Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    run_sync_value_on(
         py,
         pyo3_async_runtimes::tokio::get_runtime(),
         future,
@@ -28,14 +41,14 @@ where
     )
 }
 
-fn run_sync_on<T, F>(
+fn run_sync_value_on<T, F>(
     py: Python<'_>,
     runtime: &Runtime,
     future: F,
     map_error: fn(Error) -> PyErr,
-) -> PyResult<Py<PyAny>>
+) -> PyResult<T>
 where
-    T: Serialize + Send + 'static,
+    T: Send + 'static,
     F: Future<Output = Result<T, Error>> + Send + 'static,
 {
     if Handle::try_current().is_ok() {
@@ -45,8 +58,7 @@ where
     }
 
     let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
-    let result = map_core_result(result, map_error)?;
-    Pythonized(result).into_pyobject(py).map(Bound::unbind)
+    map_core_result(result, map_error)
 }
 
 pub(crate) fn run_async<T, F>(
@@ -65,6 +77,15 @@ where
     })
 }
 
+pub(crate) async fn run_async_value<T, F>(future: F, map_error: fn(Error) -> PyErr) -> PyResult<T>
+where
+    T: Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    let result = catch_future_panic(future).await?;
+    map_core_result(result, map_error)
+}
+
 fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
     match result {
         Ok(value) => Ok(value),
@@ -256,7 +277,7 @@ mod tests {
             .build()
             .expect("runtime should build");
         Python::attach(|py| {
-            let result = run_sync_on(
+            let result = run_sync_value_on(
                 py,
                 &runtime,
                 async {
@@ -265,7 +286,7 @@ mod tests {
                 },
                 runtime_error,
             );
-            assert!(extract_bool(py, result));
+            assert!(result.expect("route should complete"));
         });
     }
 

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

@@ -222,11 +222,7 @@ mod tests {
             let module = PyModule::new(py, "routes").expect("module should be created");
             crate::routes::register(&module).expect("routes should register");
             let routes = [
-                (
-                    "ocr",
-                    "aocr",
-                    "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
-                ),
+                ("ocr", "aocr", "(boundary)"),
                 (
                     "transcription",
                     "atranscription",
@@ -311,7 +307,7 @@ mod tests {
                 .expect("kwargs should accept extra_headers");
             let document = PyDict::new(py);
 
-            for (sync_name, async_name) in [("ocr", "aocr"), ("transcription", "atranscription")] {
+            for (sync_name, async_name) in [("transcription", "atranscription")] {
                 let sync_error = module
                     .getattr(sync_name)
                     .and_then(|function| function.call(("model", &document), Some(&kwargs)))

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

@@ -1,73 +1,181 @@
-use litellm_core::Error;
-use std::future::Future;
+//! Retained OCR route: Python owns request/response objects, Rust sequences
+//! prepare -> encode -> POST -> finish through owning `Py` handles.
 
-use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
+use litellm_core::http_utils::buffered_post::{self, Request, Response};
+use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
 use pyo3::prelude::*;
-use serde_json::Value;
+use pyo3::sync::PyOnceLock;
+use pyo3::types::{PyBytes, PyList, PyTuple};
 
-use crate::errors::ocr_error_to_pyerr;
-use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
+use crate::errors::core_error_to_pyerr;
+use crate::execution::{run_async_value, run_sync_value};
 
-fn prepare_ocr(
-    inputs: OcrInputs,
-) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
-    let document = inputs.document;
-    let options = RouteOptions::from_python(RouteOptionsInputs {
-        model: inputs.model,
-        api_key: inputs.api_key,
-        api_base: inputs.api_base,
-        custom_llm_provider: inputs.custom_llm_provider,
-        extra_headers: inputs.extra_headers,
-        timeout_seconds: inputs.timeout_seconds,
-    })?;
-    let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
+#[derive(Clone, Copy)]
+struct BoundaryStep {
+    method: &'static str,
+    awaited: bool,
+}
 
-    Ok(async move {
-        let RouteOptions {
-            model,
-            api_key,
-            api_base,
-            custom_llm_provider,
-            extra_headers,
-            timeout,
-        } = options;
-        run_ocr(OcrRequest {
-            model: &model,
-            document,
-            api_key: api_key.as_deref(),
-            api_base: api_base.as_deref(),
-            custom_llm_provider: custom_llm_provider.as_deref(),
-            extra_headers,
-            optional_params,
-            timeout,
-            callbacks: Vec::new(),
-            guardrails: Vec::new(),
-            request_metadata: Default::default(),
-            litellm_call_id: None,
-        })
-        .await
+const PREPARE_SYNC: BoundaryStep = BoundaryStep {
+    method: "prepare",
+    awaited: false,
+};
+const PREPARE_ASYNC: BoundaryStep = BoundaryStep {
+    method: "aprepare",
+    awaited: true,
+};
+const ENCODE: BoundaryStep = BoundaryStep {
+    method: "encode",
+    awaited: false,
+};
+const FINISH_SYNC: BoundaryStep = BoundaryStep {
+    method: "finish",
+    awaited: false,
+};
+const FINISH_ASYNC: BoundaryStep = BoundaryStep {
+    method: "afinish",
+    awaited: true,
+};
+
+fn invoke(
+    boundary: &Bound<'_, PyAny>,
+    step: BoundaryStep,
+    args: Bound<'_, PyTuple>,
+) -> PyResult<Py<PyAny>> {
+    let call = PreparedCall::new(
+        if step.awaited {
+            InvocationMode::Await
+        } else {
+            InvocationMode::Direct
+        },
+        boundary.getattr(step.method)?.unbind(),
+        args.unbind(),
+        None,
+    );
+    match call.invoke(boundary.py())? {
+        InvocationOutcome::Returned(value) | InvocationOutcome::Awaitable(value) => Ok(value),
+    }
+}
+
+#[pyfunction]
+fn prepare(boundary: &Bound<'_, PyAny>, asynchronous: bool) -> PyResult<Py<PyAny>> {
+    let step = if asynchronous {
+        PREPARE_ASYNC
+    } else {
+        PREPARE_SYNC
+    };
+    invoke(boundary, step, PyTuple::empty(boundary.py()))
+}
+
+fn request(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult<Request> {
+    type ByteHeaders<'py> = Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>;
+    let py = boundary.py();
+    let encoded = invoke(boundary, ENCODE, PyTuple::new(py, [roots])?)?;
+    let (url, headers, body, timeout_seconds): (String, ByteHeaders<'_>, Bound<'_, PyBytes>, f64) =
+        encoded.into_bound(py).extract()?;
+    Ok(Request {
+        url,
+        headers: headers
+            .into_iter()
+            .map(|(name, value)| (name.as_bytes().to_vec(), value.as_bytes().to_vec()))
+            .collect(),
+        body: body.as_bytes().to_vec(),
+        timeout_seconds,
     })
 }
 
-bridge_route! {
-    sync = ocr,
-    asynchronous = aocr,
-    inputs = OcrInputs,
-    required = {
-        model: String,
-        #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        document: serde_json::Value,
-    },
-    optional = {
-        api_key: Option<String>,
-        api_base: Option<String>,
-        custom_llm_provider: Option<String>,
-        #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        extra_headers: Option<serde_json::Value>,
-        #[pyo3(from_py_with = litellm_python_interop::from_py)]
-        optional_params: Option<serde_json::Value>,
-        timeout_seconds: Option<f64>,
-    },
-    prepare = prepare_ocr,
-    errors = ocr_error_to_pyerr,
+struct Wire(Response);
+
+impl<'py> IntoPyObject<'py> for Wire {
+    type Target = PyTuple;
+    type Output = Bound<'py, PyTuple>;
+    type Error = PyErr;
+
+    fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
+        let headers = PyList::new(
+            py,
+            self.0
+                .headers
+                .iter()
+                .map(|(name, value)| (PyBytes::new(py, name), PyBytes::new(py, value))),
+        )?;
+        (self.0.status, headers, PyBytes::new(py, &self.0.content)).into_pyobject(py)
+    }
+}
+
+#[pyfunction]
+fn send<'a>(
+    boundary: &'a Bound<'a, PyAny>,
+    roots: &Bound<'_, PyAny>,
+) -> PyResult<Bound<'a, PyAny>> {
+    let request = request(boundary, roots)?;
+    pyo3_async_runtimes::tokio::future_into_py(boundary.py(), async move {
+        let response = run_async_value(buffered_post::send(request), core_error_to_pyerr).await?;
+        Ok(Wire(response))
+    })
+}
+
+#[pyfunction]
+fn finish(
+    boundary: &Bound<'_, PyAny>,
+    wire: &Bound<'_, PyAny>,
+    asynchronous: bool,
+) -> PyResult<Py<PyAny>> {
+    let step = if asynchronous {
+        FINISH_ASYNC
+    } else {
+        FINISH_SYNC
+    };
+    invoke(boundary, step, PyTuple::new(boundary.py(), [wire])?)
+}
+
+#[pyfunction]
+fn ocr(boundary: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
+    let py = boundary.py();
+    let roots = prepare(boundary, false)?;
+    let request = request(boundary, roots.bind(py))?;
+    let response = run_sync_value(py, buffered_post::send(request), core_error_to_pyerr)?;
+    let wire = Wire(response).into_pyobject(py)?;
+    finish(boundary, &wire, false)
+}
+
+#[pyfunction]
+fn aocr<'a>(boundary: &'a Bound<'a, PyAny>) -> PyResult<Bound<'a, PyAny>> {
+    driver(boundary.py())?.getattr("drive")?.call1((boundary,))
+}
+
+/// The async route must await `aprepare`/`afinish` inline in the caller's
+/// Python task, so a Python driver coroutine owns the roots between steps.
+/// Compiling the driver runs Python (audit hooks can re-enter `aocr`), so
+/// compile first and publish only a finished module into the once-lock.
+fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
+    static DRIVER: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
+    if let Some(module) = DRIVER.get(py) {
+        return Ok(module.bind(py));
+    }
+    let module = PyModule::from_code(
+        py,
+        c"async def drive(boundary):
+    roots = await _prepare(boundary, True)
+    wire = await _send(boundary, roots)
+    return await _finish(boundary, wire, True)
+",
+        c"ocr_driver.py",
+        c"_ocr_driver",
+    )?;
+    module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?;
+    module.add("_send", wrap_pyfunction!(send, &module)?)?;
+    module.add("_finish", wrap_pyfunction!(finish, &module)?)?;
+    Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py))
+}
+
+pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!(ocr, module)?)?;
+    crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!(aocr, module)?)?;
+    Ok(())
+}
+
+#[cfg(feature = "trace-parity")]
+pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    register(module)
 }

litellm-rust/crates/python-interop/tests/fixtures/callback_components.py

@@ -35,6 +35,8 @@ def logger_for(callbacks=(), stream=False, input_callbacks=(), sync_callbacks=()
 async def real_pre_call_logging(owners):
     retained = []
     observed = []
+    snapshots = []
+    order = []
     ignored = {"replacement": True}
     metadata = {"secret": "private", "keep": []}
     removed = object()
@@ -42,11 +44,13 @@ async def real_pre_call_logging(owners):
 
     class Retain(CustomLogger):
         def log_pre_api_call(self, model, messages, kwargs):
+            order.append("retain")
             retained.append(kwargs)
             return ignored
 
     class Mutate(CustomLogger):
         def log_pre_api_call(self, model, messages, kwargs):
+            order.append("mutate")
             kwargs["normalized"] = "normalized"
             assert kwargs.pop("remove") is removed
             kwargs["retained_metadata"]["secret"] = "masked"
@@ -54,12 +58,24 @@ async def real_pre_call_logging(owners):
 
     class Fail(CustomLogger):
         def log_pre_api_call(self, model, messages, kwargs):
+            order.append("fail")
             kwargs["lock"] = lock
             kwargs["retained_metadata"]["keep"].append("before failure")
             raise RuntimeError("expected pre-call callback failure")
 
     class Observe(CustomLogger):
         def log_pre_api_call(self, model, messages, kwargs):
+            order.append("observe")
+            snapshots.append(
+                (
+                    kwargs.get("normalized"),
+                    "remove" in kwargs,
+                    kwargs["retained_metadata"]["secret"],
+                    tuple(kwargs["retained_metadata"]["keep"]),
+                    "lock" in kwargs,
+                    "replacement" in kwargs,
+                )
+            )
             observed.append((kwargs, messages))
 
     logger = logger_for(input_callbacks=[Retain(), Mutate(), Fail(), Observe()])
@@ -72,7 +88,9 @@ async def real_pre_call_logging(owners):
         assert owner.invoke() is None
     finally:
         owner.close()
-    assert retained == [details] and observed == [(details, messages)]
+    assert order == ["retain", "mutate", "fail", "observe"]
+    assert snapshots == [("normalized", False, "masked", ("before failure",), True, False)]
+    assert len(retained) == len(observed) == 1
     assert retained[0] is details and observed[0][0] is details
     assert observed[0][1] is messages and details["input"] is messages
     assert details["additional_args"] is additional

litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py

@@ -661,15 +661,19 @@ async def detached_work_after_error(owners):
 
 def run_checked(owners, scenario):
     baseline = owners.live
+    background_failures = []
 
     async def run():
+        asyncio.get_running_loop().set_exception_handler(lambda loop, context: background_failures.append(context))
         await asyncio.wait_for(scenario, timeout=15)
+        gc.collect()
         assert owners.live == baseline
         pending = asyncio.all_tasks() - {asyncio.current_task()}
         assert not pending, f"undrained tasks: {pending}"
 
     asyncio.run(run())
     gc.collect()
+    assert not background_failures, f"unhandled background failures: {background_failures}"
     assert owners.live == baseline
 
 

litellm-rust/crates/python-interop/tests/prepared_call.rs

@@ -6,7 +6,7 @@ use rstest::rstest;
 #[path = "support/mod.rs"]
 mod support;
 
-use support::python::{InitializedPython, initialized_python, item, scope};
+use support::python::{InitializedPython, initialized_python, item, run_fixture, scope};
 
 #[rstest]
 fn retains_aliases_mutations_and_original_result(
@@ -271,6 +271,55 @@ fn prepare_pre_call(
     ))
 }
 
+#[rstest]
+fn checked_runner_rejects_unhandled_background_failures(
+    initialized_python: &InitializedPython,
+) -> PyResult<()> {
+    let _ = initialized_python;
+    Python::attach(|py| {
+        let globals = PyDict::new(py);
+        run_fixture(
+            py,
+            &globals,
+            include_str!("fixtures/callback_lifecycle.py"),
+            concat!(
+                env!("CARGO_MANIFEST_DIR"),
+                "/tests/fixtures/callback_lifecycle.py"
+            ),
+        )?;
+        py.run(
+            c"
+async def fail():
+    raise RuntimeError('background task regression')
+
+for cyclic in (False, True):
+    for handled in (False, True):
+        async def scenario(cyclic=cyclic, handled=handled):
+            task = asyncio.create_task(fail())
+            if cyclic:
+                task.cycle = task
+            await checkpoint()
+            assert task.done()
+            if handled:
+                with TestCase().assertRaisesRegex(RuntimeError, 'background task regression'):
+                    task.result()
+            del task
+
+        owners = ReferenceFactory()
+        if handled:
+            run_checked(owners, scenario())
+        else:
+            with TestCase().assertRaisesRegex(
+                AssertionError, r'unhandled background failures: .*background task regression'
+            ):
+                run_checked(owners, scenario())
+",
+            Some(&globals),
+            None,
+        )
+    })
+}
+
 #[rstest]
 #[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"]
 fn real_ocr_logging_preserves_execution_roots_and_continues_after_error(
@@ -285,15 +334,20 @@ from datetime import datetime
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.litellm_logging import Logging
 
+order = []
+
 class Retain(CustomLogger):
     def log_pre_api_call(self, model, messages, kwargs):
+        order.append('retain')
         self.view = kwargs['additional_args']
         self.headers = self.view['headers']
         self.body = self.view['complete_input_dict']
+        self.snapshot = (self.headers['X-Trace'], self.body['document']['value'])
         return {'ignored_replacement': True}
 
 class MutateThenFail(CustomLogger):
     def log_pre_api_call(self, model, messages, kwargs):
+        order.append('mutate_then_fail')
         view = kwargs['additional_args']
         view['headers']['X-Trace'] = 'mutated'
         view['complete_input_dict']['document']['value'] = 'mutated'
@@ -303,7 +357,13 @@ class MutateThenFail(CustomLogger):
 
 class Observe(CustomLogger):
     def log_pre_api_call(self, model, messages, kwargs):
+        order.append('observe')
         self.view = kwargs['additional_args']
+        self.snapshot = (
+            tuple(sorted(self.view['headers'].items())),
+            self.view['complete_input_dict'].get('replacement'),
+            'document' in self.view['complete_input_dict'],
+        )
 
 first = Retain()
 last = Observe()
@@ -338,6 +398,9 @@ logger = Logging(
         );
         py.run(
             c"
+assert order == ['retain', 'mutate_then_fail', 'observe']
+assert first.snapshot == ('original', 'original')
+assert last.snapshot == ((('X-Trace', 'replacement'),), True, False)
 assert first.view is last.view
 assert first.body['document'] is document
 assert first.body['alias'] is document

litellm/ocr/main.py

@@ -7,7 +7,7 @@ import base64
 import mimetypes
 import os
 import re
-from collections.abc import Callable, Coroutine, Mapping
+from collections.abc import Coroutine, Mapping
 from dataclasses import dataclass
 from io import IOBase
 from typing import Any, Final, cast
@@ -29,7 +29,6 @@ from litellm.llms.base_llm.ocr.transformation import (
 )
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
 from litellm.rust_bridge import ocr as rust_ocr_bridge
-from litellm.rust_bridge.bindings import native_exception_types
 from litellm.rust_bridge.configuration import rust_enabled
 from litellm.types.router import GenericLiteLLMParams
 from litellm.utils import ProviderConfigManager, client
@@ -54,14 +53,6 @@ class _PreparedOCRRequest:
     litellm_logging_obj: LiteLLMLoggingObj
 
 
-@dataclass
-class _PreparedRustOCRCall:
-    api_key: str | None
-    api_base: str | None
-    headers: dict[str, object]
-    optional_params: dict[str, object]
-
-
 _RUST_OCR_PROVIDERS: Final = {
     "mistral",
     "azure_ai",
@@ -198,169 +189,39 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
     return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
 
 
-def _rust_bridge_optional_params(
-    prepared_request: _PreparedOCRRequest,
-    resolve_secret: Callable[[str], str | None],
-) -> dict[str, object]:
-    optional_params: Final = dict(prepared_request.optional_params)
-    if prepared_request.custom_llm_provider == "vertex_ai":
-        vertex_project: Final = (
-            prepared_request.litellm_params.get("vertex_project")
-            or prepared_request.litellm_params.get("vertex_ai_project")
-            or litellm.vertex_project
-            or resolve_secret("VERTEXAI_PROJECT")
-        )
-        vertex_location: Final = (
-            prepared_request.litellm_params.get("vertex_location")
-            or prepared_request.litellm_params.get("vertex_ai_location")
-            or litellm.vertex_location
-            or resolve_secret("VERTEXAI_LOCATION")
-            or resolve_secret("VERTEX_LOCATION")
-        )
-        if vertex_project is not None:
-            optional_params["vertex_project"] = vertex_project
-        if vertex_location is not None:
-            optional_params["vertex_location"] = vertex_location
-    return optional_params
+def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool:
+    raw_request_override: Final = prepared_request.litellm_params.get("rust")
+    request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
+    return rust_enabled(request_override=request_override)
 
 
-def _rust_bridge_api_base(
-    prepared_request: _PreparedOCRRequest,
-    resolve_secret: Callable[[str], str | None],
-) -> str | None:
-    if prepared_request.api_base is not None:
-        return prepared_request.api_base
-    if prepared_request.custom_llm_provider == "azure_ai":
-        if is_azure_document_intelligence_model(prepared_request.model):
-            return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
-        return resolve_secret("AZURE_AI_API_BASE")
-    return None
-
-
-def _prepare_rust_ocr_call(
-    prepared_request: _PreparedOCRRequest,
-    resolve_api_key: Callable[[str], str | None],
-) -> _PreparedRustOCRCall:
-    provider_config: Final = prepared_request.provider_config
-    api_key_env_var: Final = provider_config.get_api_key_env_var()
-    resolved_api_key: Final = prepared_request.api_key or (
-        resolve_api_key(api_key_env_var) if api_key_env_var is not None else None
-    )
-    resolved_headers: Final = provider_config.validate_environment(
-        headers=prepared_request.extra_headers or {},
-        model=prepared_request.model,
-        api_key=resolved_api_key,
-        api_base=prepared_request.api_base,
-        litellm_params=prepared_request.litellm_params,
-    )
-    resolved_complete_url: Final = provider_config.get_complete_url(
-        api_base=prepared_request.api_base,
+def _ocr_boundary(prepared_request: _PreparedOCRRequest) -> rust_ocr_bridge.OCRBoundary:
+    return rust_ocr_bridge.OCRBoundary(
+        handler=base_llm_http_handler,
         model=prepared_request.model,
+        document=prepared_request.document,
         optional_params=prepared_request.optional_params,
+        logging_obj=prepared_request.litellm_logging_obj,
+        api_key=prepared_request.api_key,
+        api_base=prepared_request.api_base,
+        headers=prepared_request.extra_headers,
+        provider_config=prepared_request.provider_config,
         litellm_params=prepared_request.litellm_params,
-    )
-    rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key)
-    rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key)
-    prepared_request.litellm_logging_obj.pre_call(
-        input="OCR document processing",
-        api_key=resolved_api_key,
-        additional_args={
-            "complete_input_dict": {
-                "model": prepared_request.model,
-                "document": prepared_request.document,
-                **rust_optional_params,
-            },
-            "api_base": resolved_complete_url,
-            "headers": resolved_headers,
-        },
-    )
-    return _PreparedRustOCRCall(
-        api_key=resolved_api_key,
-        api_base=rust_api_base,
-        headers=cast(dict[str, object], resolved_headers),
-        optional_params=rust_optional_params,
+        custom_llm_provider=prepared_request.custom_llm_provider,
+        timeout=prepared_request.effective_timeout,
     )
 
 
-def _map_rust_ocr_error(
-    error: Exception,
-    prepared_request: _PreparedOCRRequest,
-    exception_types: tuple[type[BaseException], type[BaseException]] | None,
-) -> Exception:
-    if exception_types is None:
-        return error
-    _, upstream_error = exception_types
-    if not isinstance(error, upstream_error):
-        return error
-    error_args: Final = cast(  # cast-ok: BaseException.args is typed with Any in the standard library stubs
-        tuple[object, ...], error.args
-    )
-    status_value: Final = error_args[0] if error_args else 0
-    message_value: Final = error_args[1] if len(error_args) > 1 else str(error)
-    status: Final = status_value if isinstance(status_value, int) else 0
-    message: Final = message_value if isinstance(message_value, str) else str(message_value)
-    error_factory: Final = cast(  # cast-ok: the legacy provider interface leaves callable parameters untyped
-        Callable[..., Exception], prepared_request.provider_config.get_error_class
-    )
-    return error_factory(
-        error_message=message,
-        status_code=status or 500,
-        headers={},  # mutable-ok: provider error factories require a concrete header dict
-    )
-def _run_rust_ocr(
-    prepared_request: _PreparedOCRRequest,
-    resolve_api_key: Callable[[str], str | None],
-) -> OCRResponse | None:
+def _run_rust_ocr(prepared_request: _PreparedOCRRequest) -> OCRResponse | None:
     if rust_ocr_bridge.load_rust_ocr() is None:
         return None
-    prepared: Final = _prepare_rust_ocr_call(
-        prepared_request=prepared_request,
-        resolve_api_key=resolve_api_key,
-    )
-    try:
-        rust_response: Final = rust_ocr_bridge.ocr(
-            model=prepared_request.model,
-            document=prepared_request.document,
-            api_key=prepared.api_key,
-            api_base=prepared.api_base,
-            custom_llm_provider=prepared_request.custom_llm_provider,
-            extra_headers=prepared.headers,
-            optional_params=prepared.optional_params,
-            timeout=prepared_request.effective_timeout,
-        )
-    except Exception as error:
-        raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error
-    if rust_response is None:
-        return None
-    return OCRResponse.model_validate(rust_response)
+    return rust_ocr_bridge.ocr(_ocr_boundary(prepared_request))
 
 
-async def _run_rust_aocr(
-    prepared_request: _PreparedOCRRequest,
-    resolve_api_key: Callable[[str], str | None],
-) -> OCRResponse | None:
+async def _run_rust_aocr(prepared_request: _PreparedOCRRequest) -> OCRResponse | None:
     if rust_ocr_bridge.load_rust_aocr() is None:
         return None
-    prepared: Final = _prepare_rust_ocr_call(
-        prepared_request=prepared_request,
-        resolve_api_key=resolve_api_key,
-    )
-    try:
-        rust_response: Final = await rust_ocr_bridge.aocr(
-            model=prepared_request.model,
-            document=prepared_request.document,
-            api_key=prepared.api_key,
-            api_base=prepared.api_base,
-            custom_llm_provider=prepared_request.custom_llm_provider,
-            extra_headers=prepared.headers,
-            optional_params=prepared.optional_params,
-            timeout=prepared_request.effective_timeout,
-        )
-    except Exception as error:
-        raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error
-    if rust_response is None:
-        return None
-    return OCRResponse.model_validate(rust_response)
+    return await rust_ocr_bridge.aocr(_ocr_boundary(prepared_request))
 
 
 @client
@@ -457,13 +318,8 @@ async def aocr(
         custom_llm_provider = prepared.custom_llm_provider
         completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
 
-        if _rust_ocr_supported(prepared) and rust_enabled():
-            from litellm.secret_managers.main import get_secret_str
-
-            rust_response: Final = await _run_rust_aocr(
-                prepared_request=prepared,
-                resolve_api_key=get_secret_str,
-            )
+        if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
+            rust_response: Final = await _run_rust_aocr(prepared_request=prepared)
             if rust_response is None:
                 verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path")
             else:
@@ -729,13 +585,8 @@ def ocr(
         custom_llm_provider = prepared.custom_llm_provider
         completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
 
-        if _rust_ocr_supported(prepared) and rust_enabled():
-            from litellm.secret_managers.main import get_secret_str
-
-            rust_response: Final = _run_rust_ocr(
-                prepared_request=prepared,
-                resolve_api_key=get_secret_str,
-            )
+        if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
+            rust_response: Final = _run_rust_ocr(prepared_request=prepared)
             if rust_response is None:
                 verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path")
             else:

litellm/rust_bridge/ocr.py

@@ -1,44 +1,145 @@
-"""Thin Python wrapper for the native Rust OCR bridge."""
+"""Retained OCR bridge: Python owns the request/response objects, Rust drives
+prepare -> encode -> POST -> finish against those same objects."""
 
 from __future__ import annotations
 
+import math
 from collections.abc import Awaitable
+from dataclasses import dataclass, field
 from typing import Final, Protocol, cast  # noqa: TID251  # native extension exposes dynamically typed callables
 
 import httpx
 
+import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, DocumentType, OCRResponse
+from litellm.llms.custom_httpx.http_handler import (
+    AsyncHTTPHandler,
+    HTTPHandler,
+    _get_httpx_client,
+    get_async_httpx_client,
+)
+from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
+from litellm.rust_bridge import configuration as _configuration
 from litellm.rust_bridge.bindings import NativeBinding
 from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds
 
+rust_ocr_enabled = _configuration.rust_enabled
+rust = _configuration.rust
+
+OCRRoots = tuple[dict[str, object], str, dict[str, object], None]
+OCRWire = tuple[int, list[tuple[bytes, bytes]], bytes]
+OCREncoded = tuple[str, list[tuple[bytes, bytes]], bytes, float]
+
+
+def _positive_timeout_seconds(timeout: float | httpx.Timeout) -> float:
+    seconds: Final = _timeout_to_seconds(timeout)
+    if seconds is None or not math.isfinite(seconds) or seconds <= 0:
+        raise ValueError("OCR bridge requires a positive finite timeout")
+    return seconds
+
+
+@dataclass(kw_only=True, slots=True)
+class OCRBoundary:
+    handler: BaseLLMHTTPHandler
+    model: str
+    document: DocumentType
+    optional_params: dict[str, object]
+    logging_obj: Logging
+    api_key: str | None
+    api_base: str | None
+    headers: dict[str, object] | None
+    provider_config: BaseOCRConfig
+    litellm_params: dict[str, object]
+    custom_llm_provider: str
+    timeout: float | httpx.Timeout
+    client: HTTPHandler | AsyncHTTPHandler | None = None
+    request: httpx.Request | None = field(default=None, init=False)
+
+    def prepare(self) -> OCRRoots:
+        roots: Final = self.handler._prepare_ocr_request(
+            model=self.model,
+            document=self.document,
+            optional_params=self.optional_params,
+            logging_obj=self.logging_obj,
+            api_key=self.api_key,
+            api_base=self.api_base,
+            headers=self.headers,
+            provider_config=self.provider_config,
+            litellm_params=self.litellm_params,
+        )
+        if not isinstance(self.client, HTTPHandler):
+            self.client = _get_httpx_client()
+        return roots
+
+    async def aprepare(self) -> OCRRoots:
+        roots: Final = await self.handler._async_prepare_ocr_request(
+            model=self.model,
+            document=self.document,
+            optional_params=self.optional_params,
+            logging_obj=self.logging_obj,
+            api_key=self.api_key,
+            api_base=self.api_base,
+            headers=self.headers,
+            provider_config=self.provider_config,
+            litellm_params=self.litellm_params,
+        )
+        if not isinstance(self.client, AsyncHTTPHandler):
+            self.client = get_async_httpx_client(llm_provider=litellm.LlmProviders(self.custom_llm_provider))
+        return roots
+
+    def encode(self, roots: OCRRoots) -> OCREncoded:
+        headers, url, data, _files = roots
+        seconds: Final = _positive_timeout_seconds(self.timeout)
+        if self.client is None:
+            raise RuntimeError("OCR boundary must be prepared before encoding")
+        try:
+            self.request = self.client.client.build_request(
+                "POST",
+                url,
+                headers=cast(dict[str, str], headers),
+                json=data,
+                timeout=self.timeout,
+            )
+            return str(self.request.url), self.request.headers.raw, self.request.read(), seconds
+        except Exception as e:  # noqa: BLE001  # match the Python OCR handler's encoding error mapping
+            raise self.handler._handle_error(e=e, provider_config=self.provider_config)
+
+    def _response(self, wire: OCRWire) -> httpx.Response:
+        if self.request is None:
+            raise RuntimeError("OCR boundary must be encoded before finishing")
+        status, headers, content = wire
+        try:
+            response: Final = httpx.Response(status, headers=headers, content=content, request=self.request)
+            response.raise_for_status()
+        except Exception as e:  # noqa: BLE001  # match the Python OCR handler's response error mapping
+            raise self.handler._handle_error(e=e, provider_config=self.provider_config)
+        return response
+
+    def finish(self, wire: OCRWire) -> OCRResponse:
+        return self.handler._transform_ocr_response(
+            provider_config=self.provider_config,
+            model=self.model,
+            response=self._response(wire),
+            logging_obj=self.logging_obj,
+            optional_params=self.optional_params,
+        )
+
+    async def afinish(self, wire: OCRWire) -> OCRResponse:
+        return await self.provider_config.async_transform_ocr_response(
+            model=self.model,
+            raw_response=self._response(wire),
+            logging_obj=self.logging_obj,
+            optional_params=self.optional_params,
+        )
+
 
 class RustOcr(Protocol):
-    def __call__(
-        self,
-        model: str,
-        document: dict[str, object],
-        api_key: str | None,
-        api_base: str | None,
-        custom_llm_provider: str | None,
-        extra_headers: dict[str, object] | None,
-        optional_params: dict[str, object],
-        timeout_seconds: float | None,
-    ) -> dict[str, object]:
-        raise NotImplementedError
+    def __call__(self, boundary: OCRBoundary) -> OCRResponse: ...
 
 
 class RustAocr(Protocol):
-    def __call__(
-        self,
-        model: str,
-        document: dict[str, object],
-        api_key: str | None,
-        api_base: str | None,
-        custom_llm_provider: str | None,
-        extra_headers: dict[str, object] | None,
-        optional_params: dict[str, object],
-        timeout_seconds: float | None,
-    ) -> Awaitable[dict[str, object]]:
-        raise NotImplementedError
+    def __call__(self, boundary: OCRBoundary) -> Awaitable[OCRResponse]: ...
 
 
 def _as_ocr(value: object) -> RustOcr | None:
@@ -61,53 +162,15 @@ def load_rust_aocr() -> RustAocr | None:
     return _AOCR.load()
 
 
-def ocr(
-    *,
-    model: str,
-    document: dict[str, object],
-    api_key: str | None,
-    api_base: str | None,
-    custom_llm_provider: str | None,
-    extra_headers: dict[str, object] | None,
-    optional_params: dict[str, object],
-    timeout: float | httpx.Timeout | None,
-) -> dict[str, object] | None:
+def ocr(boundary: OCRBoundary) -> OCRResponse | None:
     rust_ocr: Final = load_rust_ocr()
     if rust_ocr is None:
         return None
-    return rust_ocr(
-        model=model,
-        document=document,
-        api_key=api_key,
-        api_base=api_base,
-        custom_llm_provider=custom_llm_provider,
-        extra_headers=extra_headers,
-        optional_params=optional_params,
-        timeout_seconds=_timeout_to_seconds(timeout),
-    )
+    return rust_ocr(boundary)
 
 
-async def aocr(
-    *,
-    model: str,
-    document: dict[str, object],
-    api_key: str | None,
-    api_base: str | None,
-    custom_llm_provider: str | None,
-    extra_headers: dict[str, object] | None,
-    optional_params: dict[str, object],
-    timeout: float | httpx.Timeout | None,
-) -> dict[str, object] | None:
+async def aocr(boundary: OCRBoundary) -> OCRResponse | None:
     rust_aocr: Final = load_rust_aocr()
     if rust_aocr is None:
         return None
-    return await rust_aocr(
-        model=model,
-        document=document,
-        api_key=api_key,
-        api_base=api_base,
-        custom_llm_provider=custom_llm_provider,
-        extra_headers=extra_headers,
-        optional_params=optional_params,
-        timeout_seconds=_timeout_to_seconds(timeout),
-    )
+    return await rust_aocr(boundary)

tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py

@@ -11,7 +11,7 @@ supplied api_base is always honoured.
 from litellm.llms.azure_ai.ocr.common_utils import (
     is_azure_document_intelligence_model,
 )
-from litellm.ocr.main import _prepare_ocr_request, _rust_bridge_api_base
+from litellm.ocr.main import _prepare_ocr_request
 
 _DOC = {"type": "document_url", "document_url": "https://example.com/doc.pdf"}
 _DOC_INTELLIGENCE_ENDPOINT = "https://di.cognitiveservices.azure.com"
@@ -23,13 +23,6 @@ class _FakeLogging:
         return None
 
 
-def _resolve_secret(name: str) -> str | None:
-    return {
-        "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT": _DOC_INTELLIGENCE_ENDPOINT,
-        "AZURE_AI_API_BASE": _AZURE_AI_API_BASE,
-    }.get(name)
-
-
 def _prepare(model: str, api_base: str | None):
     return _prepare_ocr_request(
         model=model,
@@ -55,31 +48,59 @@ class TestIsAzureDocumentIntelligenceModel:
 
 
 class TestDocIntelligenceApiBaseResolution:
+    """Both the Python and boundary routes resolve endpoints through the provider
+    config's ``get_complete_url``; these pin that resolution."""
+
     def test_generic_azure_ai_base_does_not_hijack_doc_intelligence(self, monkeypatch):
         """Without an explicit api_base, the AZURE_AI_API_BASE fallback must not
         overwrite the endpoint, so it resolves to the Document Intelligence one."""
         monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE)
-        monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", raising=False)
+        monkeypatch.setenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", _DOC_INTELLIGENCE_ENDPOINT)
+        from litellm.llms.azure_ai.ocr.document_intelligence.transformation import (
+            AzureDocumentIntelligenceOCRConfig,
+        )
 
         prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", None)
 
-        assert prepared.api_base is None
-        assert _rust_bridge_api_base(prepared, _resolve_secret) == _DOC_INTELLIGENCE_ENDPOINT
+        url = AzureDocumentIntelligenceOCRConfig().get_complete_url(
+            api_base=prepared.api_base,
+            model=prepared.model,
+            optional_params=prepared.optional_params,
+            litellm_params=prepared.litellm_params,
+        )
+        assert url.startswith(_DOC_INTELLIGENCE_ENDPOINT)
 
     def test_explicit_api_base_is_honoured_for_doc_intelligence(self, monkeypatch):
         """A caller-supplied api_base must always win, even for doc-intelligence."""
         monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE)
+        from litellm.llms.azure_ai.ocr.document_intelligence.transformation import (
+            AzureDocumentIntelligenceOCRConfig,
+        )
 
         custom = "https://my-di.cognitiveservices.azure.com"
         prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", custom)
 
         assert prepared.api_base == custom
-        assert _rust_bridge_api_base(prepared, _resolve_secret) == custom
+        url = AzureDocumentIntelligenceOCRConfig().get_complete_url(
+            api_base=prepared.api_base,
+            model=prepared.model,
+            optional_params=prepared.optional_params,
+            litellm_params=prepared.litellm_params,
+        )
+        assert url.startswith(custom)
 
     def test_generic_azure_ai_base_still_applies_to_mistral_ocr(self, monkeypatch):
         """Non doc-intelligence azure_ai models keep using AZURE_AI_API_BASE."""
         monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE)
+        from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig
 
         prepared = _prepare("azure_ai/mistral-document-ai-2505", None)
 
         assert prepared.api_base == _AZURE_AI_API_BASE
+        url = AzureAIOCRConfig().get_complete_url(
+            api_base=prepared.api_base,
+            model=prepared.model,
+            optional_params=prepared.optional_params,
+            litellm_params=prepared.litellm_params,
+        )
+        assert url.startswith(_AZURE_AI_API_BASE)

tests/test_litellm/rust_bridge/native_route_wheel_test.py

@@ -10,11 +10,13 @@ import sys
 import tempfile
 import threading
 import zipfile
+from dataclasses import dataclass
 from http.client import HTTPMessage
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from pathlib import Path
 from socket import socket as Socket
 from typing import Final
+from urllib.error import HTTPError
 
 REQUEST_STARTED: Final = threading.Event()
 REQUEST_CANCELLED: Final = threading.Event()
@@ -123,20 +125,54 @@ def load_native(native_path: Path) -> object:
     return native_module
 
 
+@dataclass(frozen=True)
+class OCRBoundary:
+    api_base: str
+    outcome: str
+
+    def prepare(self) -> dict[str, object]:
+        return {
+            "model": "mistral-ocr-latest",
+            "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"},
+            "include_image_base64": True,
+        }
+
+    async def aprepare(self) -> dict[str, object]:
+        return self.prepare()
+
+    def encode(self, roots: dict[str, object]) -> tuple[str, list[tuple[bytes, bytes]], bytes, float]:
+        return (
+            f"{self.api_base}/v1/ocr",
+            [
+                (b"authorization", b"Bearer sk-native"),
+                (b"content-type", b"application/json"),
+                (b"x-test-route", b"ocr"),
+                (b"x-test-outcome", self.outcome.encode()),
+            ],
+            json.dumps(roots).encode(),
+            3.0,
+        )
+
+    def finish(self, wire: tuple[int, list[tuple[bytes, bytes]], bytes]) -> object:
+        status, headers, content = wire
+        assert (b"content-type", b"application/json") in headers
+        if status != 200:
+            assert content == b'{"error":"native-rate-limit"}'
+            raise HTTPError(f"{self.api_base}/v1/ocr", status, "native-rate-limit", HTTPMessage(), None)
+        return json.loads(content)
+
+    async def afinish(self, wire: tuple[int, list[tuple[bytes, bytes]], bytes]) -> object:
+        return self.finish(wire)
+
+
 def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
+    if route == "ocr":
+        return {"boundary": OCRBoundary(api_base, outcome)}
     common: Final = {
         "api_base": api_base,
         "extra_headers": {"x-test-outcome": outcome, "x-test-route": route},
         "timeout_seconds": 3.0,
     }
-    if route == "ocr":
-        return common | {
-            "model": "mistral-ocr-latest",
-            "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"},
-            "api_key": "sk-native",
-            "custom_llm_provider": "mistral",
-            "optional_params": {"include_image_base64": True},
-        }
     if route == "transcription":
         return common | {
             "model": "mistral.voxtral-mini-3b-2507",
@@ -192,7 +228,11 @@ def success_value(route: str, response: dict[object, object]) -> object:
 
 
 def assert_rate_limit(native: object, route: str, error: BaseException) -> None:
-    if route in {"ocr", "chat_completions"}:
+    if route == "ocr":
+        if not isinstance(error, HTTPError) or error.code != 429:
+            raise AssertionError(f"{route} returned the wrong 429 error: {error!r}")
+        return
+    if route == "chat_completions":
         upstream_error: Final = native.RustUpstreamError
         if not isinstance(error, upstream_error) or error.args[0] != 429:
             raise AssertionError(f"{route} returned the wrong 429 error: {error!r}")
@@ -207,7 +247,7 @@ def exercise_sync(native: object, api_base: str) -> None:
         assert_success(route, function(**route_kwargs(route, api_base, "success")))
         try:
             function(**route_kwargs(route, api_base, "429"))
-        except (RuntimeError, native.RustUpstreamError) as error:
+        except (HTTPError, RuntimeError, native.RustUpstreamError) as error:
             assert_rate_limit(native, route, error)
         else:
             raise AssertionError(f"{route} accepted a 429 response")
@@ -219,7 +259,7 @@ async def exercise_async(native: object, api_base: str) -> None:
         assert_success(route, await function(**route_kwargs(route, api_base, "success")))
         try:
             await function(**route_kwargs(route, api_base, "429"))
-        except (RuntimeError, native.RustUpstreamError) as error:
+        except (HTTPError, RuntimeError, native.RustUpstreamError) as error:
             assert_rate_limit(native, route, error)
         else:
             raise AssertionError(f"a{route} accepted a 429 response")
@@ -227,12 +267,7 @@ async def exercise_async(native: object, api_base: str) -> None:
 
 async def exercise_async_concurrency(native: object, api_base: str) -> None:
     responses: Final = await asyncio.wait_for(
-        asyncio.gather(
-            *(
-                native.amessages(**route_kwargs("messages", api_base, "success"))
-                for _ in range(32)
-            )
-        ),
+        asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))),
         timeout=15,
     )
     for response in responses: