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/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..c4842dd9223 --- /dev/null +++ b/litellm-rust/crates/host-python/src/fork_gate.rs @@ -0,0 +1,139 @@ +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(()) + } + + /// 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 { + let _ = + self.fork_only_pid + .compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst); + 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 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(); + 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..7d164ab7535 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,12 @@ 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..46f98736aa1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -13,7 +13,7 @@ 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] @@ -30,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::*; @@ -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/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 467c286db9d..c23b3291365 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,84 @@ 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] + refusal: 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") + 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 + ) + 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": _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] = { + "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..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": @@ -138,6 +140,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): @@ -153,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"]: @@ -161,6 +167,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/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/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/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/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 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/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/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index f4a8691f72f..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 @@ -738,6 +738,140 @@ 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", "refusal": None, "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" + 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(): + 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..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 @@ -196,6 +196,37 @@ 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.", + "refusal": None, + "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..276a67e0bd4 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,40 @@ 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_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", @@ -563,6 +597,23 @@ 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_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"), 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/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 58201bd14ce..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 @@ -2604,6 +2604,67 @@ 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 types import MappingProxyType + from typing import Final + 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: Final = MappingProxyType({"max_parallel_requests": 100, "proxy_config_reload_interval_seconds": 7}) + store: Final = SettingsStore("general_settings") + store.load_yaml(file_settings) + + 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": 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) + + 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: Final = client.post("/add/allowed_ip", json={"ip": "203.0.113.77"}) + assert resp.status_code == 200, resp.text + + 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"] + 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.""" 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..88ae017ec39 --- /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: + assert _reserve_with(monkeypatch, None) is None + + +def test_extension_built_before_reservation_existed_passes(monkeypatch: pytest.MonkeyPatch) -> None: + assert _reserve_with(monkeypatch, SimpleNamespace()) is None + + +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/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/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()) 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 new file mode 100644 index 00000000000..086397bab5c --- /dev/null +++ b/tests/test_litellm_rust/test_fork_guard.py @@ -0,0 +1,146 @@ +import os +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( + """ + 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 = run_child_interpreter(_NATIVE_CONTRACT, env=env, timeout=60) + + 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 = run_child_interpreter(_SDK_CONTRACT, env=env, timeout=120) + + assert result.returncode == 0, result.stderr 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..694ec336bac --- /dev/null +++ b/tests/unified_google_tests/test_google_genai_proxy_test_config.py @@ -0,0 +1,97 @@ +import time +from pathlib import Path +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"} +} +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 = sum( + min(INITIAL_RETRY_DELAY * 2**attempt, MAX_RETRY_DELAY) for attempt in range(CONSECUTIVE_RATE_LIMITS) +) + + +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 _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 = _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 +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_DEPLOYMENT, + 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 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]]