diff --git a/litellm-rust/crates/python-bridge/src/driver.rs b/litellm-rust/crates/python-bridge/src/driver.rs index 331eecff6a5..e1e35845850 100644 --- a/litellm-rust/crates/python-bridge/src/driver.rs +++ b/litellm-rust/crates/python-bridge/src/driver.rs @@ -2,6 +2,8 @@ use std::ffi::CString; use pyo3::prelude::*; +use crate::errors::RustBridgeDriverError; + const DRIVE: &str = r#" def drive_sync(arguments): host = Host(arguments, False) @@ -37,12 +39,46 @@ pub(crate) fn compile<'py>( route: &str, host: &str, ) -> PyResult> { - let source = CString::new(format!("{host}\n{DRIVE}")).map_err(|_| { - pyo3::exceptions::PyValueError::new_err("driver source contains a null byte") - })?; + let source = CString::new(format!("{host}\n{DRIVE}")) + .map_err(|_| RustBridgeDriverError::new_err("driver source contains a null byte"))?; let filename = CString::new(format!("{route}_driver.py")) - .map_err(|_| pyo3::exceptions::PyValueError::new_err("invalid driver route name"))?; + .map_err(|_| RustBridgeDriverError::new_err("driver route name contains a null byte"))?; let module_name = CString::new(format!("_{route}_driver")) - .map_err(|_| pyo3::exceptions::PyValueError::new_err("invalid driver route name"))?; + .map_err(|_| RustBridgeDriverError::new_err("driver route name contains a null byte"))?; PyModule::from_code(py, &source, &filename, &module_name) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn null_byte_in_route_raises_driver_error() { + Python::initialize(); + Python::attach(|py| { + let error = compile(py, "invalid\0route", "class Host: pass") + .expect_err("route names containing null bytes should fail"); + + assert!(error.is_instance_of::(py)); + assert_eq!( + error.to_string(), + "RustBridgeDriverError: driver route name contains a null byte" + ); + }); + } + + #[test] + fn null_byte_in_source_raises_driver_error() { + Python::initialize(); + Python::attach(|py| { + let error = compile(py, "test", "class Host:\0 pass") + .expect_err("driver source containing null bytes should fail"); + + assert!(error.is_instance_of::(py)); + assert_eq!( + error.to_string(), + "RustBridgeDriverError: driver source contains a null byte" + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 00cfa55c4ce..b3a049c0c31 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -16,6 +16,13 @@ pyo3::create_exception!( "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." ); +pyo3::create_exception!( + _native, + RustBridgeDriverError, + pyo3::exceptions::PyRuntimeError, + "The Rust bridge could not initialize a route driver." +); + pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { match err { Error::InvalidProvider(_) => PyValueError::new_err("Invalid provider configuration"), @@ -81,7 +88,11 @@ pub(crate) fn messages_provider_error_to_pyerr(err: Error) -> PyErr { pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { let py = module.py(); module.add("RustBridgeDeclined", py.get_type::())?; - module.add("RustUpstreamError", py.get_type::()) + module.add("RustUpstreamError", py.get_type::())?; + module.add( + "RustBridgeDriverError", + py.get_type::(), + ) } #[cfg(test)] diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 2a14beb41bd..b52dfde7e00 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -35,6 +35,7 @@ mod tests { let expected = [ "RustBridgeDeclined", "RustUpstreamError", + "RustBridgeDriverError", "ocr", "aocr", "transcription",