fix(python-bridge): harden sync and async route boundaries (#39332)

.github/workflows/test-rust.yml

@@ -9,6 +9,7 @@ on:
       - "rust-toolchain.toml"
       - ".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:
@@ -23,6 +24,7 @@ on:
       - "rust-toolchain.toml"
       - ".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:
@@ -121,3 +123,6 @@ jobs:
         env:
           RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
         run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
+
+      - name: Test native route wheel
+        run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl

litellm-rust/Cargo.lock

@@ -1450,6 +1450,7 @@ name = "litellm-python-bridge"
 version = "0.1.0"
 dependencies = [
  "criterion",
+ "futures-util",
  "litellm-ai-gateway",
  "litellm-core",
  "litellm-python-interop",
@@ -1458,6 +1459,7 @@ dependencies = [
  "serde",
  "serde_json",
  "tokio",
+ "tokio-tungstenite",
  "tracing",
  "tracing-subscriber",
 ]
@@ -1680,9 +1682,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3"
-version = "0.29.0"
+version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c"
+checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
 dependencies = [
  "libc",
  "once_cell",
@@ -1708,18 +1710,18 @@ dependencies = [
 
 [[package]]
 name = "pyo3-build-config"
-version = "0.29.0"
+version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078"
+checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9"
 dependencies = [
  "target-lexicon",
 ]
 
 [[package]]
 name = "pyo3-ffi"
-version = "0.29.0"
+version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b"
+checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327"
 dependencies = [
  "libc",
  "pyo3-build-config",
@@ -1727,9 +1729,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros"
-version = "0.29.0"
+version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771"
+checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c"
 dependencies = [
  "proc-macro2",
  "pyo3-macros-backend",
@@ -1739,9 +1741,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros-backend"
-version = "0.29.0"
+version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362"
+checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952"
 dependencies = [
  "heck",
  "proc-macro2",

litellm-rust/Cargo.toml

@@ -20,7 +20,7 @@ litellm-core = { path = "crates/core" }
 litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
 litellm-python-interop = { path = "crates/python-interop" }
 axum = "0.7"
-pyo3 = "0.29.0"
+pyo3 = "0.29.2"
 pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
 pythonize = "0.29.0"
 rand = "0.8"

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

@@ -16,7 +16,7 @@ extension-module = ["pyo3/extension-module"]
 panic-test = []
 
 [dependencies]
-serde.workspace = true
+futures-util.workspace = true
 tracing.workspace = true
 tracing-subscriber.workspace = true
 litellm-core = { workspace = true, features = ["bedrock-auth"] }
@@ -24,11 +24,13 @@ litellm-ai-gateway = { workspace = true, default-features = false }
 litellm-python-interop.workspace = true
 pyo3.workspace = true
 pyo3-async-runtimes.workspace = true
+serde.workspace = true
 serde_json.workspace = true
 tokio.workspace = true
 
 [dev-dependencies]
 criterion = "0.8.2"
+tokio-tungstenite.workspace = true
 
 [[bench]]
 name = "serialization"

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

@@ -0,0 +1,423 @@
+use std::future::Future;
+use std::panic::AssertUnwindSafe;
+use std::time::Duration;
+
+use futures_util::FutureExt;
+use litellm_core::error::Error;
+use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use serde::Serialize;
+use tokio::runtime::{Handle, Runtime};
+use tokio::time::{self, MissedTickBehavior};
+
+pub(crate) fn run_sync<T, F>(
+    py: Python<'_>,
+    future: F,
+    map_error: fn(Error) -> PyErr,
+) -> PyResult<Py<PyAny>>
+where
+    T: Serialize + Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    run_sync_on(
+        py,
+        pyo3_async_runtimes::tokio::get_runtime(),
+        future,
+        map_error,
+    )
+}
+
+fn run_sync_on<T, F>(
+    py: Python<'_>,
+    runtime: &Runtime,
+    future: F,
+    map_error: fn(Error) -> PyErr,
+) -> PyResult<Py<PyAny>>
+where
+    T: Serialize + Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    if Handle::try_current().is_ok() {
+        return Err(PyRuntimeError::new_err(
+            "synchronous native routes cannot run from a Tokio context; use the async route",
+        ));
+    }
+
+    let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
+    let result = map_core_result(result, map_error)?;
+    Pythonized(result).into_pyobject(py).map(Bound::unbind)
+}
+
+pub(crate) fn run_async<T, F>(
+    py: Python<'_>,
+    future: F,
+    map_error: fn(Error) -> PyErr,
+) -> PyResult<Bound<'_, PyAny>>
+where
+    T: Serialize + Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        let result = catch_future_panic(future).await?;
+        let result = map_core_result(result, map_error)?;
+        Ok(Pythonized(result))
+    })
+}
+
+fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
+    match result {
+        Ok(value) => Ok(value),
+        Err(error) => Err(
+            std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error)))
+                .map_err(panic_to_pyerr)?,
+        ),
+    }
+}
+
+async fn catch_future_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
+where
+    F: Future<Output = Result<T, Error>>,
+{
+    AssertUnwindSafe(future)
+        .catch_unwind()
+        .await
+        .map_err(panic_to_pyerr)
+}
+
+async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
+where
+    F: Future<Output = Result<T, Error>>,
+{
+    let future = catch_future_panic(future);
+    tokio::pin!(future);
+
+    let signal_interval = Duration::from_millis(50);
+    let mut signal_checks =
+        time::interval_at(time::Instant::now() + signal_interval, signal_interval);
+    signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay);
+    loop {
+        tokio::select! {
+            result = &mut future => return result,
+            _ = signal_checks.tick() => Python::attach(|py| py.check_signals())?,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::ffi::CString;
+    use std::future::poll_fn;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::{Arc, mpsc};
+    use std::task::Poll;
+    use std::thread;
+    use std::time::Instant;
+
+    use pyo3::panic::PanicException;
+    use pyo3::types::{PyDict, PyModule};
+    use serde::Serializer;
+    use tokio::runtime::Builder;
+
+    use super::*;
+
+    fn runtime_error(error: Error) -> PyErr {
+        PyRuntimeError::new_err(error.to_string())
+    }
+
+    fn panicking_error_mapper(_error: Error) -> PyErr {
+        panic!("error mapper panicked")
+    }
+
+    struct PanickingOutput;
+
+    static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0);
+
+    impl Serialize for PanickingOutput {
+        fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
+        where
+            S: Serializer,
+        {
+            panic!("serializer panicked")
+        }
+    }
+
+    #[pyfunction]
+    fn async_serialization_panic(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
+        run_async(py, async { Ok(PanickingOutput) }, runtime_error)
+    }
+
+    #[pyfunction]
+    fn async_runtime_probe(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
+        run_async(
+            py,
+            async {
+                ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst);
+                Ok(true)
+            },
+            runtime_error,
+        )
+    }
+
+    #[pyfunction]
+    fn runtime_worker_count() -> usize {
+        pyo3_async_runtimes::tokio::get_runtime()
+            .metrics()
+            .num_workers()
+    }
+
+    #[pyfunction]
+    fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
+        let completion_deadline = Instant::now() + Duration::from_secs(2);
+        while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
+            if Instant::now() >= completion_deadline {
+                return false;
+            }
+            thread::sleep(Duration::from_millis(1));
+        }
+
+        let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
+        pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
+            let _ = heartbeat_tx.send(());
+        });
+        heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
+    }
+
+    fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
+        result
+            .expect("route should complete")
+            .bind(py)
+            .extract()
+            .expect("result should convert")
+    }
+
+    #[test]
+    fn sync_runner_polls_future_on_the_caller_thread() {
+        Python::initialize();
+        Python::attach(|py| {
+            let caller_thread = std::thread::current().id();
+            let result = run_sync(
+                py,
+                async move { Ok(std::thread::current().id() == caller_thread) },
+                runtime_error,
+            );
+
+            assert!(extract_bool(py, result));
+        });
+    }
+
+    #[test]
+    fn sync_runner_releases_gil_while_waiting() {
+        Python::initialize();
+        Python::attach(|py| {
+            let result = run_sync(
+                py,
+                async {
+                    let gil_acquired = tokio::time::timeout(
+                        Duration::from_secs(2),
+                        tokio::task::spawn_blocking(|| Python::attach(|_| true)),
+                    )
+                    .await;
+                    Ok(matches!(gil_acquired, Ok(Ok(true))))
+                },
+                runtime_error,
+            );
+
+            assert!(extract_bool(py, result));
+        });
+    }
+
+    #[test]
+    fn sync_runner_rejects_calls_from_a_tokio_context() {
+        Python::initialize();
+        let runtime = Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .expect("runtime should build");
+
+        let error = runtime.block_on(async {
+            Python::attach(|py| {
+                run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
+                    .expect_err("sync route should reject a nested Tokio runtime")
+            })
+        });
+
+        assert_eq!(
+            error.to_string(),
+            "RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route"
+        );
+    }
+
+    #[test]
+    fn sync_runner_can_drive_a_current_thread_runtime() {
+        Python::initialize();
+        let runtime = Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .expect("runtime should build");
+        Python::attach(|py| {
+            let result = run_sync_on(
+                py,
+                &runtime,
+                async {
+                    tokio::task::yield_now().await;
+                    Ok(true)
+                },
+                runtime_error,
+            );
+            assert!(extract_bool(py, result));
+        });
+    }
+
+    #[test]
+    fn sync_runner_maps_a_panicked_future() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = run_sync::<bool, _>(
+                py,
+                poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
+                runtime_error,
+            )
+            .expect_err("panicked route should become a Python exception");
+
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: route future panicked");
+        });
+    }
+
+    #[test]
+    fn sync_runner_maps_a_panicked_error_mapper() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = run_sync::<bool, _>(
+                py,
+                async { Err(Error::InvalidRequest("invalid".to_string())) },
+                panicking_error_mapper,
+            )
+            .expect_err("panicked mapper should become a Python exception");
+
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: error mapper panicked");
+        });
+    }
+
+    #[test]
+    fn sync_runner_surfaces_serializer_panics() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error)
+                .expect_err("serializer panic should become a Python exception");
+
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: serializer panicked");
+        });
+    }
+
+    #[test]
+    fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
+        Python::initialize();
+        let barrier = Arc::new(tokio::sync::Barrier::new(2));
+        let callers: Vec<_> = (0..2)
+            .map(|_| {
+                let barrier = Arc::clone(&barrier);
+                thread::spawn(move || {
+                    Python::attach(|py| {
+                        extract_bool(
+                            py,
+                            run_sync(
+                                py,
+                                async move {
+                                    Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait())
+                                        .await
+                                        .is_ok())
+                                },
+                                runtime_error,
+                            ),
+                        )
+                    })
+                })
+            })
+            .collect();
+        let results: Vec<_> = callers
+            .into_iter()
+            .map(|caller| caller.join().expect("caller should not panic"))
+            .collect();
+
+        assert_eq!(results, vec![true, true]);
+    }
+
+    #[test]
+    fn async_runner_surfaces_serializer_panics() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "runtime").expect("module should be created");
+            module
+                .add_function(
+                    wrap_pyfunction!(async_serialization_panic, &module)
+                        .expect("function should wrap"),
+                )
+                .expect("function should register");
+            let locals = PyDict::new(py);
+            locals
+                .set_item("runtime", &module)
+                .expect("module should enter Python locals");
+            let code = CString::new(
+                r#"
+import asyncio
+
+async def exercise():
+    try:
+        await runtime.async_serialization_panic()
+    except BaseException as error:
+        assert type(error).__name__ == "PanicException"
+        assert str(error) == "serializer panicked"
+    else:
+        raise AssertionError("serializer panic was not raised")
+
+asyncio.run(exercise())
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("serializer panic should reach the Python awaiter");
+        });
+    }
+
+    #[test]
+    fn async_result_delivery_does_not_stall_tokio_workers() {
+        Python::initialize();
+        ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst);
+        Python::attach(|py| {
+            let module = PyModule::new(py, "runtime").expect("module should be created");
+            for function in [
+                wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"),
+                wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"),
+                wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"),
+            ] {
+                module
+                    .add_function(function)
+                    .expect("function should register");
+            }
+            let locals = PyDict::new(py);
+            locals
+                .set_item("runtime", &module)
+                .expect("module should enter Python locals");
+            let code = CString::new(
+                r#"
+import asyncio
+
+async def exercise():
+    worker_count = runtime.runtime_worker_count()
+    awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)]
+    assert runtime.runtime_is_responsive(worker_count)
+    assert await asyncio.gather(*awaitables) == [True] * worker_count
+
+asyncio.run(exercise())
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("result delivery should leave Tokio workers responsive");
+        });
+    }
+}

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

@@ -1,6 +1,7 @@
 mod constants;
 mod diagnostics;
 mod errors;
+mod execution;
 pub mod function_trace;
 mod marshal;
 mod routes;
@@ -8,6 +9,7 @@ mod routes;
 use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
 use pyo3::prelude::*;
 use pyo3::types::PyAny;
+use serde_json::Value;
 
 use crate::errors::core_error_to_pyerr;
 use crate::marshal::{marshal_headers, optional_timeout};
@@ -25,16 +27,16 @@ impl ResponsesWebSocketConnection {
         _cls: &Bound<'py, pyo3::types::PyType>,
         py: Python<'py>,
         url: String,
-        headers: Option<Py<PyAny>>,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)] headers: Option<Value>,
         timeout_seconds: Option<f64>,
     ) -> PyResult<Bound<'py, PyAny>> {
-        let headers = marshal_headers(py, headers)?;
+        let headers = marshal_headers(headers)?;
         let timeout = optional_timeout(timeout_seconds);
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
             let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
                 .await
                 .map_err(core_error_to_pyerr)?;
-            Python::attach(|py| Py::new(py, ResponsesWebSocketConnection { inner }))
+            Ok(ResponsesWebSocketConnection { inner })
         })
     }
 
@@ -60,24 +62,36 @@ impl ResponsesWebSocketConnection {
     }
 }
 
-#[pymodule]
-fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    errors::register(module)?;
-    routes::register(module)?;
-    module.add_class::<ResponsesWebSocketConnection>()?;
-    diagnostics::register(module)
+#[pymodule(gil_used = false)]
+mod _native {
+    use pyo3::prelude::*;
+
+    #[pymodule_init]
+    fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
+        super::errors::register(module)?;
+        super::routes::register(module)?;
+        module.add_class::<super::ResponsesWebSocketConnection>()?;
+        super::diagnostics::register(module)
+    }
 }
 
 #[cfg(test)]
 mod tests {
+    use std::ffi::CString;
+    use std::time::Duration;
+
+    use futures_util::{SinkExt, StreamExt};
+    use pyo3::types::PyDict;
+    use tokio::net::TcpListener;
+    use tokio_tungstenite::{accept_async, tungstenite::Message};
+
     use super::*;
 
     #[test]
     fn module_registration_preserves_the_public_surface() {
         Python::initialize();
         Python::attach(|py| {
-            let module = PyModule::new(py, "_native").expect("module should be created");
-            _native(&module).expect("module should register");
+            let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py);
 
             let expected = [
                 "RustBridgeDeclined",
@@ -95,13 +109,79 @@ mod tests {
                 "gil_stats",
             ];
 
-            for name in expected {
-                assert!(
-                    module
-                        .hasattr(name)
-                        .expect("attribute lookup should succeed")
-                );
-            }
+            let public_names: Vec<String> = module
+                .dict()
+                .keys()
+                .extract::<Vec<String>>()
+                .expect("module names should be strings")
+                .into_iter()
+                .filter(|name| !name.starts_with("__"))
+                .collect();
+            assert_eq!(public_names, expected);
         });
     }
+
+    #[test]
+    fn responses_websocket_connection_round_trips_through_python() {
+        Python::initialize();
+        let runtime = pyo3_async_runtimes::tokio::get_runtime();
+        let listener = runtime
+            .block_on(TcpListener::bind("127.0.0.1:0"))
+            .expect("listener should bind");
+        let address = listener
+            .local_addr()
+            .expect("listener should have an address");
+        let server = runtime.spawn(async move {
+            let (stream, _) = listener.accept().await.expect("server should accept");
+            let mut socket = accept_async(stream)
+                .await
+                .expect("handshake should succeed");
+
+            let message = socket
+                .next()
+                .await
+                .expect("client should send a frame")
+                .expect("client frame should be valid");
+            assert_eq!(message, Message::Text("from-python".into()));
+            socket
+                .send(Message::Text("from-server".into()))
+                .await
+                .expect("server should reply");
+            assert!(matches!(socket.next().await, Some(Ok(Message::Close(_)))));
+        });
+
+        Python::attach(|py| {
+            let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py);
+            let locals = PyDict::new(py);
+            locals
+                .set_item("native", &module)
+                .expect("module should enter Python locals");
+            locals
+                .set_item("url", format!("ws://{address}"))
+                .expect("URL should enter Python locals");
+            let code = CString::new(
+                r#"
+import asyncio
+
+async def exercise():
+    connection = await native.ResponsesWebSocketConnection.connect(url)
+    assert type(connection) is native.ResponsesWebSocketConnection
+    await connection.send_text("from-python")
+    assert await connection.recv_text() == "from-server"
+    await connection.close()
+    assert await connection.recv_text() is None
+
+asyncio.run(asyncio.wait_for(exercise(), timeout=5))
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("Python WebSocket methods should round trip");
+        });
+
+        runtime
+            .block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await })
+            .expect("server should finish")
+            .expect("server task should not panic");
+    }
 }

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

@@ -40,12 +40,9 @@ pub(crate) fn optional_timeout(timeout_seconds: Option<f64>) -> Option<Duration>
     })
 }
 
-pub(crate) fn marshal_headers(
-    py: Python<'_>,
-    headers: Option<Py<PyAny>>,
-) -> PyResult<HashMap<String, String>> {
+pub(crate) fn marshal_headers(headers: Option<Value>) -> PyResult<HashMap<String, String>> {
     let value = match headers {
-        Some(headers) => from_py(headers.bind(py))?,
+        Some(headers) => headers,
         None => Value::Object(Map::new()),
     };
     let Value::Object(headers) = value else {

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

@@ -9,10 +9,10 @@ use pyo3::prelude::*;
 use serde_json::{Map, Value};
 
 use crate::errors::core_error_to_pyerr;
+use crate::execution::{run_async, run_sync};
+use crate::function_trace::trace_call;
 use crate::marshal::{optional_object, optional_object_to_map, optional_timeout};
 
-use super::{block_on, into_py_future};
-
 struct TranscriptionInputs {
     model: String,
     audio: Value,
@@ -88,7 +88,7 @@ fn transcription(
         optional_params,
         timeout_seconds,
     )?;
-    block_on(py, call(inputs), trace, core_error_to_pyerr)
+    run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr)
 }
 
 #[pyfunction]
@@ -117,7 +117,7 @@ fn atranscription(
         optional_params,
         timeout_seconds,
     )?;
-    into_py_future(py, call(inputs), trace, core_error_to_pyerr)
+    run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr)
 }
 
 pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {

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

@@ -11,10 +11,10 @@ use pyo3::prelude::*;
 use serde_json::{Map, Value};
 
 use crate::errors::chat_completions_error_to_pyerr;
+use crate::execution::{run_async, run_sync};
+use crate::function_trace::trace_call;
 use crate::marshal::{optional_object, optional_object_to_map, optional_timeout};
 
-use super::{block_on, into_py_future};
-
 struct ChatCompletionsInputs {
     model: String,
     messages: Value,
@@ -117,7 +117,11 @@ fn chat_completions(
         extra_headers,
         timeout_seconds,
     )?;
-    block_on(py, call(inputs), trace, chat_completions_error_to_pyerr)
+    run_sync(
+        py,
+        trace_call(call(inputs), trace),
+        chat_completions_error_to_pyerr,
+    )
 }
 
 #[pyfunction]
@@ -146,7 +150,11 @@ fn achat_completions(
         extra_headers,
         timeout_seconds,
     )?;
-    into_py_future(py, call(inputs), trace, chat_completions_error_to_pyerr)
+    run_async(
+        py,
+        trace_call(call(inputs), trace),
+        chat_completions_error_to_pyerr,
+    )
 }
 
 pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {

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

@@ -9,10 +9,10 @@ use pyo3::prelude::*;
 use serde_json::{Map, Value};
 
 use crate::errors::core_error_to_pyerr;
+use crate::execution::{run_async, run_sync};
+use crate::function_trace::trace_call;
 use crate::marshal::{optional_object, optional_timeout};
 
-use super::{block_on, into_py_future};
-
 struct MessagesInputs {
     model: String,
     body: Value,
@@ -86,7 +86,7 @@ fn messages(
         extra_headers,
         timeout_seconds,
     )?;
-    block_on(py, call(inputs), trace, core_error_to_pyerr)
+    run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr)
 }
 
 #[pyfunction]
@@ -113,7 +113,7 @@ fn amessages(
         extra_headers,
         timeout_seconds,
     )?;
-    into_py_future(py, call(inputs), trace, core_error_to_pyerr)
+    run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr)
 }
 
 pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {

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

@@ -1,11 +1,4 @@
-use std::future::Future;
-
-use litellm_core::error::Error;
-use litellm_python_interop::{release_gil, to_py};
 use pyo3::prelude::*;
-use serde::Serialize;
-
-use crate::function_trace::trace_call;
 
 mod audio_transcription;
 mod chat_completions;
@@ -18,36 +11,3 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
     messages::register(module)?;
     chat_completions::register(module)
 }
-
-fn block_on<T>(
-    py: Python<'_>,
-    call: impl Future<Output = Result<T, Error>> + Send,
-    trace: bool,
-    map_err: fn(Error) -> PyErr,
-) -> PyResult<Py<PyAny>>
-where
-    T: Serialize + Send,
-{
-    let result = release_gil(py, || {
-        pyo3_async_runtimes::tokio::get_runtime().block_on(trace_call(call, trace))
-    });
-    match result {
-        Ok(response) => to_py(py, &response),
-        Err(err) => Err(map_err(err)),
-    }
-}
-
-fn into_py_future<'py, T>(
-    py: Python<'py>,
-    call: impl Future<Output = Result<T, Error>> + Send + 'static,
-    trace: bool,
-    map_err: fn(Error) -> PyErr,
-) -> PyResult<Bound<'py, PyAny>>
-where
-    T: Serialize + Send + 'static,
-{
-    pyo3_async_runtimes::tokio::future_into_py(py, async move {
-        let response = trace_call(call, trace).await.map_err(map_err)?;
-        Python::attach(|py| to_py(py, &response))
-    })
-}

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

@@ -7,10 +7,10 @@ use pyo3::prelude::*;
 use serde_json::{Map, Value};
 
 use crate::errors::core_error_to_pyerr;
+use crate::execution::{run_async, run_sync};
+use crate::function_trace::trace_call;
 use crate::marshal::{optional_object, optional_object_to_map, optional_timeout};
 
-use super::{block_on, into_py_future};
-
 struct OcrInputs {
     model: String,
     document: Value,
@@ -90,7 +90,7 @@ fn ocr(
         optional_params,
         timeout_seconds,
     )?;
-    block_on(py, call(inputs), trace, core_error_to_pyerr)
+    run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr)
 }
 
 #[pyfunction]
@@ -119,7 +119,7 @@ fn aocr(
         optional_params,
         timeout_seconds,
     )?;
-    into_py_future(py, call(inputs), trace, core_error_to_pyerr)
+    run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr)
 }
 
 pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {

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

@@ -2,4 +2,4 @@ mod gil;
 mod marshal;
 
 pub use gil::{release_count, release_gil};
-pub use marshal::{from_py, to_py};
+pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py};

litellm-rust/crates/python-interop/src/marshal.rs

@@ -1,4 +1,8 @@
+use std::any::Any;
+use std::panic::{AssertUnwindSafe, catch_unwind};
+
 use pyo3::exceptions::PyValueError;
+use pyo3::panic::PanicException;
 use pyo3::prelude::*;
 use serde::Serialize;
 use serde::de::DeserializeOwned;
@@ -18,3 +22,71 @@ where
         .map(Bound::unbind)
         .map_err(|error| PyValueError::new_err(error.to_string()))
 }
+
+pub struct Pythonized<T>(pub T);
+
+impl<'py, T> IntoPyObject<'py> for Pythonized<T>
+where
+    T: Serialize,
+{
+    type Target = PyAny;
+    type Output = Bound<'py, PyAny>;
+    type Error = PyErr;
+
+    fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
+        catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, &self.0)))
+            .map_err(panic_to_pyerr)?
+            .map_err(|error| PyValueError::new_err(error.to_string()))
+    }
+}
+
+pub fn panic_to_pyerr(payload: Box<dyn Any + Send>) -> PyErr {
+    let message = payload
+        .downcast_ref::<String>()
+        .map(String::as_str)
+        .or_else(|| payload.downcast_ref::<&str>().copied())
+        .unwrap_or("panic from Rust code");
+    PanicException::new_err(message.to_string())
+}
+
+#[cfg(test)]
+mod tests {
+    use serde::Serializer;
+
+    use super::*;
+
+    struct PanickingSerializer;
+
+    impl Serialize for PanickingSerializer {
+        fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
+        where
+            S: Serializer,
+        {
+            panic!("serializer panicked")
+        }
+    }
+
+    #[test]
+    fn pythonized_converts_on_the_attached_thread() {
+        Python::initialize();
+        Python::attach(|py| {
+            let value: Vec<i32> = Pythonized(vec![1, 2, 3])
+                .into_pyobject(py)
+                .and_then(|value| value.extract())
+                .expect("value should convert");
+            assert_eq!(value, vec![1, 2, 3]);
+        });
+    }
+
+    #[test]
+    fn pythonized_maps_serializer_panics_to_a_base_exception() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = Pythonized(PanickingSerializer)
+                .into_pyobject(py)
+                .expect_err("serializer panic should become a Python exception");
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: serializer panicked");
+        });
+    }
+}

tests/test_litellm/rust_bridge/native_route_wheel_test.py

@@ -0,0 +1,349 @@
+from __future__ import annotations
+
+import asyncio
+import importlib.util
+import json
+import os
+import signal
+import subprocess
+import sys
+import tempfile
+import threading
+import zipfile
+from http.client import HTTPMessage
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from socket import socket as Socket
+from typing import Final
+
+REQUEST_STARTED: Final = threading.Event()
+REQUEST_CANCELLED: Final = threading.Event()
+
+ANTHROPIC_RESPONSE: Final = (
+    b'{"id":"msg_native","type":"message","role":"assistant",'
+    b'"model":"claude-sonnet-4-5","content":[{"type":"text","text":"native-message"}],'
+    b'"stop_reason":"end_turn","stop_sequence":null,'
+    b'"usage":{"input_tokens":2,"output_tokens":3}}'
+)
+
+
+class NativeRouteHandler(BaseHTTPRequestHandler):
+    protocol_version = "HTTP/1.1"
+
+    def do_POST(self) -> None:
+        content_length: Final = int(self.headers.get("content-length", "0"))
+        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")
+        assert_native_request(route, outcome, self.path, self.headers, body)
+        if outcome == "hang":
+            REQUEST_STARTED.set()
+            self.connection.settimeout(5)
+            if connection_was_cancelled(self.connection):
+                REQUEST_CANCELLED.set()
+            return
+
+        status: Final = 429 if outcome == "429" else 200
+        response_body: Final = native_response(status, route)
+
+        self.send_response(status)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", str(len(response_body)))
+        self.send_header("connection", "close")
+        self.end_headers()
+        self.wfile.write(response_body)
+
+    def log_message(self, _message_format: str, *_args: object) -> None:
+        pass
+
+
+def connection_was_cancelled(connection: Socket) -> bool:
+    try:
+        return connection.recv(1) == b""
+    except TimeoutError:
+        return False
+    except OSError:
+        return True
+
+
+def assert_native_request(
+    route: str | None,
+    outcome: str | None,
+    path: str,
+    headers: HTTPMessage,
+    body: object,
+) -> None:
+    if route not in {"ocr", "transcription", "messages", "chat_completions"}:
+        raise AssertionError(f"unexpected route marker: {route!r}")
+    if outcome not in {"success", "429", "hang"}:
+        raise AssertionError(f"unexpected outcome marker: {outcome!r}")
+    if not isinstance(body, dict):
+        raise TypeError(f"{route} sent {type(body).__name__}, expected a JSON object")
+    if route == "ocr":
+        assert path == "/v1/ocr"
+        assert headers.get("authorization") == "Bearer sk-native"
+        assert body["model"] == "mistral-ocr-latest"
+        assert body["document"]["document_url"] == "https://example.com/document.pdf"
+        assert body["include_image_base64"] is True
+        return
+    if route == "transcription":
+        assert path == "/model/mistral.voxtral-mini-3b-2507/converse"
+        assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ")
+        assert headers.get("x-amz-date")
+        assert body["messages"][0]["content"][0]["audio"]["source"]["bytes"] == "AQI="
+        assert "The audio language is en" in body["messages"][0]["content"][1]["text"]
+        return
+    assert path == "/v1/messages"
+    assert headers.get("x-api-key") == "sk-native"
+    assert body["model"] == "claude-sonnet-4-5"
+    if route == "messages":
+        assert body["max_tokens"] == 16
+        assert body["messages"][0]["content"] == "hello-from-messages"
+        return
+    assert body["max_tokens"] == 17
+    assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}]
+
+
+def native_response(status: int, route: str | None) -> bytes:
+    if status == 429:
+        return b'{"error":"native-rate-limit"}'
+    if route == "ocr":
+        return b'{"pages":[{"index":0,"markdown":"native-ocr"}]}'
+    if route == "transcription":
+        return b'{"output":{"message":{"content":[{"text":"native-transcription"}]}}}'
+    return ANTHROPIC_RESPONSE
+
+
+def load_native(native_path: Path) -> object:
+    module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path)
+    if module_spec is None or module_spec.loader is None:
+        raise RuntimeError("cannot create native extension import specification")
+    native_module: Final = importlib.util.module_from_spec(module_spec)
+    module_spec.loader.exec_module(native_module)
+    return native_module
+
+
+def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
+    common: Final = {
+        "api_base": api_base,
+        "extra_headers": {"x-test-outcome": outcome, "x-test-route": route},
+        "timeout_seconds": 3.0,
+    }
+    if route == "ocr":
+        return common | {
+            "model": "mistral-ocr-latest",
+            "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"},
+            "api_key": "sk-native",
+            "custom_llm_provider": "mistral",
+            "optional_params": {"include_image_base64": True},
+        }
+    if route == "transcription":
+        return common | {
+            "model": "mistral.voxtral-mini-3b-2507",
+            "audio": {"data": "AQI=", "format": "wav", "filename": "audio.wav"},
+            "custom_llm_provider": "bedrock",
+            "optional_params": {
+                "aws_access_key_id": "native-access-key",
+                "aws_secret_access_key": "native-secret-key",
+                "aws_region_name": "us-east-1",
+                "language": "en",
+            },
+        }
+    if route == "messages":
+        return common | {
+            "model": "claude-sonnet-4-5",
+            "body": {
+                "model": "claude-sonnet-4-5",
+                "max_tokens": 16,
+                "messages": [{"role": "user", "content": "hello-from-messages"}],
+            },
+            "api_key": "sk-native",
+            "custom_llm_provider": "anthropic",
+        }
+    if route == "chat_completions":
+        return common | {
+            "model": "anthropic/claude-sonnet-4-5",
+            "messages": [{"role": "user", "content": "hello-from-chat"}],
+            "optional_params": {"max_tokens": 17},
+            "api_key": "sk-native",
+        }
+    raise AssertionError(f"unknown route: {route}")
+
+
+def assert_success(route: str, response: object) -> None:
+    if not isinstance(response, dict):
+        raise TypeError(f"{route} returned {type(response).__name__}, expected dict")
+    actual: Final = success_value(route, response)
+    expected: Final = (
+        "native-ocr" if route == "ocr" else "native-transcription" if route == "transcription" else "native-message"
+    )
+    if actual != expected:
+        raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}")
+
+
+def assert_traced_success(route: str, response: object) -> None:
+    if not isinstance(response, dict):
+        raise TypeError(f"{route} returned {type(response).__name__}, expected a traced dict")
+    assert_success(route, response["response"])
+    expected_function: Final = "audio_transcription" if route == "transcription" else route
+    assert response["trace"][0] == {"function": expected_function, "depth": 0}
+
+
+def success_value(route: str, response: dict[object, object]) -> object:
+    if route == "ocr":
+        return response["pages"][0]["markdown"]
+    if route == "transcription":
+        return response["text"]
+    if route == "messages":
+        return response["content"][0]["text"]
+    return response["choices"][0]["message"]["content"]
+
+
+def assert_rate_limit(native: object, route: str, error: BaseException) -> None:
+    if route == "chat_completions":
+        upstream_error: Final = native.RustUpstreamError
+        if not isinstance(error, upstream_error) or error.args[0] != 429:
+            raise AssertionError(f"{route} returned the wrong 429 error: {error!r}")
+        return
+    if not isinstance(error, RuntimeError) or "429" not in str(error):
+        raise AssertionError(f"{route} returned the wrong 429 error: {error!r}")
+
+
+def exercise_sync(native: object, api_base: str) -> None:
+    for route in ("ocr", "transcription", "messages", "chat_completions"):
+        function: Final = getattr(native, route)
+        assert_success(route, function(**route_kwargs(route, api_base, "success")))
+        assert_traced_success(route, function(**route_kwargs(route, api_base, "success"), trace=True))
+        try:
+            function(**route_kwargs(route, api_base, "429"))
+        except (RuntimeError, native.RustUpstreamError) as error:
+            assert_rate_limit(native, route, error)
+        else:
+            raise AssertionError(f"{route} accepted a 429 response")
+
+
+async def exercise_async(native: object, api_base: str) -> None:
+    for route in ("ocr", "transcription", "messages", "chat_completions"):
+        function: Final = getattr(native, f"a{route}")
+        assert_success(route, await function(**route_kwargs(route, api_base, "success")))
+        assert_traced_success(route, await function(**route_kwargs(route, api_base, "success"), trace=True))
+        try:
+            await function(**route_kwargs(route, api_base, "429"))
+        except (RuntimeError, native.RustUpstreamError) as error:
+            assert_rate_limit(native, route, error)
+        else:
+            raise AssertionError(f"a{route} accepted a 429 response")
+
+
+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)
+            )
+        ),
+        timeout=15,
+    )
+    for response in responses:
+        assert_success("messages", response)
+
+
+def exercise_routes(native_path: Path, api_base: str) -> object:
+    native: Final = load_native(native_path)
+    exercise_sync(native, api_base)
+    asyncio.run(exercise_async(native, api_base))
+    asyncio.run(exercise_async_concurrency(native, api_base))
+    return native
+
+
+def exercise_signal(native: object, api_base: str) -> int:
+    try:
+        native.messages(
+            **route_kwargs("messages", api_base, "hang"),
+        )
+    except KeyboardInterrupt:
+        sys.stdout.write("KeyboardInterrupt\n")
+        sys.stdout.flush()
+        sys.stdin.read(1)
+        return 0
+    raise AssertionError("sync native route ignored SIGINT")
+
+
+def verify_sigint(native_path: Path, api_base: str) -> None:
+    REQUEST_STARTED.clear()
+    REQUEST_CANCELLED.clear()
+    process: Final = subprocess.Popen(
+        (sys.executable, __file__, "child", str(native_path), api_base),
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    try:
+        if not REQUEST_STARTED.wait(30):
+            process.kill()
+            stdout, stderr = process.communicate(timeout=5)
+            raise AssertionError(
+                f"native route matrix did not reach the hanging upstream\nstdout:\n{stdout}\nstderr:\n{stderr}"
+            )
+        os.kill(process.pid, signal.SIGINT)
+        if not REQUEST_CANCELLED.wait(5):
+            raise AssertionError("interrupted native route did not cancel its upstream future")
+        if process.poll() is not None:
+            raise AssertionError("signal child exited before cancellation was observed")
+        stdout, stderr = process.communicate(input="\n", timeout=5)
+        if process.returncode != 0 or stdout != "KeyboardInterrupt\n":
+            raise AssertionError(
+                f"signal child exited with status {process.returncode}\nstdout:\n{stdout}\nstderr:\n{stderr}"
+            )
+    finally:
+        if process.poll() is None:
+            process.kill()
+            process.wait(timeout=5)
+
+
+def verify_wheel(wheel: Path) -> int:
+    with tempfile.TemporaryDirectory() as temporary_directory, zipfile.ZipFile(wheel) as archive:
+        wheel_root: Final = Path(temporary_directory)
+        for member in archive.infolist():
+            target: Final = wheel_root / member.filename
+            if member.is_dir():
+                target.mkdir(parents=True, exist_ok=True)
+            else:
+                target.parent.mkdir(parents=True, exist_ok=True)
+                target.write_bytes(archive.read(member))
+        native_members: Final = tuple(
+            member
+            for member in archive.infolist()
+            if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so")
+        )
+        if len(native_members) != 1:
+            raise AssertionError(f"expected one native extension, found {len(native_members)}")
+        native_path: Final = wheel_root / native_members[0].filename
+
+        server: Final = ThreadingHTTPServer(("127.0.0.1", 0), NativeRouteHandler)
+        server_thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
+        server_thread.start()
+        api_base: Final = f"http://127.0.0.1:{server.server_address[1]}"
+        try:
+            verify_sigint(native_path, api_base)
+        finally:
+            server.shutdown()
+            server.server_close()
+            server_thread.join(timeout=5)
+    return 0
+
+
+def main() -> int:
+    if len(sys.argv) == 2:
+        return verify_wheel(Path(sys.argv[1]))
+    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])
+    sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n")
+    return 2
+
+
+if __name__ == "__main__":
+    sys.exit(main())