From bb44fe5292bd8f967bfa9823d593203bb445b35f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 10:36:59 -0700 Subject: [PATCH 01/16] wip --- litellm-rust/Cargo.lock | 1 - litellm-rust/Cargo.toml | 2 +- litellm-rust/clippy.toml | 10 ++ .../crates/host-python/src/execution.rs | 102 +++++++++--- .../crates/host-python/src/fork_gate.rs | 121 ++++++++++++++ litellm-rust/crates/host-python/src/lib.rs | 7 +- .../crates/python-bridge/src/diagnostics.rs | 18 ++- litellm-rust/crates/python-bridge/src/lib.rs | 8 +- .../python-bridge/src/routes/responses.rs | 12 +- litellm/proxy/proxy_cli.py | 5 + litellm/rust_bridge/_native.pyi | 8 + litellm/rust_bridge/fork_guard.py | 47 ++++++ tests/test_litellm/proxy/test_proxy_cli.py | 35 ++++ .../rust_bridge/test_fork_guard.py | 36 +++++ tests/test_litellm_rust/test_fork_guard.py | 150 ++++++++++++++++++ 15 files changed, 534 insertions(+), 28 deletions(-) create mode 100644 litellm-rust/clippy.toml create mode 100644 litellm-rust/crates/host-python/src/fork_gate.rs create mode 100644 litellm/rust_bridge/fork_guard.py create mode 100644 tests/test_litellm/rust_bridge/test_fork_guard.py create mode 100644 tests/test_litellm_rust/test_fork_guard.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 860f01c4ad1..ebab2a118fc 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3046,7 +3046,6 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ "base64 0.22.1", "bytes", - "futures-channel", "futures-core", "futures-util", "h2 0.4.15", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 8634dce92d0..fa2bdb4224c 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -34,7 +34,7 @@ pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" -reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" rstest_reuse = "0.7.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } diff --git a/litellm-rust/clippy.toml b/litellm-rust/clippy.toml new file mode 100644 index 00000000000..f7e3293069b --- /dev/null +++ b/litellm-rust/clippy.toml @@ -0,0 +1,10 @@ +# The Tokio runtime is reached only through `host-python/src/execution.rs`, whose fork gate +# must see every entry. Going around it makes a fork-after-use hang instead of raising. +disallowed-methods = [ + { path = "pyo3_async_runtimes::tokio::get_runtime", reason = "use litellm_host_python::run_sync / run_sync_value" }, + { path = "pyo3_async_runtimes::tokio::future_into_py", reason = "use litellm_host_python::run_async / run_async_value" }, + { path = "pyo3_async_runtimes::tokio::future_into_py_with_locals", reason = "use litellm_host_python::run_async / run_async_value" }, + { path = "pyo3_async_runtimes::tokio::local_future_into_py", reason = "use litellm_host_python::run_async / run_async_value" }, + { path = "pyo3_async_runtimes::tokio::run", reason = "use litellm_host_python::run_sync / run_sync_value" }, + { path = "pyo3_async_runtimes::tokio::run_until_complete", reason = "use litellm_host_python::run_sync / run_sync_value" }, +] diff --git a/litellm-rust/crates/host-python/src/execution.rs b/litellm-rust/crates/host-python/src/execution.rs index 45a1183acf5..083c184e37e 100644 --- a/litellm-rust/crates/host-python/src/execution.rs +++ b/litellm-rust/crates/host-python/src/execution.rs @@ -4,6 +4,7 @@ use std::pin::Pin; use std::task::{Context, Poll, Waker}; use std::time::Duration; +use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted}; use crate::{Pythonized, panic_to_pyerr, release_gil}; use futures_util::FutureExt; use pyo3::exceptions::PyRuntimeError; @@ -12,6 +13,67 @@ use serde::Serialize; use tokio::runtime::{Handle, Runtime}; use tokio::time::{self, MissedTickBehavior}; +pyo3::create_exception!( + _native, + ForkedAfterNativeRuntimeStarted, + PyRuntimeError, + "This process was forked after the native runtime started. Runtime threads do not survive fork(), so native routes cannot run here." +); + +pyo3::create_exception!( + _native, + ProcessReservedForForking, + PyRuntimeError, + "This process was reserved for forking workers, so native routes cannot run here." +); + +static FORK_GATE: ForkGate = ForkGate::new(); + +/// Whether this process has started the Tokio runtime. +pub fn runtime_started() -> bool { + FORK_GATE.started(std::process::id()) +} + +/// Declares that this process exists to fork workers, so it must never start the runtime. +/// Fails if it already has. Workers are unaffected: the reservation is keyed by pid. +pub fn reserve_process_for_forking() -> Result<(), RuntimeAlreadyStarted> { + FORK_GATE.reserve(std::process::id()) +} + +/// The only door to the Tokio runtime: every route reaches it through this module, which is +/// what lets the gate speak for the whole extension. `clippy.toml` disallows going around it. +fn enter_runtime() -> PyResult<()> { + FORK_GATE + .enter(std::process::id()) + .map_err(|refused| match refused { + Refused::ReservedForForking => ProcessReservedForForking::new_err( + "this process is reserved for forking workers and cannot run native routes; \ + move the call into a worker, after the fork", + ), + Refused::ForkedAfterStart => ForkedAfterNativeRuntimeStarted::new_err( + "this process was forked after the native runtime started, and runtime threads \ + do not survive fork(); start workers with spawn or forkserver, or fork before \ + the first native call", + ), + }) +} + +#[expect(clippy::disallowed_methods, reason = "this is the gated door")] +fn runtime() -> PyResult<&'static Runtime> { + enter_runtime()?; + Ok(pyo3_async_runtimes::tokio::get_runtime()) +} + +#[expect(clippy::disallowed_methods, reason = "this is the gated door")] +fn future_into_py(py: Python<'_>, future: F) -> PyResult> +where + F: Future> + Send + 'static, + T: for<'py> IntoPyObject<'py> + Send + 'static, +{ + enter_runtime()?; + pyo3_async_runtimes::tokio::future_into_py(py, future) +} + pub fn run_sync( py: Python<'_>, future: F, @@ -22,12 +84,7 @@ where E: Send + 'static, F: Future> + Send + 'static, { - run_sync_on( - py, - pyo3_async_runtimes::tokio::get_runtime(), - future, - map_error, - ) + run_sync_on(py, runtime()?, future, map_error) } pub fn run_sync_value(py: Python<'_>, future: F) -> PyResult @@ -35,7 +92,7 @@ where T: Send + 'static, F: Future> + Send + 'static, { - run_sync_value_on(py, pyo3_async_runtimes::tokio::get_runtime(), future) + run_sync_value_on(py, runtime()?, future) } fn run_sync_value_on(py: Python<'_>, runtime: &Runtime, future: F) -> PyResult @@ -83,7 +140,7 @@ where E: Send + 'static, F: Future> + Send + 'static, { - pyo3_async_runtimes::tokio::future_into_py(py, async move { + future_into_py(py, async move { let result = catch_future_panic(future).await?; let result = map_core_result(result, map_error)?; Ok(Pythonized(result)) @@ -95,7 +152,7 @@ where T: for<'py> IntoPyObject<'py> + Send + 'static, F: Future> + Send + 'static, { - pyo3_async_runtimes::tokio::future_into_py(py, async move { catch_future_panic(future).await? }) + future_into_py(py, async move { catch_future_panic(future).await? }) } pub fn poll_async_value(py: Python<'_>, future: Pin<&mut F>) -> PyResult> @@ -103,8 +160,9 @@ where T: Send, F: Future> + Send, { + let runtime = runtime()?; let result = release_gil(py, || { - let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter(); + let _runtime = runtime.enter(); std::panic::catch_unwind(AssertUnwindSafe(|| { future.poll(&mut Context::from_waker(Waker::noop())) })) @@ -286,27 +344,25 @@ mod tests { } #[pyfunction] - fn runtime_worker_count() -> usize { - pyo3_async_runtimes::tokio::get_runtime() - .metrics() - .num_workers() + fn runtime_worker_count() -> PyResult { + Ok(runtime()?.metrics().num_workers()) } #[pyfunction] - fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool { + fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> PyResult { 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; + return Ok(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 { + runtime()?.spawn(async move { let _ = heartbeat_tx.send(()); }); - heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok() + Ok(heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()) } fn extract_bool(py: Python<'_>, result: PyResult>) -> bool { @@ -317,6 +373,16 @@ mod tests { .expect("result should convert") } + #[rstest] + fn reaching_the_runtime_marks_the_process_as_started( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { + run_sync_value(py, async { Ok(()) }).unwrap(); + assert!(runtime_started()); + }); + } + #[rstest] fn inline_poll_releases_gil_and_enters_runtime( #[from(initialized_python)] python: &InitializedPython, diff --git a/litellm-rust/crates/host-python/src/fork_gate.rs b/litellm-rust/crates/host-python/src/fork_gate.rs new file mode 100644 index 00000000000..62284e978ff --- /dev/null +++ b/litellm-rust/crates/host-python/src/fork_gate.rs @@ -0,0 +1,121 @@ +use std::sync::atomic::{AtomicU32, Ordering}; + +const UNSET: u32 = 0; + +/// Decides which process may use the Tokio runtime. Its worker threads do not survive +/// `fork()`: a child forked after they started hangs on its first native call. The gate turns +/// both halves of that hazard into errors, keyed by pid so a fork needs no hook to be seen: +/// a process reserved for forking can never start the runtime, and a child of a process that +/// did start it is refused instead of hanging. +pub(crate) struct ForkGate { + runtime_pid: AtomicU32, + fork_only_pid: AtomicU32, +} + +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum Refused { + ReservedForForking, + ForkedAfterStart, +} + +#[derive(Debug, PartialEq, Eq)] +pub struct RuntimeAlreadyStarted; + +impl ForkGate { + pub(crate) const fn new() -> Self { + Self { + runtime_pid: AtomicU32::new(UNSET), + fork_only_pid: AtomicU32::new(UNSET), + } + } + + /// Claims the runtime for `pid`. Claim first, then look for a reservation: `reserve` does + /// the mirror image, so when the two race at least one of them sees the other. + pub(crate) fn enter(&self, pid: u32) -> Result<(), Refused> { + match self + .runtime_pid + .compare_exchange(UNSET, pid, Ordering::SeqCst, Ordering::SeqCst) + { + Err(owner) if owner != pid => return Err(Refused::ForkedAfterStart), + _ => {} + } + + if self.fork_only_pid.load(Ordering::SeqCst) == pid { + // Nothing was started, so the workers forked from here must still find it unclaimed. + let _ = + self.runtime_pid + .compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst); + return Err(Refused::ReservedForForking); + } + + Ok(()) + } + + pub(crate) fn reserve(&self, pid: u32) -> Result<(), RuntimeAlreadyStarted> { + self.fork_only_pid.store(pid, Ordering::SeqCst); + if self.runtime_pid.load(Ordering::SeqCst) == pid { + return Err(RuntimeAlreadyStarted); + } + Ok(()) + } + + pub(crate) fn started(&self, pid: u32) -> bool { + self.runtime_pid.load(Ordering::SeqCst) == pid + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const MASTER: u32 = 100; + const WORKER: u32 = 101; + + #[test] + fn unreserved_process_starts_the_runtime_and_stays_started() { + let gate = ForkGate::new(); + + assert!(!gate.started(MASTER)); + assert_eq!(gate.enter(MASTER), Ok(())); + assert_eq!(gate.enter(MASTER), Ok(())); + assert!(gate.started(MASTER)); + } + + #[test] + fn reserved_process_can_never_start_the_runtime() { + let gate = ForkGate::new(); + + assert_eq!(gate.reserve(MASTER), Ok(())); + assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking)); + assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking)); + assert!(!gate.started(MASTER)); + } + + #[test] + fn workers_forked_from_a_reserved_process_start_their_own_runtime() { + let gate = ForkGate::new(); + gate.reserve(MASTER).unwrap(); + gate.enter(MASTER).unwrap_err(); + + assert_eq!(gate.enter(WORKER), Ok(())); + assert!(gate.started(WORKER)); + } + + #[test] + fn reserving_after_the_runtime_started_is_refused() { + let gate = ForkGate::new(); + gate.enter(MASTER).unwrap(); + + assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted)); + } + + #[test] + fn child_forked_after_the_runtime_started_is_refused_instead_of_hanging() { + let gate = ForkGate::new(); + gate.enter(MASTER).unwrap(); + + assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart)); + assert!(!gate.started(WORKER)); + assert_eq!(gate.enter(MASTER), Ok(())); + } +} diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 583a4eb91b6..4e6337d916d 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -8,6 +8,7 @@ mod argument; mod callable; mod driver; mod execution; +mod fork_gate; mod gil; mod handle; mod marshal; @@ -18,7 +19,11 @@ pub use adapter::{ pub use argument::lookup; pub use callable::wrap_failure; pub use driver::run_call; -pub use execution::{poll_async_value, run_async, run_async_value, run_sync, run_sync_value}; +pub use execution::{ + ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, poll_async_value, reserve_process_for_forking, run_async, run_async_value, run_sync, + run_sync_value, runtime_started, +}; +pub use fork_gate::RuntimeAlreadyStarted; pub use gil::{release_count, release_gil}; pub use handle::{Execution, ExecutionBody, ExecutionStep}; pub use marshal::{Pythonized, from_py, from_py_argument, panic_to_pyerr, to_py}; diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs index 39fa8bc3596..687a090e768 100644 --- a/litellm-rust/crates/python-bridge/src/diagnostics.rs +++ b/litellm-rust/crates/python-bridge/src/diagnostics.rs @@ -1,5 +1,5 @@ -use litellm_host_python::release_count; -use pyo3::{prelude::*, types::PyDict}; +use litellm_host_python::{release_count, runtime_started}; +use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict}; #[pyfunction] pub(crate) fn gil_stats(py: Python<'_>) -> PyResult> { @@ -8,6 +8,20 @@ pub(crate) fn gil_stats(py: Python<'_>) -> PyResult> { Ok(stats.into_any().unbind()) } +/// True once this process has started the native runtime, which does not survive `fork()`. +#[pyfunction] +pub(crate) fn process_state_started() -> bool { + runtime_started() +} + +/// Declares that this process only forks workers: from now on every native route raises here, +/// so the runtime can never start. Raises if it already has. Forked workers are unaffected. +#[pyfunction] +pub(crate) fn reserve_process_for_forking() -> PyResult<()> { + litellm_host_python::reserve_process_for_forking() + .map_err(|_| PyRuntimeError::new_err("the native runtime already started in this process")) +} + #[cfg(feature = "panic-test")] #[pyfunction] pub(crate) fn _panic_for_test() { diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 7eba0d201be..a41e1500f04 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -13,10 +13,12 @@ mod _native { #[pymodule_export] use crate::diagnostics::_panic_for_test; #[pymodule_export] - use crate::diagnostics::gil_stats; + use crate::diagnostics::{gil_stats, process_state_started, reserve_process_for_forking}; #[pymodule_export] use crate::errors::{RustBridgeDeclined, RustUpstreamError}; #[pymodule_export] + use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{ @@ -50,6 +52,8 @@ mod tests { let mut expected = vec![ "RustBridgeDeclined", "RustUpstreamError", + "ForkedAfterNativeRuntimeStarted", + "ProcessReservedForForking", "ocr", "aocr", "transcription", @@ -62,6 +66,8 @@ mod tests { "ResponsesWebSocketConnection", "TokenCounter", "gil_stats", + "process_state_started", + "reserve_process_for_forking", ]; expected.sort_unstable(); diff --git a/litellm-rust/crates/python-bridge/src/routes/responses.rs b/litellm-rust/crates/python-bridge/src/routes/responses.rs index 9c10d58de4f..2e7e8fcbc21 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses.rs @@ -25,7 +25,7 @@ impl ResponsesWebSocketConnection { ) -> PyResult> { let headers = marshal_headers(headers)?; let timeout = optional_timeout(timeout_seconds); - pyo3_async_runtimes::tokio::future_into_py(py, async move { + litellm_host_python::run_async_value(py, async move { let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) .await .map_err(responses_error_to_pyerr)?; @@ -35,7 +35,7 @@ impl ResponsesWebSocketConnection { fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult> { let inner = self.inner.clone(); - pyo3_async_runtimes::tokio::future_into_py(py, async move { + litellm_host_python::run_async_value(py, async move { inner .send_text(text) .await @@ -45,14 +45,14 @@ impl ResponsesWebSocketConnection { fn recv_text<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); - pyo3_async_runtimes::tokio::future_into_py(py, async move { + litellm_host_python::run_async_value(py, async move { inner.recv_text().await.map_err(responses_error_to_pyerr) }) } fn close<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); - pyo3_async_runtimes::tokio::future_into_py(py, async move { + litellm_host_python::run_async_value(py, async move { inner.close().await.map_err(responses_error_to_pyerr) }) } @@ -68,6 +68,10 @@ mod tests { use tokio_tungstenite::{accept_async, tungstenite::Message}; #[test] + #[expect( + clippy::disallowed_methods, + reason = "the test server shares the routes' runtime" + )] fn responses_websocket_connection_round_trips_through_python() { Python::initialize(); let runtime = pyo3_async_runtimes::tokio::get_runtime(); diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 9f2e4c9802e..0477b6c62e9 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -589,6 +589,11 @@ class ProxyInitializationHelpers: gunicorn_options["certfile"] = ssl_certfile_path gunicorn_options["keyfile"] = ssl_keyfile_path + # The master preloads the app and then forks every worker, so native routes are + # forbidden in it: their runtime threads would not survive the fork. + from litellm.rust_bridge.fork_guard import reserve_process_for_forking + + reserve_process_for_forking("the gunicorn master") start_query_engine_reaper() StandaloneApplication(app=app, options=gunicorn_options).run() # Run gunicorn diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 9f959c056de..c0a06364261 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -9,6 +9,8 @@ from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMe class RustBridgeDeclined(Exception): ... class RustUpstreamError(Exception): ... +class ForkedAfterNativeRuntimeStarted(RuntimeError): ... +class ProcessReservedForForking(RuntimeError): ... def ocr( request: LiteLLMOcrRequest, @@ -101,8 +103,12 @@ class TokenCounter: def acount_request(self, body: bytes) -> Future[dict[str, object]]: ... def gil_stats() -> dict[str, int]: ... +def process_state_started() -> bool: ... +def reserve_process_for_forking() -> None: ... __all__ = [ + "ForkedAfterNativeRuntimeStarted", + "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", "RustUpstreamError", @@ -116,5 +122,7 @@ __all__ = [ "gil_stats", "messages", "ocr", + "process_state_started", + "reserve_process_for_forking", "transcription", ] diff --git a/litellm/rust_bridge/fork_guard.py b/litellm/rust_bridge/fork_guard.py new file mode 100644 index 00000000000..c94665fb8db --- /dev/null +++ b/litellm/rust_bridge/fork_guard.py @@ -0,0 +1,47 @@ +"""Fork safety of the Rust extension. + +Its runtime threads do not survive ``fork``, so a child forked after the first native call +cannot run native routes: it raises ``ForkedAfterNativeRuntimeStarted`` instead of hanging. +Fork before the first native call, or start workers with ``spawn`` / ``forkserver``. + +A process whose job is to fork workers (the gunicorn master under ``preload``) reserves itself: +from then on any native route called in it raises ``ProcessReservedForForking`` at the call +site, so the runtime can never start there. Workers forked from it are unaffected. +""" + +from __future__ import annotations + +from typing import Final + +from litellm.rust_bridge.loader import get_native_bridge + + +class NativeStateStartedBeforeFork(RuntimeError): + pass + + +class _NeverRaised(RuntimeError): + """Stands in for a native exception when the extension is unavailable or predates it.""" + + +_native: Final = get_native_bridge() +ForkedAfterNativeRuntimeStarted: Final[type[RuntimeError]] = getattr( + _native, "ForkedAfterNativeRuntimeStarted", _NeverRaised +) +ProcessReservedForForking: Final[type[RuntimeError]] = getattr(_native, "ProcessReservedForForking", _NeverRaised) + + +def reserve_process_for_forking(where: str) -> None: + """Forbid native routes in this process. Raises if one already ran here.""" + native: Final = get_native_bridge() + reserve: Final = getattr(native, "reserve_process_for_forking", None) + if not callable(reserve): + return + try: + reserve() + except RuntimeError as error: + raise NativeStateStartedBeforeFork( + f"The LiteLLM Rust extension already ran a native route in {where}, and its runtime " + "threads do not survive fork(). Move the native call (warm-up, health check, " + "import-time initialization) into the worker, after the fork." + ) from error diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 712c526b244..c806725d594 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -21,6 +21,15 @@ from uvicorn.importer import import_from_string from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server +@pytest.fixture(autouse=True) +def fork_reservation(): + """Reserving is irreversible: it would forbid native routes in this pytest worker for good""" + with patch( # test-quality-ok: process-global native state, a real reservation would poison every later test in the worker + "litellm.rust_bridge.fork_guard.reserve_process_for_forking" + ) as reserve: + yield reserve + + @pytest.mark.xdist_group("proxy_cli") class TestProxyInitializationHelpers: @patch("importlib.metadata.version") @@ -1574,6 +1583,32 @@ class TestProxyInitializationHelpers: assert captured["options"]["max_requests"] == 1000 assert captured["options"]["max_requests_jitter"] == 50 + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") + def test_gunicorn_master_is_reserved_for_forking_before_it_runs(self, fork_reservation): + """preload forks workers from the master, so native routes are forbidden there first""" + pytest.importorskip("gunicorn") + reserved_before_run: list = [] + + def capture_run(self): + reserved_before_run.append(fork_reservation.call_args) + + with ( + patch("gunicorn.app.base.BaseApplication.run", capture_run), + patch( # test-quality-ok: option tests must not start a thread or change the pytest worker's child ownership + "litellm.proxy.proxy_cli.start_query_engine_reaper" + ), + ): + ProxyInitializationHelpers._run_gunicorn_server( + host="127.0.0.1", + port=4012, + app=MagicMock(), + num_workers=2, + ssl_certfile_path=None, + ssl_keyfile_path=None, + ) + + assert [call.args for call in reserved_before_run] == [("the gunicorn master",)] + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") def test_gunicorn_jitter_without_base_warns(self): """gunicorn path warns when jitter is set without --max_requests_before_restart""" diff --git a/tests/test_litellm/rust_bridge/test_fork_guard.py b/tests/test_litellm/rust_bridge/test_fork_guard.py new file mode 100644 index 00000000000..54bfd54c230 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_fork_guard.py @@ -0,0 +1,36 @@ +from types import SimpleNamespace + +import pytest + +from litellm.rust_bridge import fork_guard + + +def _reserve_with(monkeypatch: pytest.MonkeyPatch, native: object) -> None: + monkeypatch.setattr(fork_guard, "get_native_bridge", lambda: native) + fork_guard.reserve_process_for_forking("the gunicorn master") + + +def test_missing_extension_has_nothing_to_reserve(monkeypatch: pytest.MonkeyPatch) -> None: + _reserve_with(monkeypatch, None) + + +def test_extension_built_before_reservation_existed_passes(monkeypatch: pytest.MonkeyPatch) -> None: + _reserve_with(monkeypatch, SimpleNamespace()) + + +def test_unused_extension_is_reserved(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[None] = [] + + _reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=lambda: calls.append(None))) + + assert calls == [None] + + +def test_used_extension_refuses_and_names_the_place(monkeypatch: pytest.MonkeyPatch) -> None: + def reserve() -> None: + raise RuntimeError("the native runtime already started in this process") + + with pytest.raises(fork_guard.NativeStateStartedBeforeFork, match="the gunicorn master") as raised: + _reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=reserve)) + + assert isinstance(raised.value.__cause__, RuntimeError) diff --git a/tests/test_litellm_rust/test_fork_guard.py b/tests/test_litellm_rust/test_fork_guard.py new file mode 100644 index 00000000000..b92095cbaaa --- /dev/null +++ b/tests/test_litellm_rust/test_fork_guard.py @@ -0,0 +1,150 @@ +import os +import subprocess +import sys +import textwrap + +import pytest + +pytestmark = pytest.mark.requires_rust_extension + +_NATIVE_CONTRACT = textwrap.dedent( + """ + import os + from litellm.rust_bridge import _native + from litellm.rust_bridge.fork_guard import reserve_process_for_forking + + def native_route_error(): + import asyncio + + async def call(): + await _native.ResponsesWebSocketConnection.connect("ws://127.0.0.1:1", {}, 0.2) + + try: + asyncio.run(call()) + except Exception as error: + return f"{type(error).__name__}: {error}" + return "" + + assert _native.process_state_started() is False + reserve_process_for_forking("the test master") + assert native_route_error().startswith("ProcessReservedForForking: ") + assert _native.process_state_started() is False + + pid = os.fork() + if pid == 0: + error = native_route_error() + started = _native.process_state_started() + os._exit(0 if started and "reserved" not in error and "forked" not in error else 1) + assert os.waitpid(pid, 0)[1] == 0 + + pid = os.fork() + if pid == 0: + native_route_error() + grandchild = os.fork() + if grandchild == 0: + os._exit(0 if native_route_error().startswith("ForkedAfterNativeRuntimeStarted: ") else 1) + os._exit(os.waitpid(grandchild, 0)[1]) + assert os.waitpid(pid, 0)[1] == 0 + """ +) + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork only") +def test_compiled_extension_forbids_the_master_and_frees_its_workers() -> None: + env = {**os.environ, "OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES"} + + result = subprocess.run( + [sys.executable, "-c", _NATIVE_CONTRACT], capture_output=True, text=True, timeout=60, env=env + ) + + assert result.returncode == 0, result.stderr + + +_SDK_CONTRACT = textwrap.dedent( + """ + import asyncio, json, multiprocessing, os, threading + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + + import litellm + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + self.rfile.read(int(self.headers["Content-Length"])) + if self.headers.get("User-Agent", "").startswith("python-httpx"): + self.send_response(418) + self.end_headers() + return + body = json.dumps({ + "pages": [{"index": 0, "markdown": "native", "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, + }).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + arguments = { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{server.server_port}", + "num_retries": 0, + } + litellm.rust(True) + + SERVED, REFUSED, OTHER = 0, 3, 4 + + def outcome(asynchronous): + try: + response = asyncio.run(litellm.aocr(**arguments)) if asynchronous else litellm.ocr(**arguments) + except ForkedAfterNativeRuntimeStarted: + return REFUSED + except Exception: + return OTHER + return SERVED if response.pages[0].markdown == "native" else OTHER + + def forked(asynchronous): + pid = os.fork() + if pid == 0: + os._exit(outcome(asynchronous)) + return os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1]) + + def pooled(asynchronous): + with multiprocessing.get_context("fork").Pool(1) as pool: + return pool.apply(outcome, (asynchronous,)) + + # Forking before the first native call is fine: the child starts its own runtime. + assert [forked(False), forked(True)] == [SERVED, SERVED] + + assert outcome(False) == SERVED + # After it, a forked child is told so instead of hanging on threads that do not exist. + assert [forked(False), forked(True)] == [REFUSED, REFUSED] + assert [pooled(False), pooled(True)] == [REFUSED, REFUSED] + # The parent is not poisoned by any of it. + assert [outcome(False), outcome(True)] == [SERVED, SERVED] + """ +) + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork only") +def test_sdk_call_in_a_child_forked_after_native_use_raises_instead_of_hanging() -> None: + env = { + **os.environ, + "OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES", + "LITELLM_RUST": "1", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + } + + result = subprocess.run( + [sys.executable, "-c", _SDK_CONTRACT], capture_output=True, text=True, timeout=120, env=env + ) + + assert result.returncode == 0, result.stderr From 752647d1467624fd794d653bed9933b6c8c8037a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 10:41:41 -0700 Subject: [PATCH 02/16] wip --- litellm-rust/crates/host-python/src/lib.rs | 5 +++-- litellm-rust/crates/python-bridge/src/lib.rs | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 4e6337d916d..7d164ab7535 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -20,8 +20,9 @@ pub use argument::lookup; pub use callable::wrap_failure; pub use driver::run_call; pub use execution::{ - ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, poll_async_value, reserve_process_for_forking, run_async, run_async_value, run_sync, - run_sync_value, runtime_started, + ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, poll_async_value, + reserve_process_for_forking, run_async, run_async_value, run_sync, run_sync_value, + runtime_started, }; pub use fork_gate::RuntimeAlreadyStarted; pub use gil::{release_count, release_gil}; diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index a41e1500f04..46f98736aa1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -17,8 +17,6 @@ mod _native { #[pymodule_export] use crate::errors::{RustBridgeDeclined, RustUpstreamError}; #[pymodule_export] - use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; - #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{ @@ -32,6 +30,8 @@ mod _native { use crate::routes::responses::ResponsesWebSocketConnection; #[pymodule_export] use crate::token_counter::TokenCounter; + #[pymodule_export] + use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; } use pyo3::prelude::*; From 1bcd8d704fe4ee0a791cec1f2e96221495f94f8e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 18:08:19 +0000 Subject: [PATCH 03/16] test: run fork-guard contract subprocesses with python -I Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm_rust/test_fork_guard.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm_rust/test_fork_guard.py b/tests/test_litellm_rust/test_fork_guard.py index b92095cbaaa..c15555ff535 100644 --- a/tests/test_litellm_rust/test_fork_guard.py +++ b/tests/test_litellm_rust/test_fork_guard.py @@ -54,7 +54,7 @@ def test_compiled_extension_forbids_the_master_and_frees_its_workers() -> None: env = {**os.environ, "OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES"} result = subprocess.run( - [sys.executable, "-c", _NATIVE_CONTRACT], capture_output=True, text=True, timeout=60, env=env + [sys.executable, "-I", "-c", _NATIVE_CONTRACT], capture_output=True, text=True, timeout=60, env=env ) assert result.returncode == 0, result.stderr @@ -144,7 +144,7 @@ def test_sdk_call_in_a_child_forked_after_native_use_raises_instead_of_hanging() } result = subprocess.run( - [sys.executable, "-c", _SDK_CONTRACT], capture_output=True, text=True, timeout=120, env=env + [sys.executable, "-I", "-c", _SDK_CONTRACT], capture_output=True, text=True, timeout=120, env=env ) assert result.returncode == 0, result.stderr From 18a1491bd2b3cb2ddc9a493e712c92393f970d4c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 18:17:54 +0000 Subject: [PATCH 04/16] test(rust): pin child interpreters to the parent's litellm and lint for it Children spawned as [sys.executable, -c, ...] put the working directory first on sys.path, so under 'make test-rust-extension' a source checkout shadows the installed wheel and the child imports a litellm with no compiled extension. A shared helper spawns them with -I and asserts the child resolved the same litellm.__file__ as the parent, and a new TQ009 rule flags un-isolated sys.executable spawns. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- scripts/check_test_quality.py | 40 +++++++++++++++++++ test-quality-budget.json | 3 ++ .../rust_bridge/test_fork_guard.py | 4 +- tests/test_litellm/test_check_test_quality.py | 25 ++++++++++++ .../support/child_interpreter.py | 36 +++++++++++++++++ tests/test_litellm_rust/test_fork_guard.py | 12 ++---- 6 files changed, 110 insertions(+), 10 deletions(-) create mode 100644 tests/test_litellm_rust/support/child_interpreter.py diff --git a/scripts/check_test_quality.py b/scripts/check_test_quality.py index 41342acd23a..1ef4aed8675 100644 --- a/scripts/check_test_quality.py +++ b/scripts/check_test_quality.py @@ -60,6 +60,13 @@ TQ007 A module global that a conftest saves before every test and restores aft names are read from the keys the conftest assigns directly and from whatever the save loop iterates, including a module-level tuple or dict it names rather than spells out. +TQ009 A child interpreter spawned as `subprocess.run([sys.executable, ...])` without + `-I`/`-P` as its first flag. Without isolation the child's sys.path leads with + the working directory, so a source checkout shadows the installed package and + the child tests a different `litellm` than the parent imported -- TQ003 is the + same working-directory hazard seen from the child's side. Use + tests.test_litellm_rust.support.child_interpreter.run_child_interpreter, which + also asserts the child resolved the same `litellm.__file__` as the parent. Every rule is suppressible with `# test-quality-ok: ` on the reported line, following the repo's `*-ok: ` convention. A suppression without a @@ -140,6 +147,9 @@ SKIP_CALLS: Final = frozenset(("pytest.skip", "skip")) CONFTEST_NAME: Final = "conftest.py" SDK_MODULE: Final = "litellm" +SUBPROCESS_SPAWNS: Final = frozenset(("run", "Popen", "check_output", "check_call", "call")) +INTERPRETER_ISOLATION_FLAGS: Final = frozenset(("-I", "-P")) + CREDENTIAL_NAME_RE: Final = re.compile( r"(?:API_KEY|_KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|DATABASE_URL|ACCESS_KEY_ID)$" ) @@ -709,6 +719,35 @@ def _snapshotted_names(tree: ast.Module) -> Iterator[tuple[str, int]]: yield from _string_members(iterable) +def iter_child_interpreter_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and node.args): + continue + if _dotted_name(node.func).rsplit(".", 1)[-1] not in SUBPROCESS_SPAWNS: + continue + argv: Final = node.args[0] + if not isinstance(argv, (ast.List, ast.Tuple)) or not argv.elts: + continue + if _dotted_name(argv.elts[0]) != "sys.executable": + continue + isolated: Final = ( + len(argv.elts) > 1 + and isinstance(argv.elts[1], ast.Constant) + and argv.elts[1].value in INTERPRETER_ISOLATION_FLAGS + ) + if isolated: + continue + yield Violation( + path, + node.lineno, + "TQ009", + "child interpreter spawned without -I/-P; the working directory lands on sys.path " + "and a source checkout can shadow the installed package, use " + "tests.test_litellm_rust.support.child_interpreter.run_child_interpreter or pass -I " + f"(suppress: `# {SUPPRESSION_TOKEN}: `)", + ) + + def iter_conftest_inventory_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: if path.name != CONFTEST_NAME: return @@ -746,6 +785,7 @@ def check_file(path: Path) -> tuple[Violation, ...]: *iter_credential_skip_violations(path, tree), *iter_conftest_inventory_violations(path, tree), *iter_internal_patch_violations(path, tree), + *iter_child_interpreter_violations(path, tree), ) if violation.line not in skip ) diff --git a/test-quality-budget.json b/test-quality-budget.json index 3c12371f02f..ae4ea4d31be 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -22,5 +22,8 @@ }, "TQ008": { "limit": 10993 + }, + "TQ009": { + "limit": 59 } } diff --git a/tests/test_litellm/rust_bridge/test_fork_guard.py b/tests/test_litellm/rust_bridge/test_fork_guard.py index 54bfd54c230..88ae017ec39 100644 --- a/tests/test_litellm/rust_bridge/test_fork_guard.py +++ b/tests/test_litellm/rust_bridge/test_fork_guard.py @@ -11,11 +11,11 @@ def _reserve_with(monkeypatch: pytest.MonkeyPatch, native: object) -> None: def test_missing_extension_has_nothing_to_reserve(monkeypatch: pytest.MonkeyPatch) -> None: - _reserve_with(monkeypatch, None) + assert _reserve_with(monkeypatch, None) is None def test_extension_built_before_reservation_existed_passes(monkeypatch: pytest.MonkeyPatch) -> None: - _reserve_with(monkeypatch, SimpleNamespace()) + assert _reserve_with(monkeypatch, SimpleNamespace()) is None def test_unused_extension_is_reserved(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/test_litellm/test_check_test_quality.py index a75b1e43fb7..bfe503e74d1 100644 --- a/tests/test_litellm/test_check_test_quality.py +++ b/tests/test_litellm/test_check_test_quality.py @@ -737,3 +737,28 @@ def test_a_fanned_out_run_reports_each_generated_file_exactly_once(tmp_path): assert len(reported) == len(paths) assert len({line.split(":")[0] for line in reported}) == len(paths) assert all(" TQ001 " in line for line in reported) + + +def test_sys_executable_child_without_isolation_flag_is_flagged(tmp_path): + source = 'import subprocess, sys\nsubprocess.run([sys.executable, "-c", "pass"])\n' + assert _codes(tmp_path, source) == ["TQ009"] + + +def test_sys_executable_child_with_dash_i_is_clean(tmp_path): + source = 'import subprocess, sys\nsubprocess.run([sys.executable, "-I", "-c", "pass"])\n' + assert _codes(tmp_path, source) == [] + + +def test_sys_executable_child_with_dash_p_is_clean(tmp_path): + source = 'import subprocess, sys\nsubprocess.run([sys.executable, "-P", "-c", "pass"])\n' + assert _codes(tmp_path, source) == [] + + +def test_non_interpreter_subprocess_call_is_untouched(tmp_path): + source = 'import subprocess\nsubprocess.run(["python", "-c", "pass"])\n' + assert _codes(tmp_path, source) == [] + + +def test_popen_sys_executable_tuple_is_flagged(tmp_path): + source = 'import subprocess, sys\nsubprocess.Popen((sys.executable, "script.py"))\n' + assert _codes(tmp_path, source) == ["TQ009"] diff --git a/tests/test_litellm_rust/support/child_interpreter.py b/tests/test_litellm_rust/support/child_interpreter.py new file mode 100644 index 00000000000..26bbe03a2d8 --- /dev/null +++ b/tests/test_litellm_rust/support/child_interpreter.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import os +import subprocess +import sys +from collections.abc import Mapping +from typing import Final + +import litellm + +PARENT_LITELLM_FILE: Final = "LITELLM_TEST_PARENT_LITELLM_FILE" + +_PROLOGUE: Final = ( + "import os as _os, litellm as _litellm; _parent = _os.environ.pop({key!r}); " + 'assert _litellm.__file__ == _parent, f"child imported litellm from {{_litellm.__file__}}, parent from {{_parent}}"; ' + "del _os, _litellm, _parent\n" +) + + +def run_child_interpreter( + source: str, *, env: Mapping[str, str] | None = None, timeout: float +) -> subprocess.CompletedProcess[str]: + """Run `source` in a fresh interpreter that imports the same `litellm` as this process. + + `-I` keeps the working directory off sys.path so a source checkout cannot shadow an + installed wheel, and the prologue fails fast with both paths if the child still + resolves a different package. + """ + environment: Final = {**(os.environ if env is None else env), PARENT_LITELLM_FILE: litellm.__file__} + return subprocess.run( + [sys.executable, "-I", "-c", _PROLOGUE.format(key=PARENT_LITELLM_FILE) + source], + capture_output=True, + text=True, + timeout=timeout, + env=environment, + ) diff --git a/tests/test_litellm_rust/test_fork_guard.py b/tests/test_litellm_rust/test_fork_guard.py index c15555ff535..086397bab5c 100644 --- a/tests/test_litellm_rust/test_fork_guard.py +++ b/tests/test_litellm_rust/test_fork_guard.py @@ -1,10 +1,10 @@ import os -import subprocess -import sys import textwrap import pytest +from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter + pytestmark = pytest.mark.requires_rust_extension _NATIVE_CONTRACT = textwrap.dedent( @@ -53,9 +53,7 @@ _NATIVE_CONTRACT = textwrap.dedent( def test_compiled_extension_forbids_the_master_and_frees_its_workers() -> None: env = {**os.environ, "OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES"} - result = subprocess.run( - [sys.executable, "-I", "-c", _NATIVE_CONTRACT], capture_output=True, text=True, timeout=60, env=env - ) + result = run_child_interpreter(_NATIVE_CONTRACT, env=env, timeout=60) assert result.returncode == 0, result.stderr @@ -143,8 +141,6 @@ def test_sdk_call_in_a_child_forked_after_native_use_raises_instead_of_hanging() "LITELLM_LOCAL_MODEL_COST_MAP": "True", } - result = subprocess.run( - [sys.executable, "-I", "-c", _SDK_CONTRACT], capture_output=True, text=True, timeout=120, env=env - ) + result = run_child_interpreter(_SDK_CONTRACT, env=env, timeout=120) assert result.returncode == 0, result.stderr From 38fa8a7f551dc0a3e37d85930f3084afab97d5a4 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 18:21:32 +0000 Subject: [PATCH 05/16] fix(rust): leave the fork gate untouched when a late reservation is refused reserve() stored fork_only_pid before noticing the runtime already ran under that pid, so a refused reservation still reserved the process: the next enter() cleared the runtime claim and children forked afterwards inherited a dead runtime and hung. Undo the reservation on the error path so the gate is exactly as it was. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/host-python/src/fork_gate.rs | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/litellm-rust/crates/host-python/src/fork_gate.rs b/litellm-rust/crates/host-python/src/fork_gate.rs index 62284e978ff..cdf269deaec 100644 --- a/litellm-rust/crates/host-python/src/fork_gate.rs +++ b/litellm-rust/crates/host-python/src/fork_gate.rs @@ -51,9 +51,18 @@ impl ForkGate { Ok(()) } + /// Reserves `pid` for forking. Reserve first, then look for a started runtime: `enter` does + /// the mirror image, so when the two race at least one of them sees the other. A refused + /// reservation leaves the gate exactly as it was, so a process already running the runtime + /// keeps refusing the children it forks. pub(crate) fn reserve(&self, pid: u32) -> Result<(), RuntimeAlreadyStarted> { self.fork_only_pid.store(pid, Ordering::SeqCst); if self.runtime_pid.load(Ordering::SeqCst) == pid { + // Nothing may change for a process that already runs the runtime: its children + // must still be refused. + let _ = + self.fork_only_pid + .compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst); return Err(RuntimeAlreadyStarted); } Ok(()) @@ -109,6 +118,17 @@ mod tests { assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted)); } + #[test] + fn a_refused_reservation_leaves_the_runtime_claimed_and_its_children_refused() { + let gate = ForkGate::new(); + gate.enter(MASTER).unwrap(); + + assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted)); + assert_eq!(gate.enter(MASTER), Ok(())); + assert!(gate.started(MASTER)); + assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart)); + } + #[test] fn child_forked_after_the_runtime_started_is_refused_instead_of_hanging() { let gate = ForkGate::new(); From cc23e5781e4d09de7c744ea45dfd197e46d2fbff Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 18:22:06 +0000 Subject: [PATCH 06/16] refactor(rust): drop a comment that repeats the reserve doc Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/host-python/src/fork_gate.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm-rust/crates/host-python/src/fork_gate.rs b/litellm-rust/crates/host-python/src/fork_gate.rs index cdf269deaec..c4842dd9223 100644 --- a/litellm-rust/crates/host-python/src/fork_gate.rs +++ b/litellm-rust/crates/host-python/src/fork_gate.rs @@ -58,8 +58,6 @@ impl ForkGate { pub(crate) fn reserve(&self, pid: u32) -> Result<(), RuntimeAlreadyStarted> { self.fork_only_pid.store(pid, Ordering::SeqCst); if self.runtime_pid.load(Ordering::SeqCst) == pid { - // Nothing may change for a process that already runs the runtime: its children - // must still be refused. let _ = self.fork_only_pid .compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst); From 364d8975456548d7e2753aa13e51ab202e6fc110 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 18:26:59 +0000 Subject: [PATCH 07/16] fix(otel v2): map Responses API output onto the Langfuse generation output Responses API calls build the generation output only from response["choices"], which Responses payloads do not carry, so Langfuse rendered a blank output. Fold output[] into one assistant choice (output_text parts concatenated, function_call and custom_tool_call items as tool_calls) and derive the finish reason from status when choices are absent. Custom tool call input is now redacted alongside function call arguments under turn_off_message_logging. Carries the behavior of #41604 by @moshemorad (issue #41591) onto current main with typed conversion and single-message output. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/model/payloads.py | 82 +++++++++++- litellm/litellm_core_utils/redact_messages.py | 4 + .../otel/test_otel_v2_sources_of_truth.py | 117 ++++++++++++++++++ .../otel/test_otel_v2_vendor_mappers.py | 30 +++++ .../test_redact_messages.py | 22 ++++ 5 files changed, 253 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 467c286db9d..484f4a4c294 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -7,9 +7,11 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from enum import Enum from types import MappingProxyType -from typing import TYPE_CHECKING, ClassVar, Final, cast +from typing import TYPE_CHECKING, ClassVar, Final, Literal, cast from urllib.parse import urlsplit +from typing_extensions import ReadOnly, TypedDict + from litellm.integrations.otel.model.metadata import RequestContext, RequestIdentity from litellm.integrations.otel.model.semconv import ( GenAIOperation, @@ -25,6 +27,7 @@ from litellm.integrations.otel.model.utils import ( as_float, as_int, as_str, + as_str_mapping, as_str_tuple, ) @@ -424,7 +427,7 @@ class LLMCallSpanData: # plain ``.get`` — no repeated ``isinstance`` guards. raw_response: Final = payload.get("response") response: Final = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {}) - choices_out: Final = _dicts(response.get("choices")) + choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response) # ``finish_reasons`` is metadata, not content, so derive it from # ``choices_out`` before gating. The raw message/choice bodies are only # retained when content capture is enabled (see ``capture_span_content``); @@ -703,6 +706,81 @@ def _finish_reasons(choices: tuple[Mapping[str, object], ...]) -> tuple[str, ... return tuple(r for c in choices if (r := as_str(c.get("finish_reason")))) +class _ToolFunction(TypedDict): + name: ReadOnly[str] + arguments: ReadOnly[str] + + +class _ToolCall(TypedDict): + id: ReadOnly[str] + type: ReadOnly[Literal["function"]] + function: ReadOnly[_ToolFunction] + + +class _AssistantMessage(TypedDict): + role: ReadOnly[str] + content: ReadOnly[str | None] + tool_calls: ReadOnly[tuple[_ToolCall, ...] | None] + + +class _Choice(TypedDict): + message: ReadOnly[_AssistantMessage] + finish_reason: ReadOnly[str | None] + + +_RESPONSES_TOOL_CALL_TYPES: Final = frozenset({"function_call", "custom_tool_call"}) + + +def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: + """A Responses API ``output`` folded into one chat-shaped assistant choice.""" + items: Final = _dicts(response.get("output")) + messages: Final = tuple(item for item in items if item.get("type") == "message") + content: Final = "".join( + text + for item in messages + for part in _dicts(item.get("content")) + if part.get("type") == "output_text" + if (text := as_str(part.get("text"))) is not None + ) + tool_calls: Final = tuple( + _responses_tool_call(item) for item in items if item.get("type") in _RESPONSES_TOOL_CALL_TYPES + ) + if not messages and not tool_calls: + return () + message: Final[_AssistantMessage] = { + "role": next((role for item in messages if (role := as_str(item.get("role")))), "assistant"), + "content": content if messages else None, + "tool_calls": tool_calls or None, + } + choice: Final[_Choice] = {"message": message, "finish_reason": _responses_finish_reason(response, bool(tool_calls))} + return (choice,) + + +def _responses_tool_call(item: Mapping[str, object]) -> _ToolCall: + custom: Final = item.get("type") == "custom_tool_call" + function: Final[_ToolFunction] = { + "name": as_str(item.get("name")) or "", + "arguments": as_str(item.get("input" if custom else "arguments")) or "", + } + tool_call: Final[_ToolCall] = { + "id": as_str(item.get("call_id")) or as_str(item.get("id")) or "", + "type": "function", + "function": function, + } + return tool_call + + +def _responses_finish_reason(response: Mapping[str, object], has_tool_calls: bool) -> str | None: + status: Final = as_str(response.get("status")) + if status == "completed": + return "tool_calls" if has_tool_calls else "stop" + if status != "incomplete": + return None + details: Final = as_str_mapping(response.get("incomplete_details")) + reason: Final = details.get("reason") if details is not None else None + return "content_filter" if reason == "content_filter" else "length" + + def _parse_error(payload: StandardLoggingPayload) -> SpanError | None: """A ``SpanError`` for a failed request, or ``None`` on success.""" if payload.get("status") != "failure": diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 9d22a5ddef5..1f9464a2a26 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -138,6 +138,8 @@ def _redact_responses_api_output(output_items): if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"): output_item.arguments = REDACTED_BY_LITELLM + if hasattr(output_item, "type") and output_item.type == "custom_tool_call" and hasattr(output_item, "input"): + output_item.input = REDACTED_BY_LITELLM def _redact_responses_api_output_dict(output_items, redacted_str: str): @@ -161,6 +163,8 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): if output_item.get("type") == "function_call" and "arguments" in output_item: output_item["arguments"] = redacted_str + if output_item.get("type") == "custom_tool_call" and "input" in output_item: + output_item["input"] = redacted_str def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[str, object]: diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index f4a8691f72f..972c91670f8 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -738,6 +738,123 @@ def test_embedding_summary_is_absent_without_vectors_and_for_chat_data_lists(): assert chat.embedding_output is None +def _responses_payload(output: list[object], status: str = "completed", **response_fields: object): + return _sample_payload( + call_type="aresponses", + model="gpt-5.4-nano", + response={"id": "resp_1", "object": "response", "status": status, "output": output, **response_fields}, + ) + + +_RESPONSES_TEXT_ITEM = { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "po", "annotations": []}, {"type": "output_text", "text": "ng"}], +} + + +def test_responses_output_text_becomes_one_assistant_choice_with_stop(): + data = LLMCallSpanData.from_standard_logging_payload( + _responses_payload([{"type": "reasoning", "summary": []}, _RESPONSES_TEXT_ITEM]), capture_content=True + ) + + assert json.loads(json.dumps(data.choices_out)) == [ + { + "message": {"role": "assistant", "content": "pong", "tool_calls": None}, + "finish_reason": "stop", + } + ] + assert data.finish_reasons == ("stop",) + assert data.response_id == "resp_1" + + +def test_responses_tool_calls_fold_into_the_assistant_message_with_tool_calls_finish_reason(): + data = LLMCallSpanData.from_standard_logging_payload( + _responses_payload( + [ + _RESPONSES_TEXT_ITEM, + {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city": "sf"}'}, + {"type": "custom_tool_call", "call_id": "call_2", "name": "grep", "input": "-r TODO"}, + ] + ), + capture_content=True, + ) + + assert len(data.choices_out) == 1 + message = data.choices_out[0]["message"] + assert message["content"] == "pong" + assert json.loads(json.dumps(message["tool_calls"])) == [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "sf"}'}}, + {"id": "call_2", "type": "function", "function": {"name": "grep", "arguments": "-r TODO"}}, + ] + assert data.finish_reasons == ("tool_calls",) + + +def test_responses_tool_call_only_output_has_no_content(): + data = LLMCallSpanData.from_standard_logging_payload( + _responses_payload([{"type": "function_call", "id": "fc_1", "name": "get_weather", "arguments": "{}"}]), + capture_content=True, + ) + + assert data.choices_out[0]["message"]["content"] is None + assert data.choices_out[0]["message"]["tool_calls"][0]["id"] == "fc_1" + + +@pytest.mark.parametrize( + ("status", "response_fields", "expected"), + [ + ("incomplete", {"incomplete_details": {"reason": "max_output_tokens"}}, ("length",)), + ("incomplete", {"incomplete_details": {"reason": "content_filter"}}, ("content_filter",)), + ("incomplete", {}, ("length",)), + ("failed", {}, ()), + ], +) +def test_responses_status_maps_to_finish_reasons(status, response_fields, expected): + data = LLMCallSpanData.from_standard_logging_payload( + _responses_payload([_RESPONSES_TEXT_ITEM], status=status, **response_fields), capture_content=True + ) + + assert data.finish_reasons == expected + assert data.choices_out[0]["message"]["content"] == "pong" + + +def test_responses_output_follows_the_content_capture_gate_but_finish_reasons_do_not(): + data = LLMCallSpanData.from_standard_logging_payload(_responses_payload([_RESPONSES_TEXT_ITEM])) + + assert data.choices_out == () + assert data.finish_reasons == ("stop",) + + +def test_responses_content_only_reads_output_text_parts(): + item = { + "type": "message", + "role": "assistant", + "content": [{"type": "refusal", "refusal": "no", "text": "not output"}, {"type": "output_text", "text": "ok"}], + } + data = LLMCallSpanData.from_standard_logging_payload(_responses_payload([item]), capture_content=True) + + assert data.choices_out[0]["message"]["content"] == "ok" + + +def test_responses_output_without_messages_or_tool_calls_stays_empty(): + data = LLMCallSpanData.from_standard_logging_payload( + _responses_payload([{"type": "reasoning", "summary": []}]), capture_content=True + ) + + assert data.choices_out == () + assert data.finish_reasons == () + + +def test_chat_choices_win_over_a_responses_output_list(): + payload = _sample_payload(response={"choices": [{"finish_reason": "stop", "message": {"content": "chat"}}]}) + payload["response"]["output"] = [_RESPONSES_TEXT_ITEM] + data = LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True) + + assert data.choices_out[0]["message"]["content"] == "chat" + assert data.finish_reasons == ("stop",) + + def test_request_identity_prefers_canonical_team_keys(): from litellm.integrations.otel.model.payloads import RequestIdentity diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py index c5ebc4bc53a..5b4d1e7a802 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py @@ -196,6 +196,36 @@ def test_langfuse_mapper_keeps_chat_output_when_no_embedding_summary(): assert json.loads(attrs["langfuse.observation.output"]) == [{"role": "assistant", "content": "Sunny."}] +def test_langfuse_mapper_renders_a_responses_api_call_from_the_standard_logging_payload(): + payload = { + "call_type": "aresponses", + "custom_llm_provider": "openai", + "model": "gpt-5.4-nano", + "messages": [{"role": "user", "content": "weather in sf?"}], + "response": { + "id": "resp_1", + "status": "completed", + "output": [ + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Checking."}]}, + {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city": "sf"}'}, + ], + }, + } + data = LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True) + attrs = LangfuseMapper().map(data) + + assert json.loads(attrs["langfuse.observation.output"]) == [ + { + "role": "assistant", + "content": "Checking.", + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "sf"}'}} + ], + } + ] + assert attrs["langfuse.observation.type"] == "generation" + + # --------------------------------------------------------------------------- # # Weave # --------------------------------------------------------------------------- # diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 584a3ac471c..c6c9a9dd2b7 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -493,6 +493,20 @@ class TestPerformRedaction: assert redacted["output"][0]["arguments"] == "redacted-by-litellm" assert redacted["output"][0]["name"] == "get_weather" + def test_redacts_responses_api_custom_tool_call_input_dict(self): + result = { + "output": [ + {"type": "custom_tool_call", "name": "grep", "input": "-r secret-token", "call_id": "call_1"}, + {"type": "function_call", "name": "get_weather", "input": "not-a-custom-input", "call_id": "call_2"}, + ] + } + + redacted = perform_redaction({}, result) + + assert redacted["output"][0]["input"] == "redacted-by-litellm" + assert redacted["output"][0]["name"] == "grep" + assert redacted["output"][1]["input"] == "not-a-custom-input" + def test_redacts_every_tool_call_in_multi_element_list(self): result = litellm.ModelResponse( id="resp-multi", @@ -563,6 +577,14 @@ class TestPerformRedaction: assert output_item.arguments == "redacted-by-litellm" assert output_item.name == "get_weather" + def test_redacts_responses_api_custom_tool_call_input_object(self): + output_item = SimpleNamespace(type="custom_tool_call", name="grep", input="-r secret-token", call_id="call_1") + + _redact_responses_api_output([output_item]) + + assert output_item.input == "redacted-by-litellm" + assert output_item.name == "grep" + def test_redacts_response_output_objects_with_top_level_text(self): output_items = [ SimpleNamespace(text="top-level output"), From d67d9984f710b90527dea9b7e1ed43b5aace0888 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 18:38:37 +0000 Subject: [PATCH 08/16] test: expect TQ009 in the shipped quality budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_test_quality_gate.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/test_litellm/test_test_quality_gate.py index 6652211a828..cde33787c6c 100644 --- a/tests/test_litellm/test_test_quality_gate.py +++ b/tests/test_litellm/test_test_quality_gate.py @@ -136,7 +136,9 @@ def test_the_shipped_budget_covers_every_rule_the_checker_can_emit(): import json budget = json.loads((_REPO_ROOT / "test-quality-budget.json").read_text()) - assert set(budget) == {"TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ008"} + assert set(budget) == { + "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ008", "TQ009" + } assert all(spec["limit"] >= 0 for spec in budget.values()) From 47d06d9fdd5973ea2e72daec07b1a38bae24b2bd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 11:56:02 -0700 Subject: [PATCH 09/16] test(unified_google_tests): use the Vertex global endpoint and retry 429s with backoff The google_generate_content_endpoint_testing job went red on main when us-central1 ran out of shared gemini-2.5-flash-lite capacity for a few hours. The suite's proxy config now sends the Vertex deployment to the global endpoint and retries rate limit errors 5 times with exponential backoff, and a regression test pins that the config rides out 3 consecutive 429s --- .../google_genai_proxy_test_config.yaml | 5 ++ .../test_google_genai_proxy_test_config.py | 67 +++++++++++++++++++ 2 files changed, 72 insertions(+) create mode 100644 tests/unified_google_tests/test_google_genai_proxy_test_config.py diff --git a/tests/unified_google_tests/google_genai_proxy_test_config.yaml b/tests/unified_google_tests/google_genai_proxy_test_config.yaml index 9913c05d434..64a83ef3d81 100644 --- a/tests/unified_google_tests/google_genai_proxy_test_config.yaml +++ b/tests/unified_google_tests/google_genai_proxy_test_config.yaml @@ -7,6 +7,11 @@ model_list: - model_name: vertex-gemini-2.5-flash-lite litellm_params: model: vertex_ai/gemini-2.5-flash-lite + vertex_location: global + +router_settings: + retry_policy: + RateLimitErrorRetries: 5 general_settings: master_key: sk-1234 diff --git a/tests/unified_google_tests/test_google_genai_proxy_test_config.py b/tests/unified_google_tests/test_google_genai_proxy_test_config.py new file mode 100644 index 00000000000..d84eefb406b --- /dev/null +++ b/tests/unified_google_tests/test_google_genai_proxy_test_config.py @@ -0,0 +1,67 @@ +import time +from pathlib import Path +from typing import Final, ReadOnly, TypedDict + +import httpx +import pytest +import respx +import yaml +from pydantic import TypeAdapter + +import litellm +from litellm import Router + +CONFIG_PATH: Final = Path(__file__).parent / "google_genai_proxy_test_config.yaml" +GEMINI_HOST: Final = "generativelanguage.googleapis.com" +GEMINI_GENERATE_CONTENT_PATH: Final = "/v1beta/models/gemini-2.5-flash-lite:generateContent" +RESOURCE_EXHAUSTED: Final = { + "error": {"code": 429, "message": "Resource exhausted. Please try again later.", "status": "RESOURCE_EXHAUSTED"} +} +PONG: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "pong"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 1, "totalTokenCount": 9}, +} +CONSECUTIVE_RATE_LIMITS: Final = 3 +MINIMUM_BACKOFF_SECONDS: Final = 0.5 + 1.0 + 2.0 + + +class _Deployment(TypedDict): + model_name: ReadOnly[str] + litellm_params: ReadOnly[dict[str, str]] + + +class _ProxyConfig(TypedDict): + model_list: ReadOnly[list[_Deployment]] + router_settings: ReadOnly[dict[str, dict[str, int]]] + + +def _router_from_ci_proxy_config() -> Router: + config: Final = TypeAdapter(_ProxyConfig).validate_python(yaml.safe_load(CONFIG_PATH.read_text())) + gemini_deployments: Final = [ + {"model_name": deployment["model_name"], "litellm_params": {**deployment["litellm_params"], "api_key": "test"}} + for deployment in config["model_list"] + if deployment["model_name"] == "gemini-2.5-flash-lite" + ] + return Router(model_list=gemini_deployments, retry_policy=config["router_settings"]["retry_policy"]) + + +@pytest.mark.asyncio +async def test_ci_proxy_config_rides_out_consecutive_429s_with_backoff( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + route: Final = respx_mock.post(host=GEMINI_HOST, path=GEMINI_GENERATE_CONTENT_PATH).mock( + side_effect=[httpx.Response(429, json=RESOURCE_EXHAUSTED)] * CONSECUTIVE_RATE_LIMITS + + [httpx.Response(200, json=PONG)] + ) + started: Final = time.monotonic() + response: Final = await _router_from_ci_proxy_config().agenerate_content( + model="gemini-2.5-flash-lite", + contents=[{"role": "user", "parts": [{"text": "Reply with only the single word: pong"}]}], + ) + elapsed: Final = time.monotonic() - started + + assert response.model_dump()["candidates"][0]["content"]["parts"][0]["text"] == "pong" + assert route.call_count == CONSECUTIVE_RATE_LIMITS + 1 + assert elapsed >= MINIMUM_BACKOFF_SECONDS From a7870a902a281e61a4842dbc4bd079fd621a21cf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 12:15:34 -0700 Subject: [PATCH 10/16] test(unified_google_tests): import ReadOnly from typing_extensions and cover the Vertex global endpoint The first commit imported ReadOnly from typing, which only exists on Python 3.13 and up. CircleCI runs this suite on 3.12, so the module failed at import and the job stopped at collection before any of its tests ran. ReadOnly and TypedDict now come from typing_extensions, like the rest of the repo A new test resolves the Vertex deployment's location from the suite's config with VERTEXAI_LOCATION set to a region, and fails if the vertex_location line is removed The expected minimum backoff is now derived from litellm's INITIAL_RETRY_DELAY and MAX_RETRY_DELAY, so the test holds when those are overridden through the environment --- .../test_google_genai_proxy_test_config.py | 50 +++++++++++++++---- 1 file changed, 40 insertions(+), 10 deletions(-) diff --git a/tests/unified_google_tests/test_google_genai_proxy_test_config.py b/tests/unified_google_tests/test_google_genai_proxy_test_config.py index d84eefb406b..694ec336bac 100644 --- a/tests/unified_google_tests/test_google_genai_proxy_test_config.py +++ b/tests/unified_google_tests/test_google_genai_proxy_test_config.py @@ -1,19 +1,26 @@ import time from pathlib import Path -from typing import Final, ReadOnly, TypedDict +from typing import Final import httpx import pytest import respx import yaml from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm import Router +from litellm.constants import INITIAL_RETRY_DELAY, MAX_RETRY_DELAY +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase CONFIG_PATH: Final = Path(__file__).parent / "google_genai_proxy_test_config.yaml" +GEMINI_DEPLOYMENT: Final = "gemini-2.5-flash-lite" +VERTEX_DEPLOYMENT: Final = "vertex-gemini-2.5-flash-lite" GEMINI_HOST: Final = "generativelanguage.googleapis.com" GEMINI_GENERATE_CONTENT_PATH: Final = "/v1beta/models/gemini-2.5-flash-lite:generateContent" +VERTEX_GLOBAL_BASE_URL: Final = "https://aiplatform.googleapis.com" RESOURCE_EXHAUSTED: Final = { "error": {"code": 429, "message": "Resource exhausted. Please try again later.", "status": "RESOURCE_EXHAUSTED"} } @@ -22,7 +29,9 @@ PONG: Final = { "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 1, "totalTokenCount": 9}, } CONSECUTIVE_RATE_LIMITS: Final = 3 -MINIMUM_BACKOFF_SECONDS: Final = 0.5 + 1.0 + 2.0 +MINIMUM_BACKOFF_SECONDS: Final = sum( + min(INITIAL_RETRY_DELAY * 2**attempt, MAX_RETRY_DELAY) for attempt in range(CONSECUTIVE_RATE_LIMITS) +) class _Deployment(TypedDict): @@ -35,14 +44,35 @@ class _ProxyConfig(TypedDict): router_settings: ReadOnly[dict[str, dict[str, int]]] +def _ci_proxy_config() -> _ProxyConfig: + return TypeAdapter(_ProxyConfig).validate_python(yaml.safe_load(CONFIG_PATH.read_text())) + + +def _litellm_params(config: _ProxyConfig, model_name: str) -> dict[str, str]: + return next( + deployment["litellm_params"] for deployment in config["model_list"] if deployment["model_name"] == model_name + ) + + def _router_from_ci_proxy_config() -> Router: - config: Final = TypeAdapter(_ProxyConfig).validate_python(yaml.safe_load(CONFIG_PATH.read_text())) - gemini_deployments: Final = [ - {"model_name": deployment["model_name"], "litellm_params": {**deployment["litellm_params"], "api_key": "test"}} - for deployment in config["model_list"] - if deployment["model_name"] == "gemini-2.5-flash-lite" - ] - return Router(model_list=gemini_deployments, retry_policy=config["router_settings"]["retry_policy"]) + config: Final = _ci_proxy_config() + return Router( + model_list=[ + { + "model_name": GEMINI_DEPLOYMENT, + "litellm_params": {**_litellm_params(config, GEMINI_DEPLOYMENT), "api_key": "test"}, + } + ], + retry_policy=config["router_settings"]["retry_policy"], + ) + + +def test_ci_proxy_config_sends_vertex_calls_to_the_global_endpoint(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5") + location: Final = VertexBase.safe_get_vertex_ai_location(_litellm_params(_ci_proxy_config(), VERTEX_DEPLOYMENT)) + + assert location == "global" + assert get_vertex_base_url(location) == VERTEX_GLOBAL_BASE_URL @pytest.mark.asyncio @@ -57,7 +87,7 @@ async def test_ci_proxy_config_rides_out_consecutive_429s_with_backoff( ) started: Final = time.monotonic() response: Final = await _router_from_ci_proxy_config().agenerate_content( - model="gemini-2.5-flash-lite", + model=GEMINI_DEPLOYMENT, contents=[{"role": "user", "parts": [{"text": "Reply with only the single word: pong"}]}], ) elapsed: Final = time.monotonic() - started From 7d93821e415bca477f3041b6f20b878d68de227f Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 19:30:57 +0000 Subject: [PATCH 11/16] fix(otel v2): keep Responses refusal text on the folded assistant message Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/model/payloads.py | 19 +++++++++++-------- .../otel/test_otel_v2_sources_of_truth.py | 19 ++++++++++++++++++- .../otel/test_otel_v2_vendor_mappers.py | 1 + 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 484f4a4c294..c23b3291365 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -720,6 +720,7 @@ class _ToolCall(TypedDict): class _AssistantMessage(TypedDict): role: ReadOnly[str] content: ReadOnly[str | None] + refusal: ReadOnly[str | None] tool_calls: ReadOnly[tuple[_ToolCall, ...] | None] @@ -735,13 +736,7 @@ def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: """A Responses API ``output`` folded into one chat-shaped assistant choice.""" items: Final = _dicts(response.get("output")) messages: Final = tuple(item for item in items if item.get("type") == "message") - content: Final = "".join( - text - for item in messages - for part in _dicts(item.get("content")) - if part.get("type") == "output_text" - if (text := as_str(part.get("text"))) is not None - ) + parts: Final = tuple(part for item in messages for part in _dicts(item.get("content"))) tool_calls: Final = tuple( _responses_tool_call(item) for item in items if item.get("type") in _RESPONSES_TOOL_CALL_TYPES ) @@ -749,13 +744,21 @@ def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: return () message: Final[_AssistantMessage] = { "role": next((role for item in messages if (role := as_str(item.get("role")))), "assistant"), - "content": content if messages else None, + "content": _responses_parts_text(parts, "output_text", "text"), + "refusal": _responses_parts_text(parts, "refusal", "refusal"), "tool_calls": tool_calls or None, } choice: Final[_Choice] = {"message": message, "finish_reason": _responses_finish_reason(response, bool(tool_calls))} return (choice,) +def _responses_parts_text(parts: tuple[Mapping[str, object], ...], part_type: str, field: str) -> str | None: + texts: Final = tuple( + text for part in parts if part.get("type") == part_type if (text := as_str(part.get(field))) is not None + ) + return "".join(texts) if texts else None + + def _responses_tool_call(item: Mapping[str, object]) -> _ToolCall: custom: Final = item.get("type") == "custom_tool_call" function: Final[_ToolFunction] = { diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 972c91670f8..17de3cf1e8a 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -761,7 +761,7 @@ def test_responses_output_text_becomes_one_assistant_choice_with_stop(): assert json.loads(json.dumps(data.choices_out)) == [ { - "message": {"role": "assistant", "content": "pong", "tool_calls": None}, + "message": {"role": "assistant", "content": "pong", "refusal": None, "tool_calls": None}, "finish_reason": "stop", } ] @@ -835,6 +835,23 @@ def test_responses_content_only_reads_output_text_parts(): data = LLMCallSpanData.from_standard_logging_payload(_responses_payload([item]), capture_content=True) assert data.choices_out[0]["message"]["content"] == "ok" + assert data.choices_out[0]["message"]["refusal"] == "no" + + +def test_responses_refusal_only_output_keeps_the_refusal_text(): + item = { + "type": "message", + "role": "assistant", + "content": [{"type": "refusal", "refusal": "I can't "}, {"type": "refusal", "refusal": "help with that."}], + } + data = LLMCallSpanData.from_standard_logging_payload(_responses_payload([item]), capture_content=True) + + assert json.loads(json.dumps(data.choices_out)) == [ + { + "message": {"role": "assistant", "content": None, "refusal": "I can't help with that.", "tool_calls": None}, + "finish_reason": "stop", + } + ] def test_responses_output_without_messages_or_tool_calls_stays_empty(): diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py index 5b4d1e7a802..4e375de0494 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py @@ -218,6 +218,7 @@ def test_langfuse_mapper_renders_a_responses_api_call_from_the_standard_logging_ { "role": "assistant", "content": "Checking.", + "refusal": None, "tool_calls": [ {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "sf"}'}} ], From e8f2ee82002683b1e7f37c6d24f4281676145e6e Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 19:35:08 +0000 Subject: [PATCH 12/16] fix(redaction): redact Responses refusal parts under turn_off_message_logging Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/redact_messages.py | 4 +++ .../test_redact_messages.py | 29 +++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 1f9464a2a26..b409b181a79 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -128,6 +128,8 @@ def _redact_responses_api_output(output_items): for content_part in output_item.content: if getattr(content_part, "text", None) is not None: content_part.text = REDACTED_BY_LITELLM + if getattr(content_part, "refusal", None) is not None: + content_part.refusal = REDACTED_BY_LITELLM # Redact reasoning items in output array if hasattr(output_item, "type") and output_item.type == "reasoning": @@ -155,6 +157,8 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): for content_item in output_item["content"]: if isinstance(content_item, dict) and content_item.get("text") is not None: content_item["text"] = redacted_str + if isinstance(content_item, dict) and content_item.get("refusal") is not None: + content_item["refusal"] = redacted_str if output_item.get("type") == "reasoning" and isinstance(output_item.get("summary"), list): for summary_item in output_item["summary"]: diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index c6c9a9dd2b7..276a67e0bd4 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -507,6 +507,26 @@ class TestPerformRedaction: assert redacted["output"][0]["name"] == "grep" assert redacted["output"][1]["input"] == "not-a-custom-input" + def test_redacts_responses_api_refusal_parts_dict(self): + result = { + "output": [ + { + "type": "message", + "role": "assistant", + "content": [ + {"type": "refusal", "refusal": "I cannot share the secret"}, + {"type": "output_text", "text": "ok"}, + ], + } + ] + } + + redacted = perform_redaction({}, result) + + assert redacted["output"][0]["content"][0]["refusal"] == "redacted-by-litellm" + assert redacted["output"][0]["content"][0]["type"] == "refusal" + assert redacted["output"][0]["content"][1]["text"] == "redacted-by-litellm" + def test_redacts_every_tool_call_in_multi_element_list(self): result = litellm.ModelResponse( id="resp-multi", @@ -585,6 +605,15 @@ class TestPerformRedaction: assert output_item.input == "redacted-by-litellm" assert output_item.name == "grep" + def test_redacts_responses_api_refusal_parts_object(self): + refusal = SimpleNamespace(type="refusal", refusal="I cannot share the secret") + output_item = SimpleNamespace(type="message", role="assistant", content=[refusal]) + + _redact_responses_api_output([output_item]) + + assert refusal.refusal == "redacted-by-litellm" + assert refusal.type == "refusal" + def test_redacts_response_output_objects_with_top_level_text(self): output_items = [ SimpleNamespace(text="top-level output"), From bf9c717d77105b206edbcc5aa43e5c07adcf81ac Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 19 Sep 2026 12:50:32 -0700 Subject: [PATCH 13/16] test(e2e): stop the config suite locking itself out of the shared proxy Two tests in the config/misc management suite were failing every run against the Buildkite e2e stack, and one of them took the rest of the build with it. test_add_allowed_ip_does_not_store_unrelated_config_value posted 127.0.0.1 to /add/allowed_ip. That route sets the live general_settings["allowed_ips"] that auth_utils._check_valid_ip reads before it persists anything, and the check is exact string membership with no CIDR support, so from the moment the POST returns only 127.0.0.1 can reach the proxy. The runner 403s on its very next call, and the deferred /delete/allowed_ip sits behind the same auth dependency, so the cleanup is locked out too and every later test in the build 403s. Build 254's first attempt lost 459 of its 465 failures to that one cascade. There is no safe way to exercise the route against a shared proxy: nothing reports the caller's address as the proxy sees it, so a test cannot allowlist itself first. Move the claim to the route's own TestClient suite, where the auth dependency is overridden and general_settings is per-test, and record the route in the module docstring beside /cache/settings and the Vault override so it is not re-added. save_config's end of the contract was already covered by test_ProxyConfig_save_config_merges_changed_keys_without_copying_file_settings; the new test covers the route's end, that what it hands save_config differs from the loaded config in allowed_ips and nothing else. The unrelated-key probe also only ever worked on one lane: max_parallel_requests was added to tests/e2e/gateway/stage_mirror_ci_config.yml and never to the Buildkite stack's config, where resolve() reports it as "unset" rather than "config". That key is now unused, so drop it again. test_config_update_persists_router_setting_to_get wrote router_settings. num_retries, which both lanes declare in their config file, so the config- ownership work correctly refuses it with a 400. Switch to retry_after, which is declared by neither lane, is accepted by /config/update, and is reported back by GET /router/settings. Verified against a live proxy: max_fallbacks also takes the write but never reads back, so the read-back poll is what picks the key. --- tests/e2e/coverage_registry/mgmt.yaml | 1 - tests/e2e/gateway/stage_mirror_ci_config.yml | 1 - .../test_config_misc_endpoints_e2e.py | 149 +++++------------- .../test_proxy_setting_endpoints.py | 63 ++++++++ 4 files changed, 102 insertions(+), 112 deletions(-) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 9890902fa5e..85fbd0acd91 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -72,7 +72,6 @@ - {id: mgmt.callback.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "callback_management_endpoints.py", rationale: "Callback config (smoke)"} - {id: mgmt.cost_tracking.estimate.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cost_tracking_settings.py", rationale: "Cost estimate (smoke)"} - {id: mgmt.router_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "router_settings_endpoints.py", rationale: "Router config (smoke)"} -- {id: mgmt.config.allowed_ip.changed_key_only, module: mgmt, tier: P2, surface: api, assertions: [persists], source: "proxy_setting_endpoints.py:496", rationale: "An allowed-IP change leaves unrelated file settings out of the DB row"} - {id: mgmt.jwt_key_mapping.new.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "jwt_key_mapping_endpoints.py", rationale: "JWT->key mapping (smoke)"} - {id: mgmt.compliance.gdpr.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "compliance_endpoints.py", rationale: "GDPR ops (smoke)"} - {id: mgmt.tool_management.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "tool_management_endpoints.py", rationale: "Tool inventory (smoke)"} diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml index 1b6ae93f461..8c8e64443cb 100644 --- a/tests/e2e/gateway/stage_mirror_ci_config.yml +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -1,5 +1,4 @@ general_settings: - max_parallel_requests: 100 proxy_batch_write_at: 5 enable_jwt_auth: true litellm_jwtauth: diff --git a/tests/e2e/management/test_config_misc_endpoints_e2e.py b/tests/e2e/management/test_config_misc_endpoints_e2e.py index 20e98e993d4..a3be0a64e7f 100644 --- a/tests/e2e/management/test_config_misc_endpoints_e2e.py +++ b/tests/e2e/management/test_config_misc_endpoints_e2e.py @@ -7,13 +7,18 @@ so a read-back reflects the change. Router settings, which mutate global proxy state, are exercised with a benign, self-restoring change so a shared proxy is left as it was found. -Cache settings and the Vault config override are deliberately not covered here. -Both routes reconfigure the whole proxy: /cache/settings persists what it receives -into a row that outranks the YAML cache_params and is re-applied on a timer, and -/config_overrides/hashicorp_vault swaps the process-wide secret manager. Neither can -be exercised safely against the shared proxy the suites run on, so they need an -isolated proxy before a test lands. Do not add a read-then-write-back test for -either one. +Cache settings, the Vault config override and the allowed-IP routes are deliberately +not covered here. All three reconfigure the whole proxy: /cache/settings persists what +it receives into a row that outranks the YAML cache_params and is re-applied on a timer, +/config_overrides/hashicorp_vault swaps the process-wide secret manager, and +/add/allowed_ip mutates the live general_settings["allowed_ips"] that +auth_utils._check_valid_ip reads, so the first call locks every other client out of the +shared proxy. The allowlist is an exact string match with no CIDR support, and no route +reports the caller's address as the proxy sees it, so a test cannot allowlist itself +first; /delete/allowed_ip sits behind the same auth dependency, so the cleanup is locked +out too and the proxy stays poisoned for the rest of the build. None of the three can be +exercised safely against the shared proxy the suites run on, so they need an isolated +proxy before a test lands. Do not add a read-then-write-back test for any of them. """ from __future__ import annotations @@ -21,10 +26,9 @@ from __future__ import annotations import math import time from collections.abc import Callable -from typing import Final import pytest -from pydantic import BaseModel, JsonValue, RootModel +from pydantic import BaseModel from e2e_config import unique_marker from e2e_http import NoBody, Success, unwrap, unwrap_status @@ -188,7 +192,7 @@ class JwtKeyMappingResponse(BaseModel): class RouterSettingsPatch(BaseModel): - num_retries: int + retry_after: int class ConfigUpdateBody(BaseModel): @@ -199,39 +203,8 @@ class ConfigUpdateResponse(BaseModel): message: str -class AllowedIpBody(BaseModel): - ip: str - - -class ConfigFieldInfoParams(BaseModel): - field_name: str - - -class ConfigFieldInfoResponse(BaseModel): - field_name: str - field_value: JsonValue - source: str - editable: bool - - -class ConfigListParams(BaseModel): - config_type: str - - -class ConfigListEntry(BaseModel): - field_name: str - field_value: JsonValue - stored_in_db: bool | None - source: str - editable: bool - - -class ConfigListResponse(RootModel[list[ConfigListEntry]]): - pass - - class RouterCurrentValues(BaseModel): - num_retries: int | None = None + retry_after: int | None = None class RouterSettingsResponse(BaseModel): @@ -493,17 +466,25 @@ class TestRouterSettings: ) -> None: """/config/update is the only write path for router_settings (there is no dedicated router-settings write route). The change is restored on teardown so - the shared proxy keeps its original retry policy.""" - original = self._read_num_retries(client) - assert original is not None, "GET /router/settings did not report num_retries; cannot prove a change" - resources.defer(lambda: self._write_num_retries(client, original)) + the shared proxy keeps its original retry policy. - target = original + 5 + retry_after is the subject because it satisfies all three constraints at once: + no lane's config file declares it, so the database owns it and the write is not + refused as config-owned; it is in RUNTIME_UPDATABLE_ROUTER_SETTINGS, so + /config/update accepts it; and it is in ROUTER_SETTINGS_FIELDS backed by an + always-set Router attribute, so GET /router/settings reports it for the + read-back. Bumping it by one second is the smallest change that proves the + round-trip without slowing a concurrent test that hits a retry.""" + original = self._read_retry_after(client) + assert original is not None, "GET /router/settings did not report retry_after; cannot prove a change" + resources.defer(lambda: self._write_retry_after(client, original)) + + target = original + 1 response = unwrap( client.proxy.transport.post( "/config/update", headers=client.proxy.transport.master, - json=ConfigUpdateBody(router_settings=RouterSettingsPatch(num_retries=target)), + json=ConfigUpdateBody(router_settings=RouterSettingsPatch(retry_after=target)), response_type=ConfigUpdateResponse, ) ) @@ -513,20 +494,20 @@ class TestRouterSettings: _ = _poll( client, - lambda: True if self._read_num_retries(client) == target else None, - f"GET /router/settings never reported num_retries {target} after /config/update", + lambda: True if self._read_retry_after(client) == target else None, + f"GET /router/settings never reported retry_after {target} after /config/update", ) - self._write_num_retries(client, original) + self._write_retry_after(client, original) restored = _poll( client, - lambda: original if self._read_num_retries(client) == original else None, - f"GET /router/settings never returned to the original num_retries {original} after the restore", + lambda: original if self._read_retry_after(client) == original else None, + f"GET /router/settings never returned to the original retry_after {original} after the restore", ) - assert restored == original, f"router num_retries left at {restored}, expected the original {original}" + assert restored == original, f"router retry_after left at {restored}, expected the original {original}" @staticmethod - def _read_num_retries(client: ManagementClient) -> int | None: + def _read_retry_after(client: ManagementClient) -> int | None: return unwrap( client.proxy.transport.get( "/router/settings", @@ -534,72 +515,20 @@ class TestRouterSettings: params=NoBody(), response_type=RouterSettingsResponse, ) - ).current_values.num_retries + ).current_values.retry_after @staticmethod - def _write_num_retries(client: ManagementClient, value: int) -> None: + def _write_retry_after(client: ManagementClient, value: int) -> None: _ = unwrap( client.proxy.transport.post( "/config/update", headers=client.proxy.transport.master, - json=ConfigUpdateBody(router_settings=RouterSettingsPatch(num_retries=value)), + json=ConfigUpdateBody(router_settings=RouterSettingsPatch(retry_after=value)), response_type=ConfigUpdateResponse, ) ) -class TestConfigPersistence: - @pytest.mark.covers("mgmt.config.allowed_ip.changed_key_only") - def test_add_allowed_ip_does_not_store_unrelated_config_value( - self, client: ManagementClient, resources: ResourceManager - ) -> None: - allowed_ip: Final = "127.0.0.1" - added: Final = unwrap( - client.proxy.transport.post( - "/add/allowed_ip", - headers=client.proxy.transport.master, - json=AllowedIpBody(ip=allowed_ip), - response_type=ConfigUpdateResponse, - ) - ) - resources.defer( - lambda: unwrap( - client.proxy.transport.post( - "/delete/allowed_ip", - headers=client.proxy.transport.master, - json=AllowedIpBody(ip=allowed_ip), - response_type=ConfigUpdateResponse, - ) - ) - ) - assert added.message == f"IP {allowed_ip} address added successfully" - - listed: Final = unwrap( - client.proxy.transport.get( - "/config/list", - headers=client.proxy.transport.master, - params=ConfigListParams(config_type="general_settings"), - response_type=ConfigListResponse, - ) - ) - unrelated: Final = next(entry for entry in listed.root if entry.field_name == "max_parallel_requests") - assert unrelated.stored_in_db is not True - assert unrelated.source == "config" - assert unrelated.editable is False - - field_info: Final = unwrap( - client.proxy.transport.get( - "/config/field/info", - headers=client.proxy.transport.master, - params=ConfigFieldInfoParams(field_name="max_parallel_requests"), - response_type=ConfigFieldInfoResponse, - ) - ) - assert field_info.source == "config" - assert field_info.editable is False - assert field_info.field_value == unrelated.field_value - - class TestMcpServerSubmission: @pytest.mark.covers("mgmt.mcp_server.register.happy_path") def test_register_submits_pending_server(self, client: ManagementClient, resources: ResourceManager) -> None: diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 58201bd14ce..b01544e54e3 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -2604,6 +2604,69 @@ def test_add_allowed_ip_writes_audit_log(mock_proxy_config, monkeypatch): app.dependency_overrides.pop(user_api_key_auth, None) +def test_add_allowed_ip_hands_save_config_only_the_changed_general_setting(monkeypatch): + """An allowed-IP write must not drag the config file's own general_settings into + the database row. This covers the route end of that contract: what /add/allowed_ip + hands save_config differs from the loaded config in allowed_ips and nothing else. + save_config's end -- that the row it writes holds only those changed keys -- is + covered by test_ProxyConfig_save_config_merges_changed_keys_without_copying_file_settings. + + This lives here rather than in the e2e suite because /add/allowed_ip mutates the + live general_settings["allowed_ips"] that auth_utils._check_valid_ip reads, so on a + shared proxy the first call locks every later request out, cleanup included. + """ + from copy import deepcopy + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.config_resolvers.changed_section_keys import changed_section_keys + from litellm.proxy.config_resolvers.settings_store import SettingsStore + + file_settings = {"max_parallel_requests": 100, "proxy_config_reload_interval_seconds": 7} + store = SettingsStore("general_settings") + store.load_yaml(file_settings) + saved = [] + + fake_prisma = MagicMock() + fake_prisma.db.litellm_auditlog.create = AsyncMock() + + async def _get_config(): + return {"general_settings": deepcopy(file_settings)} + + async def _save_config(new_config=None): + saved.append(new_config) + return new_config + + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server_module, "store_model_in_db", True) + monkeypatch.setattr(proxy_server_module, "premium_user", True) + monkeypatch.setattr(proxy_server_module, "general_settings", store) + monkeypatch.setattr(proxy_server_module.proxy_config, "get_config", _get_config) + monkeypatch.setattr(proxy_server_module.proxy_config, "save_config", _save_config) + + async def _admin_auth(): + return UserAPIKeyAuth( + user_id="config-admin", + api_key="hashed-admin-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + app.dependency_overrides[user_api_key_auth] = _admin_auth + try: + resp = client.post("/add/allowed_ip", json={"ip": "203.0.113.77"}) + assert resp.status_code == 200, resp.text + + assert len(saved) == 1, f"expected exactly one save_config call, got {len(saved)}" + changed, removed = changed_section_keys(file_settings, saved[0]["general_settings"]) + assert dict(changed) == {"allowed_ips": ["203.0.113.77"]} + assert removed == frozenset() + assert store["allowed_ips"] == ["203.0.113.77"] + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + def test_delete_allowed_ip_writes_deleted_audit_log(monkeypatch): """Removing an allowed IP must be audited as a deletion, symmetric with the add path.""" From 3c6a2f258a8017425fc0d53587c9b97daf3133c0 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 19 Sep 2026 12:56:51 -0700 Subject: [PATCH 14/16] test(proxy): capture the saved config with an AsyncMock instead of a mutable list Greptile flagged the unannotated list and append against the repository's immutable-state and Final-local rules (LIT001/LIT010). Recording the call on an AsyncMock removes the accumulator entirely and matches how the neighbouring audit-log tests in this file read their captured arguments. --- .../test_proxy_setting_endpoints.py | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index b01544e54e3..47bb1ad5a81 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -2615,7 +2615,8 @@ def test_add_allowed_ip_hands_save_config_only_the_changed_general_setting(monke live general_settings["allowed_ips"] that auth_utils._check_valid_ip reads, so on a shared proxy the first call locks every later request out, cleanup included. """ - from copy import deepcopy + from types import MappingProxyType + from typing import Final from unittest.mock import AsyncMock, MagicMock import litellm.proxy.proxy_server as proxy_server_module @@ -2624,27 +2625,23 @@ def test_add_allowed_ip_hands_save_config_only_the_changed_general_setting(monke from litellm.proxy.config_resolvers.changed_section_keys import changed_section_keys from litellm.proxy.config_resolvers.settings_store import SettingsStore - file_settings = {"max_parallel_requests": 100, "proxy_config_reload_interval_seconds": 7} - store = SettingsStore("general_settings") + file_settings: Final = MappingProxyType({"max_parallel_requests": 100, "proxy_config_reload_interval_seconds": 7}) + store: Final = SettingsStore("general_settings") store.load_yaml(file_settings) - saved = [] - fake_prisma = MagicMock() + fake_prisma: Final = MagicMock() fake_prisma.db.litellm_auditlog.create = AsyncMock() + save_config: Final = AsyncMock(side_effect=lambda new_config: new_config) async def _get_config(): - return {"general_settings": deepcopy(file_settings)} - - async def _save_config(new_config=None): - saved.append(new_config) - return new_config + return {"general_settings": dict(file_settings)} monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) monkeypatch.setattr(proxy_server_module, "store_model_in_db", True) monkeypatch.setattr(proxy_server_module, "premium_user", True) monkeypatch.setattr(proxy_server_module, "general_settings", store) monkeypatch.setattr(proxy_server_module.proxy_config, "get_config", _get_config) - monkeypatch.setattr(proxy_server_module.proxy_config, "save_config", _save_config) + monkeypatch.setattr(proxy_server_module.proxy_config, "save_config", save_config) async def _admin_auth(): return UserAPIKeyAuth( @@ -2655,11 +2652,12 @@ def test_add_allowed_ip_hands_save_config_only_the_changed_general_setting(monke app.dependency_overrides[user_api_key_auth] = _admin_auth try: - resp = client.post("/add/allowed_ip", json={"ip": "203.0.113.77"}) + resp: Final = client.post("/add/allowed_ip", json={"ip": "203.0.113.77"}) assert resp.status_code == 200, resp.text - assert len(saved) == 1, f"expected exactly one save_config call, got {len(saved)}" - changed, removed = changed_section_keys(file_settings, saved[0]["general_settings"]) + save_config.assert_awaited_once() + persisted: Final = save_config.await_args.kwargs["new_config"]["general_settings"] + changed, removed = changed_section_keys(file_settings, persisted) assert dict(changed) == {"allowed_ips": ["203.0.113.77"]} assert removed == frozenset() assert store["allowed_ips"] == ["203.0.113.77"] From c02399b29dbc6b3a243679c888302caa47245d73 Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Sat, 19 Sep 2026 12:23:59 -0700 Subject: [PATCH 15/16] fix(terraform): unlink the registry docs entries that 404 on click The resource and data source links on the provider's registry docs overview page 404 when clicked. They are written as relative paths like ./resources/team, and the registry serves the overview at .../latest/docs with no trailing slash and passes hrefs through unrewritten, so the browser resolves them to .../latest/resources/team. Drops the link markup and keeps both lists and their descriptions. No relative form works in both places: only a docs/-prefixed target resolves correctly on the registry, and that same path is wrong when reading the file on GitHub. The registry sidebar already links every resource and data source for the version being read. Co-Authored-By: Claude Opus 5 --- terraform/provider/docs/index.md | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/terraform/provider/docs/index.md b/terraform/provider/docs/index.md index e6641782a4d..c446567549d 100644 --- a/terraform/provider/docs/index.md +++ b/terraform/provider/docs/index.md @@ -43,22 +43,22 @@ resource "litellm_team" "dev_team" { The LiteLLM provider supports the following resources: -* [`litellm_model`](./resources/model) - Manage LiteLLM model configurations -* [`litellm_team`](./resources/team) - Manage teams and their permissions -* [`litellm_team_member`](./resources/team_member) - Manage team member configurations -* [`litellm_team_member_add`](./resources/team_member_add) - Add members to teams -* [`litellm_key`](./resources/key) - Manage API keys -* [`litellm_mcp_server`](./resources/mcp_server) - Manage MCP (Model Context Protocol) servers -* [`litellm_credential`](./resources/credential) - Manage credentials for various providers -* [`litellm_vector_store`](./resources/vector_store) - Manage vector stores -* [`litellm_jwt_key_mapping`](./resources/jwt_key_mapping) - Map JWT claim values to virtual keys +* `litellm_model` - Manage LiteLLM model configurations +* `litellm_team` - Manage teams and their permissions +* `litellm_team_member` - Manage team member configurations +* `litellm_team_member_add` - Add members to teams +* `litellm_key` - Manage API keys +* `litellm_mcp_server` - Manage MCP (Model Context Protocol) servers +* `litellm_credential` - Manage credentials for various providers +* `litellm_vector_store` - Manage vector stores +* `litellm_jwt_key_mapping` - Map JWT claim values to virtual keys ## Available Data Sources The LiteLLM provider supports the following data sources: -* [`litellm_credential`](./data-sources/credential) - Retrieve credential information -* [`litellm_vector_store`](./data-sources/vector_store) - Retrieve vector store information +* `litellm_credential` - Retrieve credential information +* `litellm_vector_store` - Retrieve vector store information ## Authentication From 8767f1279489ddbae97108b4d00d318efb57f3f0 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 19 Sep 2026 14:27:21 -0700 Subject: [PATCH 16/16] bump: litellm-enterprise 0.1.68 -> 0.1.69, litellm-proxy-extras 0.4.99 -> 0.4.100 --- enterprise/pyproject.toml | 4 ++-- litellm-proxy-extras/pyproject.toml | 4 ++-- pyproject.toml | 4 ++-- uv.lock | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 06b1da7ea76..729f3264706 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.68" +version = "0.1.69" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.68" +version = "0.1.69" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 604ffc3abd4..fb9022f89a5 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.99" +version = "0.4.100" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.99" +version = "0.4.100" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/pyproject.toml b/pyproject.toml index f2ee1d92d7f..1295feabb43 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,8 +71,8 @@ proxy = [ "mcp>=2.2.0,<3", "httpx2>=2.5.0,<3", "pydantic>=2.12.0,<3", - "litellm-proxy-extras==0.4.99", - "litellm-enterprise==0.1.68", + "litellm-proxy-extras==0.4.100", + "litellm-enterprise==0.1.69", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", diff --git a/uv.lock b/uv.lock index db2fb11c6e3..f1a58500a61 100644 --- a/uv.lock +++ b/uv.lock @@ -4942,12 +4942,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.68" +version = "0.1.69" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.99" +version = "0.4.100" source = { editable = "litellm-proxy-extras" } [[package]]