cleanup

.github/workflows/test-rust.yml

@@ -4,7 +4,6 @@ on:
   push:
     paths:
       - "litellm-rust/**"
-      - "tests/test_litellm/**"
       - ".cargo/**"
       - "pyproject.toml"
       - "uv.lock"
@@ -12,6 +11,7 @@ on:
       - ".github/actions/setup-uv-with-retries/**"
       - ".github/scripts/smoke_test_native_wheel.py"
       - ".github/scripts/verify_linux_native_wheel.py"
+      - "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
       - ".github/workflows/test-rust.yml"
   pull_request:
     branches:
@@ -21,7 +21,6 @@ on:
       - "litellm_**"
     paths:
       - "litellm-rust/**"
-      - "tests/test_litellm/**"
       - ".cargo/**"
       - "pyproject.toml"
       - "uv.lock"
@@ -29,6 +28,7 @@ on:
       - ".github/actions/setup-uv-with-retries/**"
       - ".github/scripts/smoke_test_native_wheel.py"
       - ".github/scripts/verify_linux_native_wheel.py"
+      - "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
       - ".github/workflows/test-rust.yml"
 
 permissions:
@@ -102,7 +102,7 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-cargo-${{ github.job }}-
 
-      - run: cargo test --workspace --locked --exclude litellm-python-interop --exclude litellm-python-bridge
+      - run: cargo test --workspace --locked
         working-directory: litellm-rust
 
       - run: cargo test -p litellm-core --features bedrock-auth --locked
@@ -131,5 +131,5 @@ jobs:
       - name: Check Python fixtures for Cargo tests
         run: make lint-rust-python-fixtures
 
-      - name: Test retained callbacks and OCR with repo Python
+      - name: Test retained callbacks with repo Python
         run: make test-rust-python

Makefile

@@ -314,11 +314,11 @@ test-rust-python: install-rust-python-test-deps
 	PYTHONPATH="$(CURDIR):$$site_packages$${PYTHONPATH:+:$$PYTHONPATH}" \
 	LITELLM_LOCAL_MODEL_COST_MAP=True \
 	cargo test --manifest-path litellm-rust/Cargo.toml \
-		-p litellm-python-interop -p litellm-python-bridge --tests --locked -- --include-ignored
+		-p litellm-python-interop --tests --locked -- --include-ignored
 
 lint-rust-python-fixtures:
-	$(UV) tool run --from ruff==0.15.3 ruff check --config ruff-tests.toml litellm-rust/crates/python-interop/tests litellm-rust/crates/python-bridge/tests
-	$(UV) tool run --from ruff==0.15.3 ruff format --check --config ruff-tests.toml litellm-rust/crates/python-interop/tests litellm-rust/crates/python-bridge/tests
+	$(UV) tool run --from ruff==0.15.3 ruff check --config ruff-tests.toml litellm-rust/crates/python-interop/tests
+	$(UV) tool run --from ruff==0.15.3 ruff format --check --config ruff-tests.toml litellm-rust/crates/python-interop/tests
 
 test: install-test-deps
 	$(UV_RUN) pytest tests/

litellm-rust/Cargo.lock

@@ -1471,10 +1471,8 @@ dependencies = [
  "litellm-python-interop",
  "pyo3",
  "pyo3-async-runtimes",
- "rstest",
  "serde",
  "serde_json",
- "serial_test",
  "tokio",
  "tokio-tungstenite",
  "tracing",

litellm-rust/README.md

@@ -58,8 +58,8 @@ runs for changes under `litellm-rust/`.
 ### Python-Integrated Tests
 
 From the repository root, run the ignored Cargo tests that need the repository's
-Python dependencies and the pinned Ruff checks over both crates' `tests/`
-directories:
+Python dependencies and the pinned Ruff checks over the interop crate's Python
+test fixtures:
 
 ```bash
 make test-rust-python
@@ -68,10 +68,16 @@ make lint-rust-python-fixtures
 
 `test-rust-python` installs the locked SDK dependencies with uv, points
 `PYO3_PYTHON` at the project interpreter, and runs
-`cargo test -p litellm-python-interop -p litellm-python-bridge --tests --locked`.
-`lint-rust-python-fixtures` runs pinless linting only, no sync.
+`cargo test -p litellm-python-interop --tests --locked -- --include-ignored`.
+`lint-rust-python-fixtures` runs pinned Ruff lint and formatting checks without
+syncing the project environment
 
-The callback lifecycle and retained OCR scenarios use
+These 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
+
+The callback lifecycle scenarios use
 `#[serial(python_interpreter)]` to isolate CPython GC and interpreter-wide
 LiteLLM settings under `cargo test`. Compatible tests in the same binary use
 `#[parallel(python_interpreter)]`: they may overlap each other, but not an

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

@@ -32,8 +32,6 @@ 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,7 +1,5 @@
 //! 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

@@ -1,77 +0,0 @@
-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,
-    })
-}
-
-#[cfg(test)]
-mod tests;

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

@@ -1,44 +0,0 @@
-use super::*;
-
-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/AGENTS.md

@@ -1,6 +1,3 @@
 litellm-python-bridge is the PyO3 cdylib that exposes LiteLLM Rust APIs to the Python SDK. Keep API registration, domain dependency wiring, request assembly, and Python exception mapping here. Put domain-neutral Python/Serde conversion and GIL primitives in litellm-python-interop.
 
 Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call the core entrypoint.
-
-The retained routes (`ocr_retained`, and any future `retained_http` variants) are a Python-compatibility adapter. Python owns prepare, transform, and logging; Rust owns only the buffered POST and the boundary call marshaling. There is no retry, billing, guardrail, or logging orchestration here, and none of it should be added.
-

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

@@ -7,7 +7,7 @@ repository.workspace = true
 
 [lib]
 name = "_native"
-crate-type = ["cdylib", "rlib"]
+crate-type = ["cdylib"]
 
 [features]
 default = ["abi3"]
@@ -34,8 +34,6 @@ tokio.workspace = true
 
 [dev-dependencies]
 criterion = "0.8.2"
-rstest.workspace = true
-serial_test.workspace = true
 tokio-tungstenite.workspace = true
 tracing.workspace = true
 

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

@@ -37,37 +37,6 @@ fn run_sync_on<T, F>(
 where
     T: Serialize + Send + 'static,
     F: Future<Output = Result<T, Error>> + Send + 'static,
-{
-    let result = run_sync_value_on(py, runtime, 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,
-        map_error,
-    )
-}
-
-fn run_sync_value_on<T, F>(
-    py: Python<'_>,
-    runtime: &Runtime,
-    future: F,
-    map_error: fn(Error) -> PyErr,
-) -> PyResult<T>
-where
-    T: Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
 {
     if Handle::try_current().is_ok() {
         return Err(PyRuntimeError::new_err(
@@ -76,7 +45,8 @@ where
     }
 
     let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
-    map_core_result(result, map_error)
+    let result = map_core_result(result, map_error)?;
+    Pythonized(result).into_pyobject(py).map(Bound::unbind)
 }
 
 pub(crate) fn run_async<T, F>(
@@ -89,19 +59,12 @@ where
     F: Future<Output = Result<T, Error>> + Send + 'static,
 {
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
-        let result = run_async_value(future, map_error).await?;
+        let result = catch_future_panic(future).await?;
+        let result = map_core_result(result, map_error)?;
         Ok(Pythonized(result))
     })
 }
 
-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),
@@ -112,7 +75,7 @@ fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -
     }
 }
 
-pub(crate) async fn catch_future_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
+async fn catch_future_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
 where
     F: Future<Output = Result<T, Error>>,
 {

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

@@ -63,7 +63,7 @@ impl ResponsesWebSocketConnection {
 }
 
 #[pymodule(gil_used = false)]
-pub mod _native {
+mod _native {
     use pyo3::prelude::*;
 
     #[pymodule_init]
@@ -98,8 +98,6 @@ mod tests {
                 "RustUpstreamError",
                 "ocr",
                 "aocr",
-                "ocr_retained",
-                "aocr_retained",
                 "transcription",
                 "atranscription",
                 "messages",

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

@@ -10,12 +10,9 @@ mod audio_transcription;
 mod chat_completions;
 mod messages;
 mod ocr;
-mod ocr_retained;
-mod retained_http;
 
 pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
     ocr::register(module)?;
-    ocr_retained::register(module)?;
     audio_transcription::register(module)?;
     messages::register(module)?;
     chat_completions::register(module)?;

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

@@ -1,18 +0,0 @@
-use pyo3::prelude::*;
-
-use super::retained_http;
-
-#[pyfunction]
-fn ocr_retained(boundary: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
-    retained_http::run_sync(boundary)
-}
-
-#[pyfunction]
-fn aocr_retained(boundary: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
-    retained_http::run_async(boundary)
-}
-
-pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    module.add_function(wrap_pyfunction!(ocr_retained, module)?)?;
-    module.add_function(wrap_pyfunction!(aocr_retained, module)?)
-}

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

@@ -1,203 +0,0 @@
-use litellm_core::http_utils::buffered_post::{self, Request, Response};
-use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall};
-use pyo3::prelude::*;
-use pyo3::sync::PyOnceLock;
-use pyo3::types::{PyBytes, PyList, PyTuple};
-
-use crate::errors::core_error_to_pyerr;
-use crate::execution::{run_async_value, run_sync_value};
-
-pub(crate) struct MethodBinding {
-    pub(crate) name: &'static str,
-    pub(crate) mode: InvocationMode,
-}
-
-pub(crate) enum BoundaryMethod {
-    Prepare,
-    Encode,
-    Finish,
-}
-
-impl BoundaryMethod {
-    pub(crate) fn resolve(self, asynchronous: bool) -> MethodBinding {
-        match (self, asynchronous) {
-            (Self::Prepare, true) => MethodBinding {
-                name: "aprepare",
-                mode: InvocationMode::Await,
-            },
-            (Self::Prepare, false) => MethodBinding {
-                name: "prepare",
-                mode: InvocationMode::Direct,
-            },
-            (Self::Encode, _) => MethodBinding {
-                name: "encode",
-                mode: InvocationMode::Direct,
-            },
-            (Self::Finish, true) => MethodBinding {
-                name: "afinish",
-                mode: InvocationMode::Await,
-            },
-            (Self::Finish, false) => MethodBinding {
-                name: "finish",
-                mode: InvocationMode::Direct,
-            },
-        }
-    }
-}
-
-fn invoke(
-    boundary: &Bound<'_, PyAny>,
-    binding: MethodBinding,
-    args: Bound<'_, PyTuple>,
-) -> PyResult<Py<PyAny>> {
-    let call = PreparedCall::new(
-        binding.mode,
-        boundary.getattr(binding.name)?.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>> {
-    invoke(
-        boundary,
-        BoundaryMethod::Prepare.resolve(asynchronous),
-        PyTuple::empty(boundary.py()),
-    )
-}
-
-#[pyfunction]
-fn encode(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
-    invoke(
-        boundary,
-        BoundaryMethod::Encode.resolve(false),
-        PyTuple::new(boundary.py(), [roots])?,
-    )
-}
-
-fn request(execution: &Bound<'_, PyAny>) -> PyResult<Request> {
-    type ByteHeaders<'py> = Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>;
-    let encoded = execution.call_method0("encode")?;
-    let (url, headers, body, timeout_seconds): (String, ByteHeaders<'_>, Bound<'_, PyBytes>, f64) =
-        encoded.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,
-    })
-}
-
-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<'py>(execution: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
-    let request = request(execution)?;
-    pyo3_async_runtimes::tokio::future_into_py(execution.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>> {
-    invoke(
-        boundary,
-        BoundaryMethod::Finish.resolve(asynchronous),
-        PyTuple::new(boundary.py(), [wire])?,
-    )
-}
-
-pub(crate) fn run_sync(boundary: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
-    let py = boundary.py();
-    let execution = execution_module(py)?
-        .getattr("RetainedExecution")?
-        .call1((boundary,))?;
-    execution.call_method0("prepare")?;
-    let request = request(&execution)?;
-    let response = run_sync_value(py, buffered_post::send(request), core_error_to_pyerr)?;
-    let wire = Wire(response).into_pyobject(py)?;
-    execution.call_method1("finish", (wire,)).map(Bound::unbind)
-}
-
-pub(crate) fn run_async(boundary: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
-    let module = execution_module(boundary.py())?;
-    let execution = module.getattr("RetainedExecution")?.call1((boundary,))?;
-    module
-        .getattr("drive")?
-        .call1((execution,))
-        .map(Bound::unbind)
-}
-
-fn execution_module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
-    static MODULE: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
-    if MODULE.get(py).is_none() {
-        let module = PyModule::from_code(
-            py,
-            c"class RetainedExecution:
-    __slots__ = ('binding', 'roots')
-
-    def __init__(self, binding):
-        self.binding = binding
-        self.roots = None
-
-    def prepare(self):
-        self.roots = _prepare(self.binding, False)
-
-    async def aprepare(self):
-        self.roots = await _prepare(self.binding, True)
-
-    def encode(self):
-        return _encode(self.binding, self.roots)
-
-    def finish(self, wire):
-        return _finish(self.binding, wire, False)
-
-    async def afinish(self, wire):
-        return await _finish(self.binding, wire, True)
-
-async def drive(execution):
-    await execution.aprepare()
-    wire = await _send(execution)
-    return await execution.afinish(wire)
-",
-            c"retained_execution.py",
-            c"_retained_execution",
-        )?;
-        module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?;
-        module.add("_encode", wrap_pyfunction!(encode, &module)?)?;
-        module.add("_finish", wrap_pyfunction!(finish, &module)?)?;
-        module.add("_send", wrap_pyfunction!(send, &module)?)?;
-        let _ = MODULE.set(py, module.unbind());
-    }
-    let module = MODULE.get(py).unwrap();
-    Ok(module.bind(py).clone())
-}

litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py

@@ -1,724 +0,0 @@
-"""Executed by the PyO3 retained test with the actual built module in `native`."""
-
-import asyncio
-import contextvars
-import gc
-import inspect
-import json
-import threading
-import unittest
-import weakref
-from copy import deepcopy
-from datetime import datetime
-from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
-from types import ModuleType
-from unittest.mock import patch
-
-import httpx
-
-import litellm
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.litellm_core_utils.litellm_logging import Logging
-from litellm.llms.base_llm.chat.transformation import BaseLLMException
-from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
-from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
-from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
-from litellm.rust_bridge.ocr_retained import OCREncoded, OCRRetainedBoundary, OCRRoots
-
-native: ModuleType = globals()["native"]
-context = contextvars.ContextVar("retained-real-boundary", default="unset")
-MODEL = "mistral-ocr-latest"
-RESPONSE = {"pages": [{"index": 0, "markdown": "local OCR"}], "model": MODEL, "usage_info": {"pages_processed": 1}}
-
-
-class Graph(dict):
-    pass
-
-
-class Header(str):
-    pass
-
-
-class PreCallAbort(BaseException):
-    pass
-
-
-class CopiedDocumentBoundary(OCRRetainedBoundary):
-    def prepare(self) -> OCRRoots:
-        self.document = deepcopy(self.document)
-        return super().prepare()
-
-    async def aprepare(self) -> OCRRoots:
-        self.document = deepcopy(self.document)
-        return await super().aprepare()
-
-
-class ReboundBodyBoundary(OCRRetainedBoundary):
-    def encode(self, roots: OCRRoots) -> OCREncoded:
-        headers, url, _body, files = roots
-        view = self.logging_obj.model_call_details["additional_args"]
-        return super().encode((headers, url, view["complete_input_dict"], files))
-
-
-class ReboundHeadersBoundary(OCRRetainedBoundary):
-    def encode(self, roots: OCRRoots) -> OCREncoded:
-        _headers, url, body, files = roots
-        view = self.logging_obj.model_call_details["additional_args"]
-        return super().encode((view["headers"], url, body, files))
-
-
-class Callback(CustomLogger):
-    def __init__(self, action):
-        super().__init__()
-        self.action = action
-        self.failures = []
-        self.calls = 0
-
-    def log_pre_api_call(self, model, messages, kwargs):
-        self.calls += 1
-        try:
-            return self.action(kwargs["additional_args"])
-        except AssertionError as error:
-            self.failures.append(str(error))
-            raise
-
-
-class Server:
-    def __init__(self):
-        self.requests = []
-        self.started = threading.Event()
-        self.release = threading.Event()
-        self.finished = threading.Event()
-        owner = self
-
-        class Handler(BaseHTTPRequestHandler):
-            def log_message(self, *args):
-                pass
-
-            def do_POST(self):
-                body = self.rfile.read(int(self.headers["Content-Length"]))
-                owner.requests.append((self.path, sorted((k.lower(), v) for k, v in self.headers.items()), body))
-                if self.path == "/blocked/v1/ocr":
-                    owner.started.set()
-                    if not owner.release.wait(10):
-                        owner.finished.set()
-                        return
-                failed = self.path == "/error/v1/ocr"
-                payload = b'{"message":"controlled HTTP failure"}' if failed else json.dumps(RESPONSE).encode()
-                self.send_response(429 if failed else 200)
-                self.send_header("Content-Type", "application/json")
-                self.send_header("Content-Length", str(len(payload)))
-                self.end_headers()
-                try:
-                    self.wfile.write(payload)
-                except (BrokenPipeError, ConnectionResetError):
-                    pass
-                finally:
-                    if self.path == "/blocked/v1/ocr":
-                        owner.finished.set()
-
-        self.httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
-        self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True)
-        self.thread.start()
-        self.url = f"http://127.0.0.1:{self.httpd.server_port}"
-
-    def close(self):
-        self.release.set()
-        self.httpd.shutdown()
-        self.httpd.server_close()
-        self.thread.join(5)
-        assert not self.thread.is_alive()
-
-
-def inputs(server, callbacks=(), *, document=None, optional=None, path="", client=None):
-    return {
-        "model": MODEL,
-        "document": Graph(type="document_url", document_url="https://example.test/original.pdf")
-        if document is None
-        else document,
-        "optional_params": {} if optional is None else optional,
-        "logging_obj": Logging(
-            model=MODEL,
-            messages=[],
-            stream=False,
-            call_type="ocr",
-            start_time=datetime.now(),
-            litellm_call_id="retained-real",
-            function_id="retained-real",
-            dynamic_input_callbacks=list(callbacks),
-            supports_correlation_logging=False,
-        ),
-        "api_key": "local-test-key",
-        "api_base": server.url + path,
-        "headers": {"X-Proof": "original"},
-        "provider_config": MistralOCRConfig(),
-        "litellm_params": {},
-        "custom_llm_provider": "mistral",
-        "timeout": 5.0,
-        "client": client,
-    }
-
-
-def invoke(mode, kwargs, *, boundary_factory=OCRRetainedBoundary):
-    handler = BaseLLMHTTPHandler()
-    if mode == "python-sync":
-        return handler.ocr(**kwargs)
-    if mode == "python-async":
-        return handler.async_ocr(**kwargs)
-    if mode == "native-sync":
-        return native.ocr_retained(boundary_factory(handler=handler, **kwargs))
-    assert mode == "native-async"
-    return native.aocr_retained(boundary_factory(handler=handler, **kwargs))
-
-
-class RealBoundaryTests(unittest.TestCase):
-    def setUp(self):
-        self.server = Server()
-        self.addCleanup(self.server.close)
-        self.sync_client = HTTPHandler(timeout=5.0)
-        self.addCleanup(self.sync_client.close)
-
-    def check_callbacks(self, callbacks):
-        for callback in callbacks:
-            self.assertEqual(callback.failures, [])
-            self.assertEqual(callback.calls, 1)
-
-    async def differential(self, mode, *, boundary_factory=OCRRetainedBoundary):
-        context.set("caller")
-        thread = threading.get_ident()
-        task = asyncio.current_task()
-        document = Graph(type="document_url", document_url="https://example.test/original.pdf")
-        nested = Graph(values=[1])
-        optional = {"unknown_python_json": {7: ("tuple", 2)}, "nested": nested, "nested_alias": nested["values"]}
-        retained = {}
-        events = []
-
-        def phase(name, expected):
-            self.assertEqual(threading.get_ident(), thread)
-            self.assertIs(asyncio.current_task(), task)
-            self.assertEqual(context.get(), expected)
-            events.append(name)
-
-        def mutate(view):
-            phase("mutate", "caller")
-            body, headers = view["complete_input_dict"], view["headers"]
-            self.assertIs(body["document"], document, "caller document identity was not retained")
-            self.assertIs(body["nested"], nested)
-            self.assertIs(body["nested_alias"], nested["values"])
-            self.assertIs(body["unknown_python_json"], optional["unknown_python_json"])
-            retained.update(body=body, headers=headers, view=view)
-            headers["X-Proof"] = "in-place"
-            body["body_mutation"] = True
-            nested["values"].append(2)
-            view["headers"] = {"X-Proof": "must-not-send"}
-            view["complete_input_dict"] = {"document": {"document_url": "must-not-send"}}
-            context.set("mutated")
-            child_callback = Callback(lambda _: events.append("reentry"))
-            child = invoke("native-sync", inputs(self.server, [child_callback], path="/child", client=self.sync_client))
-            self.assertEqual(child.pages[0].markdown, "local OCR")
-            self.check_callbacks([child_callback])
-            return {"headers": {"X-Proof": "ignored-return"}, "complete_input_dict": {"invalid": object()}}
-
-        def mutate_then_raise(view):
-            phase("raise", "mutated")
-            self.assertIs(view, retained["view"])
-            self.assertEqual(view["complete_input_dict"], {"document": {"document_url": "must-not-send"}})
-            self.assertEqual(view["headers"], {"X-Proof": "must-not-send"})
-            retained["body"]["before_error"] = True
-            retained["headers"]["X-Before-Error"] = "yes"
-            context.set("caught")
-            raise RuntimeError("intentional non-blocking pre_call error")
-
-        def closure_only(view):
-            phase("later", "caught")
-            self.assertIs(view, retained["view"])
-            self.assertIsNot(view["complete_input_dict"], retained["body"])
-            self.assertIsNot(view["headers"], retained["headers"])
-            self.assertTrue(retained["body"]["before_error"])
-            self.assertEqual(retained["headers"]["X-Before-Error"], "yes")
-            self.assertIs(retained["body"]["nested_alias"], nested["values"])
-            self.assertEqual(retained["body"]["nested_alias"], [1, 2])
-            view["complete_input_dict"]["observed"] = True
-            view["headers"]["X-View-Only"] = "not-on-wire"
-            document["document_url"] = "https://example.test/closure.pdf"
-            context.set("later")
-
-        callbacks = [Callback(action) for action in (mutate, mutate_then_raise, closure_only)]
-        async_client = AsyncHTTPHandler(timeout=5.0)
-        try:
-            kwargs = inputs(
-                self.server,
-                callbacks,
-                document=document,
-                optional=optional,
-                client=async_client if mode.endswith("async") else self.sync_client,
-            )
-            logging_ref = weakref.ref(kwargs["logging_obj"])
-            before = len(self.server.requests)
-            pending = invoke(mode, kwargs, boundary_factory=boundary_factory)
-            if mode.endswith("async"):
-                self.assertTrue(inspect.iscoroutine(pending))
-                self.assertEqual(events, [])
-                self.assertEqual(len(self.server.requests), before)
-                self.assertNotIn("additional_args", kwargs["logging_obj"].model_call_details)
-                response = await pending
-            else:
-                response = pending
-            self.check_callbacks(callbacks)
-            self.assertEqual(events, ["mutate", "reentry", "raise", "later"])
-            self.assertEqual(context.get(), "later")
-            self.assertEqual(len(self.server.requests), before + 2)
-            wire = self.server.requests[-1]
-            self.assertEqual(wire[0], "/v1/ocr")
-            expected = (
-                b'{"model":"mistral-ocr-latest","document":{"type":"document_url",'
-                b'"document_url":"https://example.test/closure.pdf"},"unknown_python_json":{"7":["tuple",2]},'
-                b'"nested":{"values":[1,2]},"nested_alias":[1,2],"body_mutation":true,"before_error":true}'
-            )
-            self.assertEqual(wire[2], expected, "wire body must encode retained mutations after pre_call")
-            self.assertIn(("x-proof", "in-place"), wire[1], "wire headers must use retained execution headers")
-            self.assertIn(("x-before-error", "yes"), wire[1])
-            self.assertIn(("authorization", "Bearer local-test-key"), wire[1])
-            self.assertFalse(any(name == "x-view-only" for name, _ in wire[1]))
-            self.assertEqual(
-                retained["view"]["complete_input_dict"],
-                {"document": {"document_url": "must-not-send"}, "observed": True},
-            )
-            self.assertEqual(retained["view"]["headers"], {"X-Proof": "must-not-send", "X-View-Only": "not-on-wire"})
-            del kwargs, pending
-            gc.collect()
-            self.assertIsNone(logging_ref(), "logging owner survived the completed call")
-            self.assertIs(retained["body"]["document"], document)
-            retained["body"]["nested"]["values"].append(3)
-            retained["headers"]["X-After-Return"] = "usable"
-            self.assertEqual(optional["nested"]["values"], [1, 2, 3])
-            self.assertIs(retained["body"]["nested_alias"], optional["nested"]["values"])
-            self.assertEqual(retained["body"]["nested_alias"], [1, 2, 3])
-            self.assertEqual(retained["headers"]["X-After-Return"], "usable")
-            self.assertEqual(wire[2], expected)
-            self.assertEqual(response.pages[0].markdown, "local OCR")
-            return wire, response.model_dump()
-        finally:
-            await async_client.close()
-
-    def test_differential_callbacks_wire(self):
-        async def exercise():
-            baseline = await self.differential("python-sync")
-            for mode in ("python-async", "native-sync", "native-async"):
-                with self.subTest(mode=mode):
-                    self.assertEqual(await self.differential(mode), baseline)
-
-        asyncio.run(exercise())
-
-    @unittest.expectedFailure
-    def test_known_gap_send_time_auth_and_request_hook(self):
-        """UC-HTTPX-SEND: pre_call, client auth, then request hook must determine the provider's wire headers."""
-        observations = []
-        for mode in ("python-sync", "native-sync"):
-            events = []
-
-            def request_hook(request, events=events):
-                events.append(("request_hook", request.headers["Authorization"]))
-                request.headers["X-Send-Hook"] = "present"
-
-            callback = Callback(
-                lambda view, events=events: events.append(("pre_call", view["headers"]["Authorization"]))
-            )
-            with httpx.Client(auth=("local-user", "local-password"), event_hooks={"request": [request_hook]}) as client:
-                before = len(self.server.requests)
-                response = invoke(mode, inputs(self.server, [callback], client=HTTPHandler(client=client)))
-            self.check_callbacks([callback])
-            self.assertEqual(response.pages[0].markdown, "local OCR")
-            self.assertEqual(len(self.server.requests), before + 1)
-            headers = dict(self.server.requests[-1][1])
-            observations.append((events, headers["authorization"], headers.get("x-send-hook")))
-        baseline, retained = observations
-        self.assertEqual(baseline[0][0], ("pre_call", "Bearer local-test-key"))
-        self.assertEqual(baseline[0][1], ("request_hook", baseline[1]))
-        self.assertTrue(baseline[1].startswith("Basic "))
-        self.assertEqual(baseline[2], "present")
-        self.assertEqual(retained, baseline, "UC-HTTPX-SEND: retained send omitted client auth/request hook")
-
-    @unittest.expectedFailure
-    def test_known_gap_custom_transport(self):
-        """UC-HTTPX-SEND: after pre_call, the client's transport must select the response without a network POST."""
-        observations = []
-        for mode in ("python-sync", "native-sync"):
-            events = []
-
-            def transport(request, events=events):
-                events.append("transport")
-                return httpx.Response(200, json={**RESPONSE, "pages": [{"index": 0, "markdown": "custom transport"}]})
-
-            callback = Callback(lambda view, events=events: events.append("pre_call"))
-            with httpx.Client(transport=httpx.MockTransport(transport)) as client:
-                before = len(self.server.requests)
-                response = invoke(mode, inputs(self.server, [callback], client=HTTPHandler(client=client)))
-            self.check_callbacks([callback])
-            observations.append((events, len(self.server.requests) - before, response.pages[0].markdown))
-        baseline, retained = observations
-        self.assertEqual(baseline, (["pre_call", "transport"], 0, "custom transport"))
-        self.assertEqual(retained, baseline, "UC-HTTPX-SEND: retained send bypassed the custom transport")
-
-    @unittest.expectedFailure
-    def test_known_gap_pre_call_timeout_mutation(self):
-        """UC-TIMEOUT-MUTATION: a caller's timeout is changed by pre_call; encoding must accept read=None."""
-        observations = []
-        for mode in ("python-sync", "native-sync"):
-            timeout = httpx.Timeout(5.0)
-
-            def mutate_timeout(view, timeout=timeout):
-                timeout.read = None
-
-            callback = Callback(mutate_timeout)
-            kwargs = inputs(self.server, [callback], client=self.sync_client)
-            kwargs["timeout"] = timeout
-            before = len(self.server.requests)
-            try:
-                response = invoke(mode, kwargs)
-            except ValueError as error:
-                outcome = ("error", str(error))
-            else:
-                outcome = ("success", response.pages[0].markdown)
-            self.check_callbacks([callback])
-            self.assertIsNone(timeout.read)
-            observations.append((outcome, len(self.server.requests) - before))
-        baseline, retained = observations
-        self.assertEqual(baseline, (("success", "local OCR"), 1))
-        self.assertEqual(retained, baseline, "UC-TIMEOUT-MUTATION: retained encoding rejected the callback's timeout")
-
-    def check_negative_control(self, boundary_factory, failure, expected_requests):
-        async def exercise():
-            baseline = await self.differential("python-sync")
-            for mode in ("native-sync", "native-async"):
-                with self.subTest(mode=mode):
-                    symbol = "aocr_retained" if mode.endswith("async") else "ocr_retained"
-                    self.assertTrue(inspect.isbuiltin(getattr(native, symbol)))
-                    self.assertEqual(await self.differential(mode), baseline)
-                    before = len(self.server.requests)
-                    with self.assertRaisesRegex(AssertionError, failure):
-                        await self.differential(mode, boundary_factory=boundary_factory)
-                    self.assertEqual(len(self.server.requests), before + expected_requests)
-                    self.assertEqual(self.server.requests[-1][0], "/v1/ocr")
-
-        asyncio.run(exercise())
-
-    def test_negative_control_copied_caller_document(self):
-        self.check_negative_control(CopiedDocumentBoundary, "caller document identity was not retained", 1)
-
-    def test_negative_control_rebound_logging_body(self):
-        self.check_negative_control(ReboundBodyBoundary, "wire body must encode retained mutations after pre_call", 2)
-
-    def test_negative_control_rebound_logging_headers(self):
-        self.check_negative_control(ReboundHeadersBoundary, "wire headers must use retained execution headers", 2)
-
-    def test_differential_callback_retained_mutation_after_post_received(self):
-        async def suspended(mode):
-            self.server.started.clear()
-            self.server.release.clear()
-            self.server.finished.clear()
-            retained = {}
-
-            def retain(view):
-                retained.update(body=view["complete_input_dict"], headers=view["headers"], view=view)
-                view["complete_input_dict"] = {"replacement": True}
-                view["headers"] = {"X-Proof": "must-not-send"}
-
-            callback = Callback(retain)
-            client = AsyncHTTPHandler(timeout=5.0)
-            kwargs = inputs(self.server, [callback], path="/blocked", client=client)
-            logging = kwargs["logging_obj"]
-            logging.log_raw_request_response = True
-            before = len(self.server.requests)
-            task = asyncio.create_task(invoke(mode, kwargs))
-            try:
-                self.assertTrue(await asyncio.to_thread(self.server.started.wait, 5))
-                self.assertFalse(task.done())
-                self.assertFalse(self.server.release.is_set())
-                self.assertFalse(self.server.finished.is_set())
-                self.check_callbacks([callback])
-                self.assertEqual(len(self.server.requests), before + 1)
-                path, headers, body = self.server.requests[-1]
-                wire = (path, tuple(headers), body)
-                expected = (
-                    b'{"model":"mistral-ocr-latest","document":{"type":"document_url",'
-                    b'"document_url":"https://example.test/original.pdf"}}'
-                )
-                self.assertEqual(path, "/blocked/v1/ocr")
-                self.assertEqual(body, expected)
-                self.assertIn(("x-proof", "original"), headers)
-                self.assertNotIn(("x-proof", "must-not-send"), headers)
-                logged_body = logging.model_call_details["raw_request_typed_dict"]["raw_request_body"]
-                self.assertIs(logged_body, retained["body"])
-                self.assertIs(logged_body["document"], kwargs["document"])
-                self.assertIs(logging.model_call_details["additional_args"], retained["view"])
-                self.assertIsNot(retained["view"]["complete_input_dict"], retained["body"])
-                self.assertIsNot(retained["view"]["headers"], retained["headers"])
-
-                retained["body"]["after_encoding"] = True
-                retained["body"]["document"]["document_url"] = "https://example.test/while-blocked.pdf"
-                retained["headers"]["X-Proof"] = "while-blocked"
-                self.assertTrue(logged_body["after_encoding"])
-                self.assertEqual(kwargs["document"]["document_url"], "https://example.test/while-blocked.pdf")
-                self.assertEqual(logged_body["document"]["document_url"], "https://example.test/while-blocked.pdf")
-                self.assertEqual(retained["headers"]["X-Proof"], "while-blocked")
-                self.assertEqual(retained["view"]["complete_input_dict"], {"replacement": True})
-                self.assertEqual(retained["view"]["headers"], {"X-Proof": "must-not-send"})
-                self.assertFalse(task.done())
-                self.assertFalse(self.server.finished.is_set())
-                self.assertEqual((path, tuple(headers), body), wire)
-
-                self.server.release.set()
-                response = await asyncio.wait_for(task, 5)
-                self.assertEqual(response.pages[0].markdown, "local OCR")
-                self.assertIs(logging.model_call_details["raw_request_typed_dict"]["raw_request_body"], logged_body)
-                self.assertTrue(logged_body["after_encoding"])
-                self.assertEqual(len(self.server.requests), before + 1)
-                received_path, received_headers, received_body = self.server.requests[-1]
-                self.assertEqual((received_path, tuple(received_headers), received_body), wire)
-                self.check_callbacks([callback])
-                return wire, logged_body, retained["headers"], response.model_dump()
-            finally:
-                self.server.release.set()
-                if not task.done():
-                    task.cancel()
-                await asyncio.gather(task, return_exceptions=True)
-                await client.close()
-                if self.server.started.is_set():
-                    self.assertTrue(await asyncio.to_thread(self.server.finished.wait, 5))
-
-        async def exercise():
-            baseline = await suspended("python-async")
-            self.assertEqual(await suspended("native-async"), baseline)
-
-        asyncio.run(exercise())
-
-    def test_public_rust_dispatch_wire_fallback_and_escaping_base_exception(self):
-        async def exercise():
-            for asynchronous in (False, True):
-                for outcome in ("success", "missing-symbol", "pre-call-abort", "disabled"):
-                    with self.subTest(asynchronous=asynchronous, outcome=outcome):
-                        symbol = "aocr_retained" if asynchronous else "ocr_retained"
-                        self.assertTrue(inspect.isbuiltin(getattr(native, symbol)))
-                        missing_symbol = ModuleType("native_without_" + symbol)
-                        missing_symbol.__dict__.update(
-                            (name, value) for name, value in vars(native).items() if name != symbol
-                        )
-                        escaped = PreCallAbort("public pre_call must escape unchanged")
-                        before = len(self.server.requests)
-
-                        def mutate(view, outcome=outcome, escaped=escaped):
-                            view["headers"]["X-Proof"] = "public-in-place"
-                            view["complete_input_dict"]["public_mutation"] = True
-                            view["complete_input_dict"]["document"]["document_url"] = (
-                                "https://example.test/public-mutated.pdf"
-                            )
-                            if outcome == "pre-call-abort":
-                                raise escaped
-
-                        callback = Callback(mutate)
-                        kwargs = inputs(self.server, [callback])
-                        with patch(
-                            "litellm.rust_bridge.get_native_bridge",
-                            return_value=missing_symbol if outcome == "missing-symbol" else native,
-                        ) as loader:
-                            public_kwargs = {
-                                **{
-                                    key: kwargs[key]
-                                    for key in (
-                                        "model",
-                                        "document",
-                                        "api_key",
-                                        "api_base",
-                                        "custom_llm_provider",
-                                        "timeout",
-                                    )
-                                },
-                                "extra_headers": kwargs["headers"],
-                                "litellm_logging_obj": kwargs["logging_obj"],
-                                "rust": outcome != "disabled",
-                            }
-
-                            async def call(asynchronous=asynchronous, public_kwargs=public_kwargs):
-                                if asynchronous:
-                                    return await litellm.aocr(**public_kwargs)
-                                return litellm.ocr(**public_kwargs)
-
-                            if outcome == "pre-call-abort":
-                                with self.assertRaises(PreCallAbort) as caught:
-                                    await call()
-                                self.assertIs(caught.exception, escaped)
-                            else:
-                                response = await call()
-                                self.assertEqual(response.pages[0].markdown, "local OCR")
-
-                        if outcome == "disabled":
-                            loader.assert_not_called()
-                        else:
-                            loader.assert_called_once_with()
-                        self.check_callbacks([callback])
-                        self.assertEqual(len(self.server.requests), before + int(outcome != "pre-call-abort"))
-                        if outcome != "pre-call-abort":
-                            path, headers, body = self.server.requests[-1]
-                            self.assertEqual(path, "/v1/ocr")
-                            self.assertIn(("x-proof", "public-in-place"), headers)
-                            self.assertIn(("authorization", "Bearer local-test-key"), headers)
-                            self.assertEqual(
-                                json.loads(body),
-                                {
-                                    "model": MODEL,
-                                    "document": {
-                                        "type": "document_url",
-                                        "document_url": "https://example.test/public-mutated.pdf",
-                                    },
-                                    "public_mutation": True,
-                                },
-                            )
-
-        asyncio.run(exercise())
-
-    def lifecycle_inputs(self, outcome, retained):
-        refs = []
-
-        def callback(view):
-            body, headers = view["complete_input_dict"], view["headers"]
-            body["sentinel"] = Graph(alive=True)
-            headers["X-Sentinel"] = Header("alive")
-            refs.extend((weakref.ref(body["sentinel"]), weakref.ref(headers["X-Sentinel"])))
-            if outcome == "encoding":
-                body["not_json"] = object()
-            if retained is not None:
-                retained.extend((body, headers))
-            view["complete_input_dict"] = {}
-            view["headers"] = {}
-            if outcome == "pre-call-abort":
-                raise PreCallAbort("lifecycle pre_call abort")
-
-        logger = Callback(callback)
-        kwargs = inputs(
-            self.server,
-            [logger],
-            optional={"nested": Graph(alive=True)},
-            path={"http": "/error", "cancel": "/blocked"}.get(outcome, ""),
-            client=self.sync_client,
-        )
-        refs.extend(weakref.ref(kwargs[key]) for key in ("document", "logging_obj"))
-        refs.append(weakref.ref(kwargs["optional_params"]["nested"]))
-        return kwargs, logger, refs
-
-    async def lifecycle(self, mode, outcome, retained=None):
-        kwargs, logger, refs = self.lifecycle_inputs(outcome, retained)
-        before = len(self.server.requests)
-        async_client = AsyncHTTPHandler(timeout=5.0)
-        if mode.endswith("async"):
-            kwargs["client"] = async_client
-        try:
-            try:
-                pending = invoke(mode, kwargs)
-                response = await pending if mode.endswith("async") else pending
-            except BaseLLMException as error:
-                self.assertIn(outcome, ("encoding", "http"))
-                self.assertEqual(error.status_code, 500 if outcome == "encoding" else 429)
-                signature = (type(error), error.status_code, str(error))
-            except PreCallAbort as error:
-                self.assertEqual(outcome, "pre-call-abort")
-                signature = (type(error), str(error))
-            else:
-                self.assertEqual(outcome, "success")
-                self.assertEqual(response.pages[0].markdown, "local OCR")
-                signature = response.model_dump()
-            self.check_callbacks([logger])
-            self.assertEqual(len(refs), 5)
-            self.assertEqual(len(self.server.requests), before + (outcome not in ("encoding", "pre-call-abort")))
-            return refs, signature
-        finally:
-            await async_client.close()
-
-    def test_collection_after_success_and_failures(self):
-        async def exercise():
-            for outcome in ("success", "encoding", "http", "pre-call-abort"):
-                baseline = None
-                for mode in ("python-sync", "python-async", "native-sync", "native-async"):
-                    with self.subTest(mode=mode, outcome=outcome):
-                        refs, signature = await self.lifecycle(mode, outcome)
-                        gc.collect()
-                        self.assertTrue(all(ref() is None for ref in refs), f"request graph leaked: {mode=} {outcome=}")
-                        if baseline is None:
-                            baseline = signature
-                        self.assertEqual(signature, baseline)
-
-        asyncio.run(exercise())
-
-    def test_callback_retained_graph_remains_usable_then_collects(self):
-        async def exercise():
-            for outcome in ("success", "encoding", "http", "pre-call-abort"):
-                for mode in ("python-sync", "python-async", "native-sync", "native-async"):
-                    with self.subTest(mode=mode, outcome=outcome):
-                        retained = []
-                        refs, _ = await self.lifecycle(mode, outcome, retained)
-                        gc.collect()
-                        self.assertIsNone(refs[1](), f"logging owner survived: {mode=} {outcome=}")
-                        self.assertTrue(all(refs[index]() is not None for index in (0, 2, 3, 4)))
-                        self.assertIs(retained[0]["document"], refs[0]())
-                        self.assertIs(retained[0]["nested"], refs[2]())
-                        self.assertIs(retained[0]["sentinel"], refs[3]())
-                        self.assertIs(retained[1]["X-Sentinel"], refs[4]())
-                        retained[0]["document"]["after_return"] = "usable"
-                        retained[0]["nested"]["alive"] = "nested still usable"
-                        retained[0]["sentinel"]["alive"] = "still usable"
-                        retained[1]["X-After-Return"] = "usable"
-                        self.assertEqual(refs[0]()["after_return"], "usable")
-                        self.assertEqual(refs[2]()["alive"], "nested still usable")
-                        self.assertEqual(refs[3]()["alive"], "still usable")
-                        self.assertEqual(retained[1]["X-After-Return"], "usable")
-                        retained.clear()
-                        gc.collect()
-                        self.assertTrue(all(ref() is None for ref in refs))
-
-        asyncio.run(exercise())
-
-    def test_collection_after_cancellation_during_blocked_transport(self):
-        async def cancel():
-            kwargs, logger, refs = self.lifecycle_inputs("cancel", None)
-            client = AsyncHTTPHandler(timeout=5.0)
-            kwargs["client"] = client
-            task = asyncio.create_task(invoke("native-async", kwargs))
-            try:
-                self.assertTrue(await asyncio.to_thread(self.server.started.wait, 5))
-                self.check_callbacks([logger])
-                self.assertEqual(len(refs), 5)
-                gc.collect()
-                self.assertTrue(all(ref() is not None for ref in refs))
-                task.cancel()
-                with self.assertRaises(asyncio.CancelledError):
-                    await task
-                self.assertFalse(self.server.release.is_set())
-                return refs
-            finally:
-                if not task.done():
-                    task.cancel()
-                    await asyncio.gather(task, return_exceptions=True)
-                await client.close()
-
-        async def exercise():
-            refs = await cancel()
-            barrier = asyncio.get_running_loop().create_future()
-            asyncio.get_running_loop().call_soon(barrier.set_result, None)
-            await barrier
-            gc.collect()
-            self.assertTrue(all(ref() is None for ref in refs), "cancelled native call retained the request graph")
-            self.server.release.set()
-            self.assertTrue(await asyncio.to_thread(self.server.finished.wait, 5))
-            self.assertEqual(len(self.server.requests), 1)
-
-        asyncio.run(exercise())
-
-
-def run_case(scenario):
-    suite = unittest.TestSuite([RealBoundaryTests("test_" + scenario)])
-    result = unittest.TextTestRunner(verbosity=2).run(suite)
-    for _, traceback in result.expectedFailures:
-        assert "AssertionError:" in traceback and (
-            "UC-HTTPX-SEND:" in traceback or "UC-TIMEOUT-MUTATION:" in traceback
-        ), traceback
-    assert result.wasSuccessful(), "OCR retained parity test failed or unexpectedly succeeded"

litellm-rust/crates/python-bridge/tests/fixtures/retained_http_contract.py

@@ -1,663 +0,0 @@
-import asyncio
-import contextvars
-import gc
-import http.client
-import http.server
-import inspect
-import sys
-import threading
-import weakref
-from copy import deepcopy
-from typing import NamedTuple
-
-native = globals()["native"]
-started = threading.Event()
-release = threading.Event()
-server = None
-server_thread = None
-url = None
-context = contextvars.ContextVar("proof", default="unset")
-
-
-def start_server():
-    global server, server_thread, url
-    requests = []
-
-    class Handler(http.server.BaseHTTPRequestHandler):
-        def log_message(self, *args):
-            pass
-
-        def do_POST(self):
-            body = self.rfile.read(int(self.headers["Content-Length"]))
-            requests.append((self.path, self.headers.get_all("X-Proof"), body))
-            if body == b"hold":
-                started.set()
-                release.wait(5)
-            self.send_response(429)
-            self.send_header("X-Reply", "one")
-            self.send_header("X-Reply", "two")
-            self.send_header("Content-Length", "3")
-            self.end_headers()
-            try:
-                self.wfile.write(b"\x00\xffR")
-            except (BrokenPipeError, ConnectionResetError):
-                pass
-
-    server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler)
-    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
-    server_thread.start()
-    url = "http://127.0.0.1:%s/ocr" % server.server_port
-    return requests
-
-
-requests = start_server()
-
-
-def stop_server():
-    release.set()
-    server.shutdown()
-    server.server_close()
-    server_thread.join(5)
-    assert not server_thread.is_alive()
-
-
-class Graph(dict):
-    pass
-
-
-class Boundary:
-    def __init__(self, *, asynchronous=False, nested=False, failure=None, hold=False):
-        self.asynchronous = asynchronous
-        self.nested = nested
-        self.failure = failure
-        self.hold = hold
-        self.events = []
-        self.thread = threading.get_ident()
-        self.task = asyncio.current_task() if asynchronous else None
-        self.result = object()
-        self.error = LookupError("original callback error")
-
-    def phase(self, name):
-        assert threading.get_ident() == self.thread
-        if self.asynchronous:
-            assert asyncio.current_task() is self.task
-            assert context.get() == ("initial" if name == "prepare" else "prepared")
-        self.events.append(name)
-        if self.failure == name:
-            raise self.error
-        if not self.nested and not self.hold:
-            child = Boundary(nested=True)
-            assert native.ocr_retained(child) is child.result
-            assert child.events == ["prepare", "encode", "finish"]
-
-    def prepare(self):
-        self.phase("prepare")
-        headers = Graph({"X-Proof": "original"})
-        document = object()
-        body = Graph(document=document, alias=document)
-        body["cycle"] = body
-        self.refs = (weakref.ref(headers), weakref.ref(body))
-        self.view = {"headers": headers, "body": body}
-        headers["X-Proof"] = "mutated"
-        self.view["headers"] = {"replacement": True}
-        self.view["body"] = {"replacement": True}
-        return (headers, url, body, None)
-
-    async def aprepare(self):
-        await asyncio.sleep(0)
-        roots = self.prepare()
-        context.set("prepared")
-        return roots
-
-    def encode(self, roots):
-        self.phase("encode")
-        headers, target, body, files = roots
-        assert headers is self.refs[0]() and body is self.refs[1]()
-        assert headers["X-Proof"] == "mutated"
-        assert body["document"] is body["alias"] and body["cycle"] is body
-        assert files is None
-        assert self.view == {"headers": {"replacement": True}, "body": {"replacement": True}}
-        return (
-            target,
-            [(b"X-Proof", b"mutated"), (b"X-Proof", b"duplicate")],
-            b"hold" if self.hold else b"\x00\xffQ",
-            3.0,
-        )
-
-    def finish(self, wire):
-        self.phase("finish")
-        assert type(wire) is tuple and len(wire) == 3
-        status, headers, content = wire
-        assert status == 429
-        assert type(headers) is list
-        assert all(type(pair) is tuple and all(type(v) is bytes for v in pair) for pair in headers)
-        assert [v for k, v in headers if k == b"x-reply"] == [b"one", b"two"]
-        assert type(content) is bytes and content == b"\x00\xffR"
-        assert all(ref() is not None for ref in self.refs)
-        return self.result
-
-    async def afinish(self, wire):
-        await asyncio.sleep(0)
-        return self.finish(wire)
-
-
-def collected(boundary):
-    gc.collect()
-    assert all(ref() is None for ref in boundary.refs)
-    assert boundary.view == {"headers": {"replacement": True}, "body": {"replacement": True}}
-
-
-def run_cold_cache_reentry(filename):
-    """UC-COLD-REENTRY: cold compilation permits nested native calls and subsequent transport."""
-    events = []
-    error = LookupError("cold-cache preparation error")
-
-    class FailingBoundary:
-        async def aprepare(self):
-            raise error
-
-    def invoke_failure():
-        pending = native.aocr_retained(FailingBoundary())
-        observed = None
-        try:
-            pending.send(None)
-        except LookupError as caught:
-            observed = caught
-        finally:
-            pending.close()
-            error.__traceback__ = None
-        assert observed is error
-
-    def audit(event, args):
-        if event == "compile" and args[1] == filename and not events:
-            events.append("entered")
-            print(f"UC-COLD-REENTRY: entering {filename}", flush=True)  # noqa: T201  # diagnose a deadlocked child process
-            invoke_failure()
-            events.append("nested completed")
-
-    async def successful_async():
-        context.set("initial")
-        boundary = Boundary(asynchronous=True, nested=True)
-        assert await native.aocr_retained(boundary) is boundary.result
-        assert boundary.events == ["prepare", "encode", "finish"]
-        collected(boundary)
-
-    try:
-        if filename == "retained_callback.py":
-            native.aocr_retained(object()).close()
-        sys.addaudithook(audit)
-        invoke_failure()
-        events.append("outer completed")
-        assert events == ["entered", "nested completed", "outer completed"]
-        assert requests == []
-        boundary = Boundary(nested=True)
-        assert native.ocr_retained(boundary) is boundary.result
-        assert boundary.events == ["prepare", "encode", "finish"]
-        collected(boundary)
-        asyncio.run(successful_async())
-        assert len(requests) == 2
-    finally:
-        stop_server()
-
-
-def check_error(boundary, error, phase):
-    assert error is boundary.error
-    names = []
-    traceback = error.__traceback__
-    while traceback:
-        names.append(traceback.tb_frame.f_code.co_name)
-        traceback = traceback.tb_next
-    assert phase in names and "phase" in names
-    assert boundary.events == ["prepare", "encode", "finish"][: ["prepare", "encode", "finish"].index(phase) + 1]
-
-
-async def exercise():
-    context.set("initial")
-    boundary = Boundary(asynchronous=True)
-    pending = native.aocr_retained(boundary)
-    assert inspect.iscoroutine(pending)
-    assert boundary.events == []
-    assert await pending is boundary.result
-    assert context.get() == "prepared"
-    assert boundary.events == ["prepare", "encode", "finish"]
-    collected(boundary)
-
-    unused = Boundary(asynchronous=True)
-    ref = weakref.ref(unused)
-    pending = native.aocr_retained(unused)
-    assert unused.events == []
-    del unused
-    assert ref() is not None
-    pending.close()
-    del pending
-    gc.collect()
-    assert ref() is None
-
-    for phase in ("prepare", "encode", "finish"):
-        context.set("initial")
-        boundary = Boundary(asynchronous=True, nested=True, failure=phase)
-        try:
-            await native.aocr_retained(boundary)
-        except LookupError as error:
-            check_error(boundary, error, phase)
-        else:
-            raise AssertionError("callback error was swallowed")
-        boundary.error.__traceback__ = None
-        if phase != "prepare":
-            collected(boundary)
-
-    context.set("initial")
-    boundary = Boundary(asynchronous=True, hold=True)
-
-    async def cancellable():
-        boundary.task = asyncio.current_task()
-        await native.aocr_retained(boundary)
-
-    task = asyncio.create_task(cancellable())
-    assert await asyncio.to_thread(started.wait, 2)
-    gc.collect()
-    assert all(ref() is not None for ref in boundary.refs)
-    task.cancel()
-    try:
-        await task
-    except asyncio.CancelledError:
-        pass
-    else:
-        raise AssertionError("cancellation was swallowed")
-    assert boundary.events == ["prepare", "encode"]
-    del task
-    await asyncio.sleep(0)
-    collected(boundary)
-    release.set()
-
-
-def run_ownership_contract():
-    try:
-        boundary = Boundary()
-        assert native.ocr_retained(boundary) is boundary.result
-        assert boundary.events == ["prepare", "encode", "finish"]
-        collected(boundary)
-        for phase in ("prepare", "encode", "finish"):
-            boundary = Boundary(nested=True, failure=phase)
-            try:
-                native.ocr_retained(boundary)
-            except LookupError as error:
-                check_error(boundary, error, phase)
-            else:
-                raise AssertionError("callback error was swallowed")
-            boundary.error.__traceback__ = None
-            if phase != "prepare":
-                collected(boundary)
-        boundary = Boundary(hold=True)
-        observations = []
-
-        def observe_blocked_transport():
-            received = started.wait(2)
-            gc.collect()
-            observations.append((received, tuple(ref() is not None for ref in boundary.refs)))
-            release.set()
-
-        observer = threading.Thread(target=observe_blocked_transport)
-        observer.start()
-        try:
-            assert native.ocr_retained(boundary) is boundary.result
-        finally:
-            release.set()
-            observer.join(5)
-        assert not observer.is_alive()
-        assert observations == [(True, (True, True))]
-        assert boundary.events == ["prepare", "encode", "finish"]
-        collected(boundary)
-        started.clear()
-        release.clear()
-        asyncio.run(asyncio.wait_for(exercise(), 15))
-        assert requests
-        assert all(
-            path == "/ocr" and headers == ["mutated", "duplicate"] and body in (b"\x00\xffQ", b"hold")
-            for path, headers, body in requests
-        )
-    finally:
-        stop_server()
-
-
-class TimeoutBoundary(Boundary):
-    def __init__(self, *, timeout, url, asynchronous=False):
-        super().__init__(asynchronous=asynchronous)
-        self.timeout = timeout
-        self.url = url
-
-    def prepare(self):
-        return ({}, self.url, {}, None)
-
-    async def aprepare(self):
-        return self.prepare()
-
-    def encode(self, roots):
-        headers, target, body, files = roots
-        return (target, [], b"hold", self.timeout)
-
-    def finish(self, wire):
-        raise AssertionError("client-side failure must not reach finish")
-
-
-def run_error_contract():
-    cases = [
-        ("timeout", url, 0.05),
-        ("refused", "http://127.0.0.1:1/", 1.0),
-    ]
-    try:
-        for name, target, timeout in cases:
-            boundary = TimeoutBoundary(timeout=timeout, url=target)
-            try:
-                native.ocr_retained(boundary)
-            except RuntimeError:
-                pass
-            else:
-                raise AssertionError(f"{name} did not surface as RuntimeError")
-    finally:
-        stop_server()
-
-
-class FrozenObservation(NamedTuple):
-    root_tuple_identity: bool
-    roots_identity: tuple[bool, bool]
-    document_is_caller: bool
-    document_alias: bool
-    logging_envelope_identity: bool
-    logging_roots_identity: tuple[bool, bool]
-    logging_fields: tuple[str, str]
-    retained_fields: tuple[str, str]
-    caller_value: str
-    caught_error_observation: tuple[str, ...]
-    callback_count: int
-    phases: tuple[str, ...]
-    request: tuple[str, tuple[str, ...], bytes]
-    response: tuple[int, tuple[bytes, ...], bytes]
-
-
-class ContractCallback:
-    def __init__(self, scenario, caller):
-        self.scenario = scenario
-        self.caller = caller
-        self.mutate_caller = lambda: caller.__setitem__("value", "closure")
-        self.calls = 0
-        self.error = LookupError("caught callback failure")
-        self.writer = None
-        self.write = threading.Event()
-        self.written = threading.Event()
-
-    def __call__(self, view):
-        self.calls += 1
-        self.view = view
-        self.headers, self.body = view["headers"], view["body"]
-        if self.scenario == "caller_closure":
-            self.mutate_caller()
-        elif self.scenario == "field_replace":
-            document = Graph(value="replacement")
-            view["headers"] = Graph({"X-Proof": "replacement"})
-            view["body"] = Graph(document=document, alias=document)
-        elif self.scenario == "delayed_after_prepare_before_encode_read":
-
-            def writer():
-                if self.write.wait(5):
-                    self.mutate("delayed")
-                    self.written.set()
-
-            self.writer = threading.Thread(target=writer)
-            self.writer.start()
-        else:
-            assert self.scenario in ("retain_mutate", "mutate_then_caught_error_observe")
-            self.mutate("mutated" if self.scenario == "retain_mutate" else "caught")
-            if self.scenario == "mutate_then_caught_error_observe":
-                raise self.error
-
-    def mutate(self, value):
-        self.headers["X-Proof"] = value
-        self.body["document"]["value"] = value
-
-    def before_encode_read(self):
-        if self.writer is not None:
-            self.write.set()
-            assert self.written.wait(2), "scheduled mutation did not complete before encoder read"
-            self.writer.join(2)
-            assert not self.writer.is_alive()
-
-    def close(self):
-        self.write.set()
-        if self.writer is not None:
-            self.writer.join(5)
-            assert not self.writer.is_alive()
-        self.error.__traceback__ = None
-
-
-class TransportContractBinding:
-    def __init__(self, caller, callback):
-        self.caller = caller
-        self.callback = callback
-        self.phases = []
-        self.caught = ()
-
-    def prepare(self):
-        self.phases.append("prepare")
-        headers = Graph({"X-Proof": "original"})
-        body = Graph(document=self.caller, alias=self.caller)
-        self.view = {"headers": headers, "body": body}
-        self.prepared_roots = (headers, url, body, None)
-        try:
-            self.callback(self.view)
-        except LookupError as error:
-            if error is not self.callback.error:
-                raise
-            self.caught = (str(error), headers["X-Proof"], body["document"]["value"])
-        return self.prepared_roots
-
-    async def aprepare(self):
-        await asyncio.sleep(0)
-        return self.prepare()
-
-    def encode(self, roots):
-        self.callback.before_encode_read()
-        self.phases.append("encode")
-        headers, target, body, files = roots
-        assert files is None
-        self.root_tuple_identity = roots is self.prepared_roots
-        self.roots_identity = (headers is self.callback.headers, body is self.callback.body)
-        self.document_is_caller = body["document"] is self.callback.caller
-        self.document_alias = body["document"] is body["alias"]
-        return (
-            target,
-            [(b"X-Proof", headers["X-Proof"].encode())],
-            b"\x00" + body["document"]["value"].encode() + b"\xff",
-            3.0,
-        )
-
-    def finish(self, wire):
-        self.phases.append("finish")
-        status, headers, content = wire
-        return (status, tuple(value for key, value in headers if key == b"x-reply"), content)
-
-    async def afinish(self, wire):
-        await asyncio.sleep(0)
-        return self.finish(wire)
-
-
-class BindingVariant:
-    def __init__(self, binding, variant):
-        self.binding = binding
-        self.variant = variant
-
-    def before_prepare(self):
-        if self.variant == "copy_caller_inputs_before_prepare":
-            self.binding.caller = deepcopy(self.binding.caller)
-
-    def at_prepare_return(self, roots):
-        if self.variant == "copy_roots_after_prepare":
-            return deepcopy(roots)
-        if self.variant == "reconstruct_root_tuple":
-            return tuple(list(roots))
-        if self.variant == "reconstruct_logging_envelope":
-            self.binding.view = dict(self.binding.view)
-        return roots
-
-    def prepare(self):
-        self.before_prepare()
-        return self.at_prepare_return(self.binding.prepare())
-
-    async def aprepare(self):
-        self.before_prepare()
-        return self.at_prepare_return(await self.binding.aprepare())
-
-    def encode(self, roots):
-        if self.variant == "encode_logging_replacements":
-            headers, target, body, files = roots
-            return self.binding.encode((self.binding.view["headers"], target, self.binding.view["body"], files))
-        return self.binding.encode(roots)
-
-    def finish(self, wire):
-        return self.binding.finish(wire)
-
-    async def afinish(self, wire):
-        return await self.binding.afinish(wire)
-
-
-def reference_post(encoded):
-    target, headers, body, timeout = encoded
-    assert target == url
-    connection = http.client.HTTPConnection("127.0.0.1", server.server_port, timeout=timeout)
-    try:
-        connection.putrequest("POST", "/ocr")
-        for key, value in headers:
-            connection.putheader(key.decode("ascii"), value.decode("ascii"))
-        connection.putheader("Content-Length", str(len(body)))
-        connection.endheaders(body)
-        response = connection.getresponse()
-        return (
-            response.status,
-            [(key.lower().encode(), value.encode()) for key, value in response.getheaders()],
-            response.read(),
-        )
-    finally:
-        connection.close()
-
-
-async def reference_async(binding):
-    roots = await binding.aprepare()
-    wire = await asyncio.to_thread(reference_post, binding.encode(roots))
-    return await binding.afinish(wire)
-
-
-def original_observation(value, *, header=None, logging=None, caught=()):
-    fields = (header or value, value)
-    return FrozenObservation(
-        True,
-        (True, True),
-        True,
-        True,
-        True,
-        (False, False) if logging else (True, True),
-        logging or fields,
-        fields,
-        value,
-        caught,
-        1,
-        ("prepare", "encode", "finish"),
-        ("/ocr", (fields[0],), b"\x00" + value.encode() + b"\xff"),
-        (429, (b"one", b"two"), b"\x00\xffR"),
-    )
-
-
-ORIGINAL_OBSERVATIONS = {
-    "retain_mutate": original_observation("mutated"),
-    "caller_closure": original_observation("closure", header="original"),
-    "delayed_after_prepare_before_encode_read": original_observation("delayed"),
-    "field_replace": original_observation("original", logging=("replacement", "replacement")),
-    "mutate_then_caught_error_observe": original_observation(
-        "caught", caught=("caught callback failure", "caught", "caught")
-    ),
-}
-
-VARIANT_OBSERVATIONS = {
-    ("retain_mutate", "copy_caller_inputs_before_prepare"): {
-        "caller_value": "original",
-        "document_is_caller": False,
-    },
-    ("caller_closure", "copy_caller_inputs_before_prepare"): {
-        "document_is_caller": False,
-        "logging_fields": ("original", "original"),
-        "retained_fields": ("original", "original"),
-        "request": ("/ocr", ("original",), b"\x00original\xff"),
-    },
-    ("retain_mutate", "copy_roots_after_prepare"): {
-        "root_tuple_identity": False,
-        "roots_identity": (False, False),
-        "document_is_caller": False,
-    },
-    ("delayed_after_prepare_before_encode_read", "copy_roots_after_prepare"): {
-        "root_tuple_identity": False,
-        "roots_identity": (False, False),
-        "document_is_caller": False,
-        "request": ("/ocr", ("original",), b"\x00original\xff"),
-    },
-    ("mutate_then_caught_error_observe", "copy_roots_after_prepare"): {
-        "root_tuple_identity": False,
-        "roots_identity": (False, False),
-        "document_is_caller": False,
-    },
-    ("field_replace", "encode_logging_replacements"): {
-        "root_tuple_identity": False,
-        "roots_identity": (False, False),
-        "document_is_caller": False,
-        "request": ("/ocr", ("replacement",), b"\x00replacement\xff"),
-    },
-    ("retain_mutate", "reconstruct_logging_envelope"): {"logging_envelope_identity": False},
-    ("field_replace", "reconstruct_logging_envelope"): {"logging_envelope_identity": False},
-    ("retain_mutate", "reconstruct_root_tuple"): {"root_tuple_identity": False},
-    ("delayed_after_prepare_before_encode_read", "reconstruct_root_tuple"): {"root_tuple_identity": False},
-    ("field_replace", "reconstruct_root_tuple"): {"root_tuple_identity": False},
-}
-
-
-def run_comparison_contract(scenario, variant):
-    try:
-        expected = ORIGINAL_OBSERVATIONS[scenario]
-        if variant != "original":
-            expected = expected._replace(**VARIANT_OBSERVATIONS[(scenario, variant)])
-        for mode in ("reference-sync", "reference-async", "native-sync", "native-async"):
-            caller = Graph(value="original")
-            callback = ContractCallback(scenario, caller)
-            binding = TransportContractBinding(caller, callback)
-            selected = binding if variant == "original" else BindingVariant(binding, variant)
-            offset = len(requests)
-            try:
-                if mode == "reference-sync":
-                    roots = selected.prepare()
-                    response = selected.finish(reference_post(selected.encode(roots)))
-                elif mode == "reference-async":
-                    response = asyncio.run(reference_async(selected))
-                elif mode == "native-sync":
-                    response = native.ocr_retained(selected)
-                else:
-                    response = asyncio.run(native.aocr_retained(selected))
-                assert len(requests) == offset + 1, (scenario, variant, mode, requests[offset:])
-                path, headers, body = requests[offset]
-                observed = FrozenObservation(
-                    binding.root_tuple_identity,
-                    binding.roots_identity,
-                    binding.document_is_caller,
-                    binding.document_alias,
-                    binding.view is callback.view,
-                    (binding.view["headers"] is callback.headers, binding.view["body"] is callback.body),
-                    (binding.view["headers"]["X-Proof"], binding.view["body"]["document"]["value"]),
-                    (callback.headers["X-Proof"], callback.body["document"]["value"]),
-                    caller["value"],
-                    binding.caught,
-                    callback.calls,
-                    tuple(binding.phases),
-                    (path, tuple(headers), body),
-                    response,
-                )
-                assert observed == expected, (scenario, variant, mode, observed, expected)
-            finally:
-                callback.close()
-    finally:
-        stop_server()

litellm-rust/crates/python-bridge/tests/ocr_retained.rs

@@ -1,51 +0,0 @@
-use pyo3::prelude::*;
-use rstest::rstest;
-use serial_test::serial;
-
-#[path = "support/mod.rs"]
-mod support;
-
-use support::native::{native_globals, run_fixture};
-
-#[rstest]
-#[case::differential_callbacks_wire("differential_callbacks_wire")]
-#[case::known_gap_send_time_auth_and_request_hook("known_gap_send_time_auth_and_request_hook")]
-#[case::known_gap_custom_transport("known_gap_custom_transport")]
-#[case::known_gap_pre_call_timeout_mutation("known_gap_pre_call_timeout_mutation")]
-#[case::negative_control_copied_caller_document("negative_control_copied_caller_document")]
-#[case::negative_control_rebound_logging_body("negative_control_rebound_logging_body")]
-#[case::negative_control_rebound_logging_headers("negative_control_rebound_logging_headers")]
-#[case::differential_callback_retained_mutation_after_post_received(
-    "differential_callback_retained_mutation_after_post_received"
-)]
-#[case::public_rust_dispatch_wire_fallback_and_escaping_base_exception(
-    "public_rust_dispatch_wire_fallback_and_escaping_base_exception"
-)]
-#[case::collection_after_success_and_failures("collection_after_success_and_failures")]
-#[case::callback_retained_graph_remains_usable_then_collects(
-    "callback_retained_graph_remains_usable_then_collects"
-)]
-#[case::collection_after_cancellation_during_blocked_transport(
-    "collection_after_cancellation_during_blocked_transport"
-)]
-#[ignore = "requires repo Python"]
-#[serial(python_interpreter)]
-fn retained_real_production_boundary_differential_and_lifecycle(
-    #[case] scenario: &str,
-) -> PyResult<()> {
-    Python::initialize();
-    Python::attach(|py| {
-        let globals = native_globals(py)?;
-        run_fixture(
-            py,
-            &globals,
-            include_str!("fixtures/ocr_retained.py"),
-            concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/tests/fixtures/ocr_retained.py"
-            ),
-        )?;
-        globals.get_item("run_case")?.unwrap().call1((scenario,))?;
-        Ok(())
-    })
-}

litellm-rust/crates/python-bridge/tests/retained_http.rs

@@ -1,163 +0,0 @@
-use std::process::Command;
-use std::time::{Duration, Instant};
-
-use pyo3::prelude::*;
-use rstest::rstest;
-use serial_test::serial;
-
-#[path = "support/mod.rs"]
-mod support;
-
-use support::native::{native_globals, run_fixture};
-
-#[test]
-fn uc_cold_reentry_execution_cache() -> PyResult<()> {
-    cold_cache_reentry("retained_execution.py", "uc_cold_reentry_execution_cache")
-}
-
-#[test]
-fn uc_cold_reentry_callback_cache() -> PyResult<()> {
-    cold_cache_reentry("retained_callback.py", "uc_cold_reentry_callback_cache")
-}
-
-fn cold_cache_reentry(filename: &str, test: &str) -> PyResult<()> {
-    if std::env::var("LITELLM_COLD_REENTRY_CHILD").as_deref() != Ok(filename) {
-        let mut child = Command::new(std::env::current_exe().unwrap())
-            .args(["--exact", test, "--nocapture"])
-            .env("LITELLM_COLD_REENTRY_CHILD", filename)
-            .spawn()
-            .unwrap();
-        let deadline = Instant::now() + Duration::from_secs(15);
-        loop {
-            if let Some(status) = child.try_wait().unwrap() {
-                assert!(
-                    status.success(),
-                    "{filename} reentry child failed: {status}"
-                );
-                return Ok(());
-            }
-            if Instant::now() >= deadline {
-                child.kill().unwrap();
-                child.wait().unwrap();
-                panic!("{filename} reentry did not complete within 15 seconds");
-            }
-            std::thread::sleep(Duration::from_millis(10));
-        }
-    }
-
-    Python::initialize();
-    Python::attach(|py| {
-        let globals = native_globals(py)?;
-        run_fixture(
-            py,
-            &globals,
-            include_str!("fixtures/retained_http_contract.py"),
-            concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/tests/fixtures/retained_http_contract.py"
-            ),
-        )?;
-        globals
-            .get_item("run_cold_cache_reentry")?
-            .unwrap()
-            .call1((filename,))?;
-        Ok(())
-    })
-}
-
-#[test]
-#[serial(python_interpreter)]
-fn retained_routes_preserve_callbacks_context_wire_and_ownership() -> PyResult<()> {
-    Python::initialize();
-    Python::attach(|py| {
-        let globals = native_globals(py)?;
-        run_fixture(
-            py,
-            &globals,
-            include_str!("fixtures/retained_http_contract.py"),
-            concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/tests/fixtures/retained_http_contract.py"
-            ),
-        )?;
-        globals
-            .get_item("run_ownership_contract")?
-            .unwrap()
-            .call0()?;
-        Ok(())
-    })
-}
-
-#[test]
-#[serial(python_interpreter)]
-fn retained_routes_surface_transport_failures_as_runtime_error() -> PyResult<()> {
-    Python::initialize();
-    Python::attach(|py| {
-        let globals = native_globals(py)?;
-        run_fixture(
-            py,
-            &globals,
-            include_str!("fixtures/retained_http_contract.py"),
-            concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/tests/fixtures/retained_http_contract.py"
-            ),
-        )?;
-        globals.get_item("run_error_contract")?.unwrap().call0()?;
-        Ok(())
-    })
-}
-
-#[rstest]
-#[case::original_retain_mutate("retain_mutate", "original")]
-#[case::original_caller_closure("caller_closure", "original")]
-#[case::original_delayed_after_prepare_before_encode_read(
-    "delayed_after_prepare_before_encode_read",
-    "original"
-)]
-#[case::original_field_replace("field_replace", "original")]
-#[case::original_mutate_then_caught_error_observe("mutate_then_caught_error_observe", "original")]
-#[case::copy_caller_inputs_retain_mutate("retain_mutate", "copy_caller_inputs_before_prepare")]
-#[case::copy_caller_inputs_caller_closure("caller_closure", "copy_caller_inputs_before_prepare")]
-#[case::copy_roots_retain_mutate("retain_mutate", "copy_roots_after_prepare")]
-#[case::copy_roots_delayed_after_prepare_before_encode_read(
-    "delayed_after_prepare_before_encode_read",
-    "copy_roots_after_prepare"
-)]
-#[case::copy_roots_mutate_then_caught_error_observe(
-    "mutate_then_caught_error_observe",
-    "copy_roots_after_prepare"
-)]
-#[case::encode_logging_replacements_field_replace("field_replace", "encode_logging_replacements")]
-#[case::reconstruct_logging_envelope_retain_mutate("retain_mutate", "reconstruct_logging_envelope")]
-#[case::reconstruct_logging_envelope_field_replace("field_replace", "reconstruct_logging_envelope")]
-#[case::reconstruct_root_tuple_retain_mutate("retain_mutate", "reconstruct_root_tuple")]
-#[case::reconstruct_root_tuple_delayed_after_prepare_before_encode_read(
-    "delayed_after_prepare_before_encode_read",
-    "reconstruct_root_tuple"
-)]
-#[case::reconstruct_root_tuple_field_replace("field_replace", "reconstruct_root_tuple")]
-#[serial(python_interpreter)]
-fn retained_production_transport_comparison_matrix(
-    #[case] scenario: &str,
-    #[case] variant: &str,
-) -> PyResult<()> {
-    Python::initialize();
-    Python::attach(|py| {
-        let globals = native_globals(py)?;
-        run_fixture(
-            py,
-            &globals,
-            include_str!("fixtures/retained_http_contract.py"),
-            concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/tests/fixtures/retained_http_contract.py"
-            ),
-        )?;
-        globals
-            .get_item("run_comparison_contract")?
-            .unwrap()
-            .call1((scenario, variant))?;
-        Ok(())
-    })
-}

litellm-rust/crates/python-bridge/tests/support/mod.rs

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

litellm-rust/crates/python-bridge/tests/support/native.rs

@@ -1,21 +0,0 @@
-use pyo3::prelude::*;
-use pyo3::types::PyDict;
-
-pub fn native_globals(py: Python<'_>) -> PyResult<Bound<'_, PyDict>> {
-    let module = pyo3::wrap_pymodule!(_native::_native)(py).into_bound(py);
-    let globals = PyDict::new(py);
-    globals.set_item("native", &module)?;
-    Ok(globals)
-}
-
-pub fn run_fixture(
-    py: Python<'_>,
-    globals: &Bound<'_, PyDict>,
-    source: &str,
-    filename: &str,
-) -> PyResult<()> {
-    let builtins = py.import("builtins")?;
-    let code = builtins.call_method1("compile", (source, filename, "exec"))?;
-    builtins.call_method1("exec", (code, globals))?;
-    Ok(())
-}

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

@@ -1,3 +1,6 @@
+use std::process::Command;
+use std::time::{Duration, Instant};
+
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
 use rstest::{fixture, rstest};
@@ -11,6 +14,45 @@ mod support;
 
 use support::python::{InitializedPython, initialized_python, run_fixture};
 
+#[test]
+fn cold_awaited_adapter_initialization_allows_reentry() -> PyResult<()> {
+    let test = "cold_awaited_adapter_initialization_allows_reentry";
+    let child_env = "LITELLM_INTEROP_COLD_REENTRY_CHILD";
+    if std::env::var(child_env).as_deref() != Ok(test) {
+        let mut child = Command::new(std::env::current_exe().unwrap())
+            .args(["--exact", test, "--nocapture"])
+            .env(child_env, test)
+            .spawn()
+            .unwrap();
+        let deadline = Instant::now() + Duration::from_secs(15);
+        loop {
+            if let Some(status) = child.try_wait().unwrap() {
+                assert!(
+                    status.success(),
+                    "awaited adapter reentry child failed: {status}"
+                );
+                return Ok(());
+            }
+            if Instant::now() >= deadline {
+                child.kill().unwrap();
+                child.wait().unwrap();
+                panic!("awaited adapter reentry did not complete within 15 seconds");
+            }
+            std::thread::sleep(Duration::from_millis(10));
+        }
+    }
+
+    let globals = scenario_scope(&initialized_python());
+    Python::attach(|py| {
+        globals
+            .bind(py)
+            .get_item("cold_awaited_adapter_reentry")?
+            .unwrap()
+            .call1((globals.bind(py).get_item("factory")?.unwrap(),))?;
+        Ok(())
+    })
+}
+
 #[fixture]
 fn scenario_scope(initialized_python: &InitializedPython) -> Py<PyDict> {
     let _ = initialized_python;

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

@@ -2,7 +2,9 @@ import asyncio
 import contextvars
 import copy
 import gc
+import inspect
 import json
+import sys
 import threading
 import weakref
 from unittest import TestCase
@@ -18,6 +20,65 @@ class Value:
     pass
 
 
+def cold_awaited_adapter_reentry(factory):
+    events = []
+    compilations = []
+
+    def invoke_failure():
+        error = LookupError("cold-cache callback error")
+
+        async def failing():
+            raise error
+
+        owner = factory.prepare(failing, (), awaited=True)
+        pending = owner.invoke()
+        try:
+            assert inspect.getcoroutinestate(pending) == inspect.CORO_CREATED
+            with TestCase().assertRaises(LookupError) as caught:
+                pending.send(None)
+            assert caught.exception is error
+        finally:
+            pending.close()
+            owner.close()
+            error.__traceback__ = None
+
+    def audit(event, args):
+        if event != "compile" or args[1] != "retained_callback.py":
+            return
+        compilations.append(args[1])
+        if len(compilations) == 1:
+            events.append("entered")
+            invoke_failure()
+            events.append("nested completed")
+
+    sys.addaudithook(audit)
+    invoke_failure()
+    events.append("outer completed")
+    assert events == ["entered", "nested completed", "outer completed"]
+    assert len(compilations) == 2
+    assert factory.live == 0
+
+    result = object()
+    direct = factory.prepare(lambda value: value, (result,))
+    try:
+        assert direct.invoke() is result
+    finally:
+        direct.close()
+
+    async def successful(value, *, alias):
+        assert value is alias
+        await checkpoint()
+        return value
+
+    owner = factory.prepare(successful, (result,), {"alias": result}, awaited=True)
+    try:
+        assert asyncio.run(owner.invoke()) is result
+    finally:
+        owner.close()
+    assert len(compilations) == 2
+    assert factory.live == 0
+
+
 class ReferenceFactory:
     def __init__(self):
         self.live = 0

litellm/ocr/main.py

@@ -31,7 +31,6 @@ 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.rust_bridge import ocr_retained as rust_ocr_retained_bridge
 from litellm.types.router import GenericLiteLLMParams
 from litellm.utils import ProviderConfigManager, client
 
@@ -308,44 +307,10 @@ def _map_rust_ocr_error(
         status_code=status or 500,
         headers={},  # mutable-ok: provider error factories require a concrete header dict
     )
-
-
-def _retained_ocr_boundary(prepared: _PreparedOCRRequest) -> rust_ocr_retained_bridge.OCRRetainedBoundary:
-    return rust_ocr_retained_bridge.OCRRetainedBoundary(
-        handler=base_llm_http_handler,
-        model=prepared.model,
-        document=prepared.document,
-        optional_params=prepared.optional_params,
-        logging_obj=prepared.litellm_logging_obj,
-        api_key=prepared.api_key,
-        api_base=prepared.api_base,
-        headers=prepared.extra_headers,
-        provider_config=prepared.provider_config,
-        litellm_params=prepared.litellm_params,
-        custom_llm_provider=prepared.custom_llm_provider,
-        timeout=prepared.effective_timeout,
-    )
-
-
 def _run_rust_ocr(
     prepared_request: _PreparedOCRRequest,
     resolve_api_key: Callable[[str], str | None],
-    *,
-    load_retained: Callable[[], rust_ocr_retained_bridge.RustOCRRetained | None] = (
-        rust_ocr_retained_bridge.load_rust_ocr_retained
-    ),
 ) -> OCRResponse | None:
-    if prepared_request.custom_llm_provider == "mistral" and not rust_ocr_bridge.has_rust_ocr_override():
-        retained: Final = load_retained()
-        if (
-            retained is None
-            or rust_ocr_retained_bridge.retained_timeout_seconds(prepared_request.effective_timeout) is None
-        ):
-            return None
-        retained_response: Final = retained(_retained_ocr_boundary(prepared_request))
-        if retained_response is None:
-            raise ValueError("Retained OCR returned no response after preparation")
-        return retained_response
     if rust_ocr_bridge.load_rust_ocr() is None:
         return None
     prepared: Final = _prepare_rust_ocr_call(
@@ -373,24 +338,7 @@ def _run_rust_ocr(
 async def _run_rust_aocr(
     prepared_request: _PreparedOCRRequest,
     resolve_api_key: Callable[[str], str | None],
-    *,
-    load_retained: Callable[[], rust_ocr_retained_bridge.RustAOCRRetained | None] = (
-        rust_ocr_retained_bridge.load_rust_aocr_retained
-    ),
 ) -> OCRResponse | None:
-    if prepared_request.custom_llm_provider == "mistral" and not rust_ocr_bridge.has_rust_ocr_override(
-        asynchronous=True
-    ):
-        retained: Final = load_retained()
-        if (
-            retained is None
-            or rust_ocr_retained_bridge.retained_timeout_seconds(prepared_request.effective_timeout) is None
-        ):
-            return None
-        retained_response: Final = await retained(_retained_ocr_boundary(prepared_request))
-        if retained_response is None:
-            raise ValueError("Retained OCR returned no response after preparation")
-        return retained_response
     if rust_ocr_bridge.load_rust_aocr() is None:
         return None
     prepared: Final = _prepare_rust_ocr_call(

litellm/rust_bridge/ocr.py

@@ -53,10 +53,6 @@ _OCR: Final = NativeBinding("ocr", validate=_as_ocr)
 _AOCR: Final = NativeBinding("aocr", validate=_as_aocr)
 
 
-def has_rust_ocr_override(*, asynchronous: bool = False) -> bool:
-    return _rust_aocr_impl is not None if asynchronous else _rust_ocr_impl is not None
-
-
 def load_rust_ocr() -> RustOcr | None:
     return _OCR.load()
 

litellm/rust_bridge/ocr_retained.py

@@ -1,146 +0,0 @@
-"""Python preparation, encoding, and response transforms for retained OCR calls."""
-
-from __future__ import annotations
-
-import math
-from collections.abc import Awaitable, Callable
-from dataclasses import dataclass, field
-from typing import Final, cast  # noqa: TID251  # native callables and legacy header types require boundary casts
-
-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.timeouts import timeout_to_seconds
-
-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 retained_timeout_seconds(timeout: float | httpx.Timeout) -> float | None:
-    seconds: Final = timeout_to_seconds(timeout)
-    return seconds if seconds is not None and math.isfinite(seconds) and seconds > 0 else None
-
-
-@dataclass(kw_only=True, slots=True)
-class OCRRetainedBoundary:
-    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 = retained_timeout_seconds(self.timeout)
-        if seconds is None:
-            raise ValueError("Retained OCR requires a positive finite read timeout")
-        if self.client is None:
-            raise RuntimeError("Retained OCR 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("Retained OCR 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,
-        )
-
-
-RustOCRRetained = Callable[[OCRRetainedBoundary], OCRResponse]
-RustAOCRRetained = Callable[[OCRRetainedBoundary], Awaitable[OCRResponse]]
-
-
-def load_rust_ocr_retained() -> RustOCRRetained | None:
-    from litellm.rust_bridge import get_native_bridge
-
-    native: Final = get_native_bridge()
-    return cast(RustOCRRetained | None, getattr(native, "ocr_retained", None))
-
-
-def load_rust_aocr_retained() -> RustAOCRRetained | None:
-    from litellm.rust_bridge import get_native_bridge
-
-    native: Final = get_native_bridge()
-    return cast(RustAOCRRetained | None, getattr(native, "aocr_retained", None))

tests/test_litellm/ocr/test_rust_bridge.py

@@ -816,14 +816,23 @@ def test_ocr_passes_default_request_timeout_to_rust(fake_bridge):
     assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout)
 
 
-@pytest.mark.parametrize("rust_enabled", [False, True], ids=["disabled", "unavailable"])
-def test_ocr_uses_python_when_native_execution_is_unavailable(monkeypatch, rust_enabled):
-    def load_native():
-        assert rust_enabled, "disabled OCR must not load the native module"
-        return None
+def test_ocr_does_not_route_to_rust_when_disabled():
+    """With the flag off, the bridge must not be consulted even if an impl exists."""
+    bridge = RecordingBridge()
+    litellm.rust(False)
+    rust_bridge.set_rust_ocr(ocr=bridge)
 
-    monkeypatch.setattr("litellm.rust_bridge.get_native_bridge", load_native)
-    litellm.rust(rust_enabled)
+    assert rust_bridge.rust_ocr_enabled() is False
+    # The impl stays available for injection, but the disabled flag gates usage,
+    # so ocr() never reaches the Rust path (asserted via the enabled-path test).
+    assert bridge.calls == []
+
+
+def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch):
+    """Rust enabled but no bridge available (no injected impl, no compiled wheel):
+    ocr() must degrade to the Python HTTP handler instead of raising."""
+    monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None)
+    litellm.rust(True)  # enabled, but load_rust_ocr() returns None in CI
 
     captured = {}
 

tests/test_litellm/rust_bridge/native_route_wheel_test.py

@@ -10,18 +10,14 @@ import sys
 import tempfile
 import threading
 import zipfile
-from collections import Counter
 from http.client import HTTPMessage
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from pathlib import Path
-from queue import SimpleQueue
 from socket import socket as Socket
-from types import FrameType
 from typing import Final
 
 REQUEST_STARTED: Final = threading.Event()
 REQUEST_CANCELLED: Final = threading.Event()
-PUBLIC_OCR_REQUESTS: Final[SimpleQueue[str]] = SimpleQueue()
 
 ANTHROPIC_RESPONSE: Final = (
     b'{"id":"msg_native","type":"message","role":"assistant",'
@@ -36,24 +32,10 @@ class NativeRouteHandler(BaseHTTPRequestHandler):
 
     def do_POST(self) -> None:
         content_length: Final = int(self.headers.get("content-length", "0"))
-        wire_body: Final = self.rfile.read(content_length)
-        body: Final = json.loads(wire_body)
+        body: Final = json.loads(self.rfile.read(content_length))
         route: Final = self.headers.get("x-test-route")
         outcome: Final = self.headers.get("x-test-outcome")
-        public_case: Final = self.headers.get("x-test-public-case")
-        if public_case is not None:
-            assert self.headers.get("x-test-callback") == public_case
-            assert self.headers.get("x-test-view-only") is None
-            assert self.headers.get_all("x-test-callback") == [public_case]
-            assert wire_body == (
-                b'{"model":"mistral-ocr-latest","document":{"type":"document_url",'
-                b'"document_url":"https://example.com/document.pdf"},"include_image_base64":true,'
-                b'"document_alias":{"type":"document_url","document_url":"https://example.com/document.pdf"},'
-                b'"callback_mutation":"observed"}'
-            ), wire_body
         assert_native_request(route, outcome, self.path, self.headers, body)
-        if public_case is not None:
-            PUBLIC_OCR_REQUESTS.put(public_case)
         if outcome == "hang":
             REQUEST_STARTED.set()
             self.connection.settimeout(5)
@@ -245,7 +227,12 @@ 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:
@@ -262,149 +249,6 @@ def exercise_routes(native_path: Path, api_base: str) -> object:
     return native
 
 
-def exercise_public_ocr(install_root: Path, api_base: str, case: str) -> int:
-    import litellm
-    from litellm.integrations.custom_logger import CustomLogger
-    from litellm.llms.base_llm.ocr.transformation import OCRResponse
-    from litellm.rust_bridge import get_native_bridge
-    from litellm.rust_bridge.ocr_retained import OCRRetainedBoundary
-
-    assert case in {"ocr", "aocr"}
-    assert Path(litellm.__file__).resolve().is_relative_to(install_root.resolve())
-    native: Final = get_native_bridge()
-    assert native is not None
-    assert Path(native.__file__).resolve().is_relative_to(install_root.resolve())
-    retained: Final = getattr(native, f"{case}_retained")
-    observed: Final[Counter[str]] = Counter()
-    roots: Final[dict[str, dict[str, object]]] = {}
-    phases: Final[list[str]] = []
-    document: Final = {"type": "document_url", "document_url": "https://example.com/before-callback.pdf"}
-    replacement_body: Final = {"replacement": True}
-    replacement_headers: Final = {"x-test-callback": "must-not-send", "x-test-view-only": "not-on-wire"}
-
-    def observe(frame: FrameType, event: str, arg: object) -> None:
-        if event == "c_call" and arg is retained:
-            observed["retained"] += 1
-        if event == "call" and frame.f_code is OCRRetainedBoundary.encode.__code__:
-            observed["encode"] += 1
-
-    class RetainedLogger(CustomLogger):
-        def __init__(self, phase: str) -> None:
-            super().__init__()
-            self.phase = phase
-            self.calls = 0
-
-        def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> dict[str, object]:
-            self.calls += 1
-            additional_args: Final = kwargs["additional_args"]
-            assert isinstance(additional_args, dict)
-            if self.phase == "retain":
-                assert phases == []
-                headers: Final = additional_args["headers"]
-                body: Final = additional_args["complete_input_dict"]
-                assert isinstance(headers, dict) and isinstance(body, dict)
-                assert body["model"] == "mistral-ocr-latest"
-                assert body["include_image_base64"] is False
-                assert body["document"] is document
-                assert headers["x-test-callback"] == "before-callback"
-                roots.update(body=body, headers=headers, view=additional_args)
-                body["document_alias"] = document
-                additional_args["complete_input_dict"] = replacement_body
-                additional_args["headers"] = replacement_headers
-            else:
-                assert additional_args is roots["view"]
-                assert additional_args["complete_input_dict"] is replacement_body
-                assert additional_args["headers"] is replacement_headers
-                assert roots["body"]["document_alias"] is document
-                assert roots["body"]["document"] is document
-                if self.phase == "mutate":
-                    assert phases == ["retain"]
-                    roots["headers"]["x-test-callback"] = case
-                    roots["body"]["include_image_base64"] = True
-                    document["document_url"] = "https://example.com/document.pdf"
-                    roots["body"]["callback_mutation"] = "observed"
-                else:
-                    assert self.phase == "observe"
-                    assert phases == ["retain", "mutate"]
-                    assert roots["headers"]["x-test-callback"] == case
-                    assert roots["body"]["include_image_base64"] is True
-                    assert document["document_url"] == "https://example.com/document.pdf"
-                    assert roots["body"]["callback_mutation"] == "observed"
-                    assert replacement_body == {"replacement": True}
-                    assert replacement_headers == {
-                        "x-test-callback": "must-not-send",
-                        "x-test-view-only": "not-on-wire",
-                    }
-            phases.append(self.phase)
-            return {"headers": {"x-test-callback": "ignored-return"}, "complete_input_dict": {"invalid": object()}}
-
-    callbacks: Final = tuple(RetainedLogger(phase) for phase in ("retain", "mutate", "observe"))
-    kwargs: Final = {
-        "model": "mistral/mistral-ocr-latest",
-        "document": document,
-        "api_base": api_base,
-        "api_key": "sk-native",
-        "extra_headers": {
-            "x-test-route": "ocr",
-            "x-test-outcome": "success",
-            "x-test-public-case": case,
-            "x-test-callback": "before-callback",
-        },
-        "include_image_base64": False,
-        "callbacks": list(callbacks),
-        "rust": True,
-        "timeout": 3.0,
-        "num_retries": 0,
-    }
-    previous_profile: Final = sys.getprofile()
-    sys.setprofile(observe)
-    try:
-        response: Final = asyncio.run(litellm.aocr(**kwargs)) if case == "aocr" else litellm.ocr(**kwargs)
-    finally:
-        sys.setprofile(previous_profile)
-    assert isinstance(response, OCRResponse)
-    assert_success("ocr", response.model_dump())
-    assert response.model == "mistral-ocr-latest"
-    assert tuple(callback.calls for callback in callbacks) == (1, 1, 1)
-    assert phases == ["retain", "mutate", "observe"], phases
-    assert observed == {"retained": 1, "encode": 1}, observed
-    assert roots["view"]["complete_input_dict"] is replacement_body
-    assert roots["view"]["headers"] is replacement_headers
-    assert replacement_body == {"replacement": True}
-    assert replacement_headers == {"x-test-callback": "must-not-send", "x-test-view-only": "not-on-wire"}
-    assert roots["body"] == {
-        "model": "mistral-ocr-latest",
-        "document": document,
-        "include_image_base64": True,
-        "document_alias": document,
-        "callback_mutation": "observed",
-    }
-    document["document_url"] = "https://example.com/after-return.pdf"
-    roots["headers"]["x-after-return"] = "usable"
-    assert roots["body"]["document"] is roots["body"]["document_alias"] is document
-    assert roots["headers"]["x-after-return"] == "usable"
-    assert replacement_body == {"replacement": True}
-    assert replacement_headers == {"x-test-callback": "must-not-send", "x-test-view-only": "not-on-wire"}
-    return 0
-
-
-def verify_public_ocr(wheel: Path, wheel_root: Path, api_base: str) -> None:
-    install_root: Final = wheel_root / "sdk-venv"
-    python: Final = install_root / "bin" / "python"
-    subprocess.run(("uv", "venv", "--python", sys.executable, str(install_root)), check=True)
-    subprocess.run(("uv", "pip", "install", "--python", str(python), str(wheel.resolve())), check=True)
-    for case in ("ocr", "aocr"):
-        subprocess.run(
-            (str(python), "-I", str(Path(__file__).resolve()), "public-ocr", str(install_root), api_base, case),
-            cwd=install_root,
-            env=os.environ | {"LITELLM_LOCAL_MODEL_COST_MAP": "True", "NO_PROXY": "127.0.0.1", "no_proxy": "127.0.0.1"},
-            check=True,
-            timeout=60,
-        )
-        assert PUBLIC_OCR_REQUESTS.get_nowait() == case
-        assert PUBLIC_OCR_REQUESTS.empty(), f"{case} sent more than one upstream request"
-
-
 def exercise_signal(native: object, api_base: str) -> int:
     try:
         native.messages(
@@ -476,7 +320,6 @@ def verify_wheel(wheel: Path) -> int:
         api_base: Final = f"http://127.0.0.1:{server.server_address[1]}"
         try:
             verify_sigint(native_path, api_base)
-            verify_public_ocr(wheel, wheel_root, api_base)
         finally:
             server.shutdown()
             server.server_close()
@@ -490,8 +333,6 @@ def main() -> int:
     if len(sys.argv) == 4 and sys.argv[1] == "child":
         native: Final = exercise_routes(Path(sys.argv[2]), sys.argv[3])
         return exercise_signal(native, sys.argv[3])
-    if len(sys.argv) == 5 and sys.argv[1] == "public-ocr":
-        return exercise_public_ocr(Path(sys.argv[2]), sys.argv[3], sys.argv[4])
     sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n")
     return 2