diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index b8b6291283d..47cae372c00 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -13,6 +13,8 @@ litellm-rust has four crates. A crate is a layer or shared foundation, not a rou Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. +Interpreter and runtime boundaries are enforced by `crates/python-bridge/src/architecture.rs`: GIL attach/detach and `block_on` live only in `python-bridge/src/execution.rs`, Tokio runtime construction only there and at the `#[pymodule]` init site in `src/lib.rs`, `SendWrapper` is banned, and `python-interop` stays domain-neutral. The test fails until its allowlist is updated — moving a boundary is a deliberate act that also updates the crate AGENTS.md. + ## Where a route lives A top-level LiteLLM call is a module under `crates/core/src//`, shaped like `messages`: diff --git a/litellm-rust/crates/python-bridge/src/architecture.rs b/litellm-rust/crates/python-bridge/src/architecture.rs new file mode 100644 index 00000000000..149adb76de1 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/architecture.rs @@ -0,0 +1,259 @@ +//! Enforcement: interpreter and runtime boundaries in the PyO3 crates. +//! +//! Companion to `crates/core/tests/workspace_crate_allowlist.rs`: layering +//! rules that live in the crate `AGENTS.md` files become tests that fail +//! until deliberately updated here. +//! +//! Scanned text is *production* source only: each file's content up to its +//! trailing `#[cfg(test)]` module (the codebase keeps test modules at EOF). +//! This scanner skips its own file, so the tokens below can be written out +//! literally. +//! +//! Rules, each tied to a documented failure mode: +//! +//! 1. GIL acquisition and release (`Python::attach`, `Python::with_gil`, +//! `.detach(`, `.allow_threads(`) appears in `python-bridge` only in +//! `execution.rs`. The interpreter boundary belongs to +//! `litellm-python-interop` primitives and the single execution module; +//! scattered attach/detach calls are how GIL-ordering deadlocks and +//! 5 ms-per-handoff contention creep in. +//! 2. `block_on` appears in `python-bridge` only in `execution.rs`. Blocking +//! a thread on a future anywhere else (a route body, a pyclass method) +//! stalls the calling Python thread and risks nested-runtime panics. +//! 3. Tokio runtime construction (`Runtime::new`, `tokio::runtime::Builder`, +//! `Builder::new_*`) appears in `python-bridge` only in `execution.rs` +//! and `lib.rs` (the `#[pymodule]` host-init site). One shared runtime +//! per process; per-call construction costs threads and an event loop, +//! and a second runtime silently fragments the worker budget. +//! 4. `SendWrapper` appears nowhere in `python-bridge` or +//! `litellm-python-interop`. It makes `!Send` Python-bound values `Send` +//! by panicking when touched from another Tokio worker — a latent +//! runtime bomb, not a fix. +//! 5. `litellm-python-interop` stays domain-neutral: no `litellm_core`, +//! `litellm_ai_gateway`, or `litellm_python_bridge` tokens in its +//! sources, keeping the dependency direction acyclic. +//! 6. `litellm-python-interop` never blocks on futures and never constructs +//! a Tokio runtime: it provides primitives, hosts drive runtimes. + +use std::fs; +use std::path::{Path, PathBuf}; + +const BRIDGE_SRC: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/src"); +const INTEROP_SRC: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../python-interop/src"); + +const TEST_MODULE_MARKER: &str = "\n#[cfg(test)]"; + +const GIL_TOKENS: &[&str] = &[ + "Python::attach", + "Python::with_gil", + ".detach(", + ".allow_threads(", +]; + +const RUNTIME_CONSTRUCTION_TOKENS: &[&str] = &[ + "Runtime::new", + "tokio::runtime::Builder", + "Builder::new_multi_thread", + "Builder::new_current_thread", +]; + +const INTEROP_DOMAIN_TOKENS: &[&str] = &[ + "litellm_core", + "litellm_ai_gateway", + "litellm_python_bridge", +]; + +/// Collect `*.rs` files under `root`, depth-first. +fn rust_sources(root: &str) -> Vec { + fn walk(dir: &Path, out: &mut Vec) { + let entries = fs::read_dir(dir).unwrap_or_else(|error| { + panic!("{} should be readable: {error}", dir.display()); + }); + for entry in entries.filter_map(Result::ok) { + let path = entry.path(); + if path.is_dir() { + walk(&path, out); + } else if path.extension().is_some_and(|ext| ext == "rs") { + out.push(path); + } + } + } + + let mut sources = Vec::new(); + walk(Path::new(root), &mut sources); + sources +} + +/// The production half of a file: everything before the trailing +/// `#[cfg(test)]` test module. +fn production_text(content: &str) -> &str { + match content.find(TEST_MODULE_MARKER) { + Some(offset) => &content[..offset], + None => content, + } +} + +/// Describe a violation of `token` in `relative_path`, or `None` when the +/// path is allowlisted or the token only appears in test code. +fn violation( + content: &str, + relative_path: &Path, + token: &str, + allowed_suffixes: &[&str], +) -> Option { + if !production_text(content).contains(token) { + return None; + } + if allowed_suffixes + .iter() + .any(|suffix| relative_path.ends_with(suffix)) + { + return None; + } + Some(format!( + "`{token}` in {} (allowed only in {allowed_suffixes:?})", + relative_path.display() + )) +} + +/// Assert that `token` never appears in the production half of any scanned +/// source under `root`, except in files whose path ends with an allowed +/// suffix. The scanner's own file is skipped so its literals cannot match. +fn assert_confined(root: &str, token: &str, allowed_suffixes: &[&str], rule: &str) { + let mut violations = Vec::new(); + for path in rust_sources(root) { + if path + .file_name() + .is_some_and(|name| name == "architecture.rs") + { + continue; + } + let content = fs::read_to_string(&path) + .unwrap_or_else(|error| panic!("{} should be readable: {error}", path.display())); + let relative = path.strip_prefix(root).unwrap_or(&path); + if let Some(violation) = violation(&content, relative, token, allowed_suffixes) { + violations.push(format!("{}: {violation}", path.display())); + } + } + assert!( + violations.is_empty(), + "{rule}\nviolations:\n {}", + violations.join("\n ") + ); +} + +#[test] +fn python_bridge_gil_calls_are_confined_to_execution() { + for token in GIL_TOKENS { + assert_confined( + BRIDGE_SRC, + token, + &["execution.rs"], + "interpreter attach/detach belongs to litellm-python-interop and python-bridge/src/execution.rs (rule 1 in the module docs)", + ); + } +} + +#[test] +fn python_bridge_block_on_is_confined_to_execution() { + assert_confined( + BRIDGE_SRC, + "block_on", + &["execution.rs"], + "blocking on futures belongs to python-bridge/src/execution.rs (rule 2 in the module docs)", + ); +} + +#[test] +fn python_bridge_runtime_construction_is_confined_to_execution_and_module_init() { + for token in RUNTIME_CONSTRUCTION_TOKENS { + assert_confined( + BRIDGE_SRC, + token, + &["execution.rs", "lib.rs"], + "one shared Tokio runtime, constructed only at the host-init site (rule 3 in the module docs)", + ); + } +} + +#[test] +fn send_wrapper_is_banned_in_the_pyo3_crates() { + assert_confined( + BRIDGE_SRC, + "SendWrapper", + &[], + "SendWrapper panics across Tokio workers; convert to owned types instead (rule 4 in the module docs)", + ); + assert_confined( + INTEROP_SRC, + "SendWrapper", + &[], + "SendWrapper panics across Tokio workers; convert to owned types instead (rule 4 in the module docs)", + ); +} + +#[test] +fn python_interop_stays_domain_neutral() { + for token in INTEROP_DOMAIN_TOKENS { + assert_confined( + INTEROP_SRC, + token, + &[], + "litellm-python-interop must not depend on LiteLLM domain crates (rule 5 in the module docs)", + ); + } +} + +#[test] +fn python_interop_never_blocks_or_builds_runtimes() { + assert_confined( + INTEROP_SRC, + "block_on", + &[], + "litellm-python-interop provides primitives; hosts drive runtimes (rule 6 in the module docs)", + ); + for token in RUNTIME_CONSTRUCTION_TOKENS { + assert_confined( + INTEROP_SRC, + token, + &[], + "litellm-python-interop provides primitives; hosts drive runtimes (rule 6 in the module docs)", + ); + } +} + +#[test] +fn scanner_flags_a_token_outside_the_allowlist() { + let content = "fn route() {\n Python::attach(|py| drop(py));\n}\n"; + let violation = violation( + content, + Path::new("routes/chat.rs"), + GIL_TOKENS[0], + &["execution.rs"], + ); + assert!(violation.is_some(), "token outside the allowlist must flag"); +} + +#[test] +fn scanner_accepts_a_token_inside_the_allowlist() { + let content = "fn runner() {\n Python::attach(|py| drop(py));\n}\n"; + let violation = violation( + content, + Path::new("src/execution.rs"), + GIL_TOKENS[0], + &["execution.rs"], + ); + assert!(violation.is_none(), "allowlisted file must not flag"); +} + +#[test] +fn scanner_ignores_test_modules() { + let content = "fn route() {}\n\n#[cfg(test)]\nmod tests {\n fn probe() {\n Python::attach(|py| drop(py));\n }\n}\n"; + let violation = violation( + content, + Path::new("routes/chat.rs"), + GIL_TOKENS[0], + &["execution.rs"], + ); + assert!(violation.is_none(), "test modules must not flag"); +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 5f36a22370a..dcd63d27567 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -6,6 +6,9 @@ pub mod function_trace; mod marshal; mod routes; +#[cfg(test)] +mod architecture; + use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use pyo3::prelude::*; use pyo3::types::PyAny; diff --git a/litellm-rust/crates/python-bridge/src/routes/runtime.rs b/litellm-rust/crates/python-bridge/src/routes/runtime.rs deleted file mode 100644 index 87a0c3e0104..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/runtime.rs +++ /dev/null @@ -1,423 +0,0 @@ -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(super) fn run_sync( - py: Python<'_>, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send + 'static, - F: Future> + Send + 'static, -{ - run_sync_on( - py, - pyo3_async_runtimes::tokio::get_runtime(), - future, - map_error, - ) -} - -fn run_sync_on( - py: Python<'_>, - runtime: &Runtime, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send + 'static, - F: Future> + 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(super) fn run_async( - py: Python<'_>, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send + 'static, - F: Future> + Send + 'static, -{ - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let result = catch_route_panic(future).await?; - let result = map_core_result(result, map_error)?; - Ok(Pythonized(result)) - }) -} - -fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { - 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_route_panic(future: F) -> PyResult> -where - F: Future>, -{ - AssertUnwindSafe(future) - .catch_unwind() - .await - .map_err(panic_to_pyerr) -} - -async fn wait_for_sync_result(future: F) -> PyResult> -where - F: Future>, -{ - let future = catch_route_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(&self, _serializer: S) -> Result - where - S: Serializer, - { - panic!("serializer panicked") - } - } - - #[pyfunction] - fn async_serialization_panic(py: Python<'_>) -> PyResult> { - run_async(py, async { Ok(PanickingOutput) }, runtime_error) - } - - #[pyfunction] - fn async_runtime_probe(py: Python<'_>) -> PyResult> { - 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>) -> 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::(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::( - py, - poll_fn(|_| -> Poll> { panic!("route future panicked") }), - runtime_error, - ) - .expect_err("panicked route should become a Python exception"); - - assert!(error.is_instance_of::(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::( - 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::(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::(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"); - }); - } -}