From c79d1d12ae75a8eaec1de15fe4f0de01330cc553 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:15:57 -0700 Subject: [PATCH 01/49] fix(ai-gateway): install a rustls crypto provider before dialing upstream WebSockets The gateway's dependency graph turns on two rustls crypto backends at once: reqwest's rustls-tls pulls in ring, and litellm-core's bedrock-auth pulls in aws-lc-rs through aws-config. rustls 0.23 refuses to guess between them, so ClientConfig::builder panics, and that is exactly how tokio-tungstenite builds its TLS config. Every outbound WebSocket dial killed its tokio worker and the client saw the socket vanish with no close frame. reqwest and the AWS SDK both pick a provider explicitly, so only the tungstenite path was affected. Route all three dial sites through one helper that installs ring once per process before connecting. --- litellm-rust/Cargo.lock | 1 + litellm-rust/Cargo.toml | 1 + litellm-rust/crates/ai-gateway/Cargo.toml | 3 + litellm-rust/crates/ai-gateway/src/io/mod.rs | 1 + .../crates/ai-gateway/src/io/realtime.rs | 6 +- .../crates/ai-gateway/src/io/responses_ws.rs | 12 ++-- litellm-rust/crates/ai-gateway/src/io/tls.rs | 67 +++++++++++++++++++ 7 files changed, 84 insertions(+), 7 deletions(-) create mode 100644 litellm-rust/crates/ai-gateway/src/io/tls.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b3dac5ca935..2ed998174e4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1415,6 +1415,7 @@ dependencies = [ "litellm-core", "pyo3", "reqwest", + "rustls 0.23.42", "serde", "serde_json", "sha2 0.10.9", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a13dd4c04b0..15074082981 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -26,6 +26,7 @@ pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" +rustls = { version = "0.23", default-features = false, features = ["ring", "std"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" sha2 = "0.10" diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index e3dbdf24ce6..414abc2356d 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -19,6 +19,9 @@ litellm-core = { workspace = true, features = ["bedrock-auth"] } # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the # Python proxy callbacks API. reqwest.workspace = true +# rustls is a direct dependency so `io::tls` can install a process-level +# crypto provider; see that module for why the graph needs one. +rustls.workspace = true # `sync` powers the bounded mpsc channel the realtime logger drains. tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] } tokio-tungstenite.workspace = true diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs index cce56dd2121..7098d67993f 100644 --- a/litellm-rust/crates/ai-gateway/src/io/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -3,3 +3,4 @@ pub mod ocr; pub mod realtime; pub mod realtime_pool; pub mod responses_ws; +pub(crate) mod tls; diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 662f7328982..53d87848342 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -23,10 +23,12 @@ use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; +use crate::io::tls::connect_upstream; + /// Environment variable holding the OpenAI API key (last-resort fallback). const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; @@ -84,7 +86,7 @@ pub(crate) async fn dial_upstream( .map_err(|err| Error::Auth(err.to_string()))?, ); - let (upstream, _response) = connect_async(request) + let (upstream, _response) = connect_upstream(request) .await .map_err(|err| Error::Network(err.to_string()))?; Ok(upstream) diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index 0b01747b1a5..ee181791413 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -14,7 +14,9 @@ use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName}; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +use crate::io::tls::connect_upstream; use crate::constants::{ DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS, @@ -49,14 +51,14 @@ impl ResponsesWebSocketConnection { .map_err(|error| Error::InvalidRequest(error.to_string()))?; request.headers_mut().insert(header_name, header_value); } - let connect = connect_async(request); + let connect = connect_upstream(request); let result = match timeout { Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { Error::Network("Responses WebSocket connection timed out".to_string()) })?, None => connect.await, }; - let (socket, _) = result.map_err(|error| match error { + let (socket, _) = result.map_err(|error| match *error { tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), @@ -138,13 +140,13 @@ async fn dial_upstream( ); let result = tokio::time::timeout( Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS), - connect_async(request), + connect_upstream(request), ) .await .map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?; result .map(|(socket, _)| socket) - .map_err(|error| match error { + .map_err(|error| match *error { tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs new file mode 100644 index 00000000000..e3754e44002 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs @@ -0,0 +1,67 @@ +//! Outbound WebSocket dials, with the rustls crypto provider settled first. +//! +//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth` +//! enables `rustls/aws-lc-rs`, so `ClientConfig::builder()` — which is how +//! `tokio-tungstenite` builds its TLS config — panics rather than guess between +//! them. reqwest and the AWS SDK pick a provider explicitly and never panic. +//! +//! Installing from the dial rather than from a `main` also covers the `cdylib` +//! the Python bridge loads, the tests, and the benches, none of which have one. +//! ring is what reqwest already falls back to, so installing it changes no +//! working path, and an embedder that installed its own provider first keeps it. + +use std::sync::Once; + +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::Error; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::handshake::client::Response; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; + +static INSTALL_CRYPTO_PROVIDER: Once = Once::new(); + +pub(crate) fn ensure_crypto_provider() { + INSTALL_CRYPTO_PROVIDER.call_once(|| { + let _ = rustls::crypto::ring::default_provider().install_default(); + }); +} + +pub(crate) async fn connect_upstream( + request: R, +) -> Result<(WebSocketStream>, Response), Box> +where + R: IntoClientRequest + Unpin, +{ + ensure_crypto_provider(); + connect_async(request).await.map_err(Box::new) +} + +#[cfg(test)] +mod tests { + use super::ensure_crypto_provider; + + #[test] + fn client_config_builder_works_with_both_provider_features_enabled() { + ensure_crypto_provider(); + + assert!(rustls::crypto::CryptoProvider::get_default().is_some()); + + let config = rustls::ClientConfig::builder() + .with_root_certificates(rustls::RootCertStore::empty()) + .with_no_client_auth(); + + assert!(!config.crypto_provider().cipher_suites.is_empty()); + } + + #[test] + fn ensure_crypto_provider_is_idempotent() { + ensure_crypto_provider(); + let first = rustls::crypto::CryptoProvider::get_default().cloned(); + + ensure_crypto_provider(); + let second = rustls::crypto::CryptoProvider::get_default().cloned(); + + assert!(first.is_some()); + assert!(std::sync::Arc::ptr_eq(&first.unwrap(), &second.unwrap())); + } +} From 2c08e7abf8bbf1d32aab48a330f4aadd494b5dd0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:27:00 -0700 Subject: [PATCH 02/49] test(ai-gateway): pin the crypto provider to ring and prove the dial installs it The two tests that shipped with the fix both called ensure_crypto_provider themselves, so deleting the call from connect_upstream left the whole suite green, and swapping ring for aws-lc-rs did too. Adds an integration test, which gets its own process, that dials wss:// at a local plain-TCP listener through the public Responses WebSocket entrypoint and asserts an Err plus an installed provider. Without the install in the dial it panics with the original CryptoProvider message. A unit test now compares the installed provider's cipher suites and key-exchange groups against ring's, so the choice of backend is pinned rather than assumed. Also names tls12 in the workspace rustls features: it already arrives through reqwest and tokio-rustls, so the graph is unchanged, but a direct dependency should say it needs TLS 1.2 rather than inherit it. --- litellm-rust/Cargo.toml | 2 +- litellm-rust/crates/ai-gateway/src/io/tls.rs | 43 +++++++++++++++++-- .../tests/crypto_provider_wiring.rs | 41 ++++++++++++++++++ 3 files changed, 81 insertions(+), 5 deletions(-) create mode 100644 litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 15074082981..2e2e8809b7f 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -26,7 +26,7 @@ pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" -rustls = { version = "0.23", default-features = false, features = ["ring", "std"] } +rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" sha2 = "0.10" diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs index e3754e44002..96544adffb7 100644 --- a/litellm-rust/crates/ai-gateway/src/io/tls.rs +++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs @@ -8,7 +8,7 @@ //! Installing from the dial rather than from a `main` also covers the `cdylib` //! the Python bridge loads, the tests, and the benches, none of which have one. //! ring is what reqwest already falls back to, so installing it changes no -//! working path, and an embedder that installed its own provider first keeps it. +//! working path, and whoever installs into this rustls build first still wins. use std::sync::Once; @@ -38,13 +38,32 @@ where #[cfg(test)] mod tests { + use rustls::crypto::CryptoProvider; + use super::ensure_crypto_provider; + fn fingerprint( + provider: &CryptoProvider, + ) -> (Vec, Vec) { + ( + provider + .cipher_suites + .iter() + .map(|suite| suite.suite()) + .collect(), + provider + .kx_groups + .iter() + .map(|group| group.name()) + .collect(), + ) + } + #[test] fn client_config_builder_works_with_both_provider_features_enabled() { ensure_crypto_provider(); - assert!(rustls::crypto::CryptoProvider::get_default().is_some()); + assert!(CryptoProvider::get_default().is_some()); let config = rustls::ClientConfig::builder() .with_root_certificates(rustls::RootCertStore::empty()) @@ -53,13 +72,29 @@ mod tests { assert!(!config.crypto_provider().cipher_suites.is_empty()); } + #[test] + fn installs_ring_rather_than_aws_lc_rs() { + ensure_crypto_provider(); + + let installed = CryptoProvider::get_default().expect("a provider is installed"); + + assert_eq!( + fingerprint(installed), + fingerprint(&rustls::crypto::ring::default_provider()) + ); + assert_ne!( + fingerprint(installed), + fingerprint(&rustls::crypto::aws_lc_rs::default_provider()) + ); + } + #[test] fn ensure_crypto_provider_is_idempotent() { ensure_crypto_provider(); - let first = rustls::crypto::CryptoProvider::get_default().cloned(); + let first = CryptoProvider::get_default().cloned(); ensure_crypto_provider(); - let second = rustls::crypto::CryptoProvider::get_default().cloned(); + let second = CryptoProvider::get_default().cloned(); assert!(first.is_some()); assert!(std::sync::Arc::ptr_eq(&first.unwrap(), &second.unwrap())); diff --git a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs new file mode 100644 index 00000000000..db5698f8460 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs @@ -0,0 +1,41 @@ +//! Guards the wiring, not just the helper: the dial itself has to install the +//! rustls provider, in a test binary where nothing else has installed one. + +use std::collections::HashMap; +use std::time::Duration; + +use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection; +use tokio::net::TcpListener; + +#[tokio::test] +async fn dialing_wss_returns_an_error_instead_of_panicking() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + let result = ResponsesWebSocketConnection::connect_url( + &format!("wss://127.0.0.1:{port}/"), + &HashMap::new(), + Some(Duration::from_secs(10)), + ) + .await; + + assert!( + result.is_err(), + "a plain TCP server cannot finish a TLS handshake" + ); + assert!( + rustls::crypto::CryptoProvider::get_default().is_some(), + "the dial is what installs the process-wide provider" + ); +} From f9a5f676a866012c24427fb8afdfa240031f70d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:16:39 -0700 Subject: [PATCH 03/49] docs(claude): have runs embed their own QA screenshots on visual changes --- CLAUDE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CLAUDE.md b/CLAUDE.md index d9e9e8f1586..9c525ffe420 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -37,7 +37,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank -Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR +Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: - don't use emojis From 6463994f78161ee4448380a5e9fb2ad28729b98f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:29:10 -0700 Subject: [PATCH 04/49] fix(ai-gateway): dial upstream WebSockets over an explicit rustls provider Build one ClientConfig that names ring and loads the native roots once, and hand it to every tokio-tungstenite dial as its connector instead of installing a process-wide default from the dial path. Each of the three dial sites gets a wss:// test that reproduces the panic if its connector is dropped. --- litellm-rust/Cargo.lock | 1 + litellm-rust/Cargo.toml | 1 + litellm-rust/crates/ai-gateway/Cargo.toml | 5 +- .../crates/ai-gateway/src/io/realtime.rs | 27 +++++ .../crates/ai-gateway/src/io/responses_ws.rs | 23 ++++ litellm-rust/crates/ai-gateway/src/io/tls.rs | 113 ++++++++++-------- .../tests/crypto_provider_wiring.rs | 19 ++- 7 files changed, 134 insertions(+), 55 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2ed998174e4..d214e80d818 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1416,6 +1416,7 @@ dependencies = [ "pyo3", "reqwest", "rustls 0.23.42", + "rustls-native-certs", "serde", "serde_json", "sha2 0.10.9", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 2e2e8809b7f..d3d25cbda8d 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -27,6 +27,7 @@ rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } +rustls-native-certs = "0.8" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" sha2 = "0.10" diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 414abc2356d..82bedd0c8a0 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -19,9 +19,10 @@ litellm-core = { workspace = true, features = ["bedrock-auth"] } # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the # Python proxy callbacks API. reqwest.workspace = true -# rustls is a direct dependency so `io::tls` can install a process-level -# crypto provider; see that module for why the graph needs one. +# rustls and its root store are direct dependencies so `io::tls` can build the +# one TLS config the outbound dials use; see that module for why it has to. rustls.workspace = true +rustls-native-certs.workspace = true # `sync` powers the bounded mpsc channel the realtime logger drains. tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] } tokio-tungstenite.workspace = true diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 53d87848342..207c31dffa0 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -286,6 +286,33 @@ mod tests { serde_json::from_str(raw).expect("valid event json") } + /// The realtime dial has to reach a `wss://` upstream without a process-wide + /// crypto provider installed, which is what dialing through `io::tls` buys. + #[tokio::test] + async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + let result = dial_upstream( + "gpt-realtime", + "sk-test", + Some(&format!("wss://127.0.0.1:{port}")), + ) + .await; + + assert!(matches!(result, Err(Error::Network(_)))); + } + #[test] fn resolve_api_key_prefers_param_then_blank_falls_through() { assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test"); diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index ee181791413..9df3d0c6cc5 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -326,6 +326,29 @@ mod tests { use tokio::net::TcpListener; use tokio_tungstenite::accept_async; + /// The Responses dial has to reach a `wss://` upstream without a process-wide + /// crypto provider installed, which is what dialing through `io::tls` buys. + #[tokio::test] + async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + let result = + dial_upstream("gpt-5", "sk-test", Some(&format!("wss://127.0.0.1:{port}"))).await; + + assert!(matches!(result, Err(Error::Network(_)))); + } + async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); let address = listener.local_addr().expect("local address"); diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs index 96544adffb7..16fd11e2e79 100644 --- a/litellm-rust/crates/ai-gateway/src/io/tls.rs +++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs @@ -1,29 +1,54 @@ -//! Outbound WebSocket dials, with the rustls crypto provider settled first. +//! Outbound WebSocket dials over a TLS config this crate builds once and owns. //! //! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth` -//! enables `rustls/aws-lc-rs`, so `ClientConfig::builder()` — which is how -//! `tokio-tungstenite` builds its TLS config — panics rather than guess between -//! them. reqwest and the AWS SDK pick a provider explicitly and never panic. -//! -//! Installing from the dial rather than from a `main` also covers the `cdylib` -//! the Python bridge loads, the tests, and the benches, none of which have one. -//! ring is what reqwest already falls back to, so installing it changes no -//! working path, and whoever installs into this rustls build first still wins. +//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that +//! `tokio-tungstenite` uses when handed no connector panics rather than guess +//! between them. Naming ring on a connector of our own settles that for these +//! dials without touching the process-wide default, and building the config +//! once keeps the platform trust store, which `tokio-tungstenite` would +//! otherwise re-read on every dial, off the dial path. -use std::sync::Once; +use std::io; +use std::sync::{Arc, OnceLock}; +use rustls::{ClientConfig, RootCertStore}; use tokio::net::TcpStream; use tokio_tungstenite::tungstenite::Error; use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::error::TlsError; use tokio_tungstenite::tungstenite::handshake::client::Response; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; +use tokio_tungstenite::{ + Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, +}; -static INSTALL_CRYPTO_PROVIDER: Once = Once::new(); +static TLS_CONFIG: OnceLock> = OnceLock::new(); -pub(crate) fn ensure_crypto_provider() { - INSTALL_CRYPTO_PROVIDER.call_once(|| { - let _ = rustls::crypto::ring::default_provider().install_default(); - }); +fn build_config() -> Result> { + let native = rustls_native_certs::load_native_certs(); + let roots = { + let mut store = RootCertStore::empty(); + let (added, _ignored) = store.add_parsable_certificates(native.certs); + if added == 0 { + return Err(Box::new(Error::Io(io::Error::other(format!( + "no usable native root certificates: {:?}", + native.errors + ))))); + } + store + }; + + ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) + .with_safe_default_protocol_versions() + .map(|builder| builder.with_root_certificates(roots).with_no_client_auth()) + .map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error)))) +} + +fn tls_config() -> Result, Box> { + if let Some(config) = TLS_CONFIG.get() { + return Ok(Arc::clone(config)); + } + let built = Arc::new(build_config()?); + Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built))) } pub(crate) async fn connect_upstream( @@ -32,19 +57,25 @@ pub(crate) async fn connect_upstream( where R: IntoClientRequest + Unpin, { - ensure_crypto_provider(); - connect_async(request).await.map_err(Box::new) + let request = request.into_client_request().map_err(Box::new)?; + let connector = match request.uri().scheme_str() { + Some("wss") => Some(Connector::Rustls(tls_config()?)), + _ => None, + }; + connect_async_tls_with_config(request, None, false, connector) + .await + .map_err(Box::new) } #[cfg(test)] mod tests { - use rustls::crypto::CryptoProvider; + use rustls::CipherSuite; + use rustls::NamedGroup; + use rustls::crypto::{CryptoProvider, aws_lc_rs, ring}; - use super::ensure_crypto_provider; + use super::{Arc, build_config, tls_config}; - fn fingerprint( - provider: &CryptoProvider, - ) -> (Vec, Vec) { + fn fingerprint(provider: &CryptoProvider) -> (Vec, Vec) { ( provider .cipher_suites @@ -60,43 +91,31 @@ mod tests { } #[test] - fn client_config_builder_works_with_both_provider_features_enabled() { - ensure_crypto_provider(); - - assert!(CryptoProvider::get_default().is_some()); - - let config = rustls::ClientConfig::builder() - .with_root_certificates(rustls::RootCertStore::empty()) - .with_no_client_auth(); + fn builds_a_usable_config_with_both_provider_features_enabled() { + let config = build_config().expect("a client config"); assert!(!config.crypto_provider().cipher_suites.is_empty()); } #[test] - fn installs_ring_rather_than_aws_lc_rs() { - ensure_crypto_provider(); - - let installed = CryptoProvider::get_default().expect("a provider is installed"); + fn dials_with_ring_rather_than_aws_lc_rs() { + let config = build_config().expect("a client config"); assert_eq!( - fingerprint(installed), - fingerprint(&rustls::crypto::ring::default_provider()) + fingerprint(config.crypto_provider()), + fingerprint(&ring::default_provider()) ); assert_ne!( - fingerprint(installed), - fingerprint(&rustls::crypto::aws_lc_rs::default_provider()) + fingerprint(config.crypto_provider()), + fingerprint(&aws_lc_rs::default_provider()) ); } #[test] - fn ensure_crypto_provider_is_idempotent() { - ensure_crypto_provider(); - let first = CryptoProvider::get_default().cloned(); + fn the_trust_store_is_loaded_once_and_shared() { + let first = tls_config().expect("a client config"); + let second = tls_config().expect("a client config"); - ensure_crypto_provider(); - let second = CryptoProvider::get_default().cloned(); - - assert!(first.is_some()); - assert!(std::sync::Arc::ptr_eq(&first.unwrap(), &second.unwrap())); + assert!(Arc::ptr_eq(&first, &second)); } } diff --git a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs index db5698f8460..05f7d9610d5 100644 --- a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs +++ b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs @@ -1,5 +1,6 @@ -//! Guards the wiring, not just the helper: the dial itself has to install the -//! rustls provider, in a test binary where nothing else has installed one. +//! Guards the wiring, not just the helper: a `wss://` dial through the public +//! API has to resolve its own crypto provider, in a test binary where nothing +//! has installed a process-wide one, and has to leave it uninstalled. use std::collections::HashMap; use std::time::Duration; @@ -7,8 +8,7 @@ use std::time::Duration; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection; use tokio::net::TcpListener; -#[tokio::test] -async fn dialing_wss_returns_an_error_instead_of_panicking() { +async fn dead_tls_server() -> u16 { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("bind a loopback port"); @@ -23,6 +23,13 @@ async fn dialing_wss_returns_an_error_instead_of_panicking() { } }); + port +} + +#[tokio::test] +async fn dialing_wss_returns_an_error_instead_of_panicking() { + let port = dead_tls_server().await; + let result = ResponsesWebSocketConnection::connect_url( &format!("wss://127.0.0.1:{port}/"), &HashMap::new(), @@ -35,7 +42,7 @@ async fn dialing_wss_returns_an_error_instead_of_panicking() { "a plain TCP server cannot finish a TLS handshake" ); assert!( - rustls::crypto::CryptoProvider::get_default().is_some(), - "the dial is what installs the process-wide provider" + rustls::crypto::CryptoProvider::get_default().is_none(), + "the dial settles its provider on its own connector, not process-wide" ); } From 199b44a475719500915dbd9059c0ea5dfeeae782 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:48:56 -0700 Subject: [PATCH 05/49] fix(async): move remote image fetches off the event loop for Snowflake, Bedrock invoke Claude, Mantle and Gemini --- .../prompt_templates/image_handling.py | 103 ++++++++++ litellm/llms/base_llm/chat/transformation.py | 4 + .../anthropic_claude3_transformation.py | 61 +----- .../bedrock/chat/mantle/transformation.py | 15 +- litellm/llms/custom_httpx/llm_http_handler.py | 194 ++++++++++-------- litellm/llms/snowflake/chat/transformation.py | 16 ++ .../llms/vertex_ai/gemini/transformation.py | 9 +- litellm/main.py | 8 +- tests/test_litellm/conftest.py | 35 ++++ .../litellm_core_utils/test_image_handling.py | 62 ++++++ ...ations_anthropic_claude3_transformation.py | 49 +++++ ...test_bedrock_chat_mantle_transformation.py | 51 +++++ .../custom_httpx/test_llm_http_handler.py | 106 +++++++++- .../test_snowflake_chat_transformation.py | 46 +++++ .../vertex_ai/gemini/test_transformation.py | 42 ++++ 15 files changed, 644 insertions(+), 157 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 81132fa89a5..1890d4eb682 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -2,7 +2,11 @@ Helper functions to handle images passed in messages """ +import asyncio import base64 +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType from typing import Final from httpx import Response @@ -12,6 +16,7 @@ from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get +from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 @@ -124,3 +129,101 @@ def convert_url_to_base64(url: str) -> str: raise litellm.ImageFetchError( f"Error: Unable to fetch image from URL after 3 attempts. url={url}", ) + + +_REMOTE_URL_PREFIXES: Final = ("http://", "https://") + + +@dataclass(frozen=True, slots=True) +class _RemoteImage: + part: Mapping[str, object] + image_url: Mapping[str, object] | None + url: str + + +@dataclass(frozen=True, slots=True) +class _RemoteFile: + part: Mapping[str, object] + file: Mapping[str, object] + url: str + + +def _as_mapping(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None # pyright: ignore[reportUnknownVariableType] # fields are parsed one by one + + +def _remote_url(candidate: object) -> str | None: + return candidate if isinstance(candidate, str) and candidate.startswith(_REMOTE_URL_PREFIXES) else None + + +def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | None: + fields: Final = _as_mapping(part) + if fields is None: + return None + if fields.get("type") == "image_url": + image_url: Final = fields.get("image_url") + image_url_fields: Final = _as_mapping(image_url) + url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url) + return _RemoteImage(fields, image_url_fields, url) if url is not None else None + file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None + file_url: Final = _remote_url(file.get("file_id")) if file is not None else None + return _RemoteFile(fields, file, file_url) if file is not None and file_url is not None else None + + +_PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"}) + + +def _inferred_format(file: Mapping[str, object], url: str) -> Mapping[str, str]: + return _PDF_FORMAT if "format" not in file and url.lower().endswith(".pdf") else MappingProxyType({}) + + +def _inlined_image_url(image_url: Mapping[str, object] | None, data_url: str) -> Mapping[str, object] | str: + return {**image_url, "url": data_url} if image_url is not None else data_url # mutable-ok: json-serialized part + + +def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mapping[str, object]: + kept: Final = {k: v for k, v in file.items() if k != "file_id"} # mutable-ok: json-serialized message part + return {**kept, **_inferred_format(file, url), "file_data": data_url} # mutable-ok: json-serialized part + + +def _inline(remote: _RemoteImage | _RemoteFile, data_url: str) -> Mapping[str, object]: + match remote: + case _RemoteImage(part, image_url, _): + return {**part, "image_url": _inlined_image_url(image_url, data_url)} # mutable-ok: json-serialized part + case _RemoteFile(part, file, url): + return {**part, "file": _inlined_file(file, url, data_url)} # mutable-ok: json-serialized message part + + +def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]: + content: Final = message.get("content") + return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one + + +def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> AllMessageValues: + parts: Final = _content_parts(message) + if not parts: + return message + inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks + _inline(remote, data_urls[remote.url]) if (remote := _parse_remote_part(part)) is not None else part + for part in parts + ] + inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message + return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined + + +async def async_inline_remote_media( + messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues] +) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues] + remote_urls: Final = tuple( + dict.fromkeys( + remote.url + for message in messages + for part in _content_parts(message) + if (remote := _parse_remote_part(part)) is not None + ) + ) + if not remote_urls: + return messages + data_urls: Final = await asyncio.gather(*(async_convert_url_to_base64(url) for url in remote_urls)) + inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True))) + return [_inline_message(message, inlined) for message in messages] # mutable-ok: transform_request takes a list diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index bbe1cc85df1..7bfc87a30d6 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -411,6 +411,10 @@ class BaseConfig(ABC): def has_custom_stream_wrapper(self) -> bool: return False + @property + def uses_async_transform_request(self) -> bool: + return False + @property def supports_stream_param_in_request_body(self) -> bool: """ diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 67720451c00..38f280eef03 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -8,7 +8,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( - async_convert_url_to_base64, + async_inline_remote_media, convert_url_to_base64, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -172,6 +172,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): return _anthropic_request + @property + def uses_async_transform_request(self) -> bool: + return True + async def async_transform_request( self, model: str, @@ -180,26 +184,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): litellm_params: dict, headers: dict, ) -> dict: - _anthropic_request: Final = self._build_bedrock_anthropic_request_base( + return self.transform_request( model=model, - messages=messages, + messages=await async_inline_remote_media(messages), optional_params=optional_params, litellm_params=litellm_params, headers=headers, ) - await self._async_convert_document_url_sources_to_base64(_anthropic_request) - beta_list: Final = self._compute_bedrock_invoke_beta_headers( - model=model, - messages=messages, - optional_params=optional_params, - headers=headers, - ) - if beta_list: - _anthropic_request["anthropic_beta"] = beta_list - - return _anthropic_request - def _build_bedrock_anthropic_request_base( self, model: str, @@ -321,45 +313,6 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): "data": image_chunk["data"], } - async def _async_convert_document_url_sources_to_base64(self, anthropic_request: dict) -> None: - """ - Async version of document URL conversion for async completion paths. - """ - messages: Final = anthropic_request.get("messages") - if not isinstance(messages, list): - return - - for message in messages: - if not isinstance(message, dict): - continue - content = message.get("content") - if not isinstance(content, list): - continue - - for block in content: - if not isinstance(block, dict) or block.get("type") != "document": - continue - source = block.get("source") - if not isinstance(source, dict) or source.get("type") != "url": - continue - source_url = source.get("url") - if not isinstance(source_url, str): - continue - - inferred_format: str | None = None - if source_url.lower().endswith(".pdf"): - inferred_format = "application/pdf" - base64_url = await async_convert_url_to_base64(url=source_url) - image_chunk = convert_to_anthropic_image_obj( - openai_image_url=base64_url, - format=inferred_format, - ) - block["source"] = { - "type": "base64", - "media_type": image_chunk["media_type"], - "data": image_chunk["data"], - } - def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict: """ Convert tool search entries to the format supported by the Bedrock Invoke API. diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index 31fc079c0c9..7e2037c33f1 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -10,6 +10,7 @@ at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth. from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Any, Final +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -110,21 +111,13 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): litellm_params: dict, headers: dict, ) -> dict: - model_id: Final = model.replace("mantle/", "", 1) - - request: Final = self._build_bedrock_anthropic_request_base( - model=model_id, - messages=messages, + return self.transform_request( + model=model, + messages=await async_inline_remote_media(messages), optional_params=optional_params, litellm_params=litellm_params, headers=headers, ) - await self._async_convert_document_url_sources_to_base64(request) - return self._restore_mantle_body_fields( - request=request, - model_id=model_id, - optional_params=optional_params, - ) @staticmethod def _restore_mantle_body_fields(request: dict, model_id: str, optional_params: dict) -> dict: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f281c249c72..54c21f8d5b3 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -116,6 +116,7 @@ from litellm.types.llms.anthropic_skills import ( Skill, ) from litellm.types.llms.openai import ( + AllMessageValues, CreateBatchRequest, CreateFileRequest, FileContentRequest, @@ -488,7 +489,7 @@ class BaseLLMHTTPHandler: def completion( self, model: str, - messages: list, + messages: list[AllMessageValues], api_base: str | None, custom_llm_provider: str, model_response: ModelResponse, @@ -507,7 +508,7 @@ class BaseLLMHTTPHandler: shared_session: Optional["ClientSession"] = None, ): json_mode: Final[bool] = optional_params.pop("json_mode", False) - extra_body: Final[dict | None] = optional_params.pop("extra_body", None) + extra_body: Final[Mapping[str, object] | None] = optional_params.pop("extra_body", None) provider_config = provider_config or ProviderConfigManager.get_provider_chat_config( model=model, provider=litellm.LlmProviders(custom_llm_provider) @@ -522,14 +523,17 @@ class BaseLLMHTTPHandler: ) # get config from model, custom llm provider - headers = provider_config.validate_environment( - api_key=api_key, - headers=headers or {}, - model=model, - messages=messages, - optional_params=optional_params, - api_base=api_base, - litellm_params=litellm_params, + request_headers: Final = cast( # cast-ok: validate_environment is declared as a bare dict + "dict[str, object]", + provider_config.validate_environment( + api_key=api_key, + headers=headers or {}, + model=model, + messages=messages, + optional_params=optional_params, + api_base=api_base, + litellm_params=litellm_params, + ), ) api_base = provider_config.get_complete_url( @@ -541,93 +545,117 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) - data: dict[str, object] = provider_config.transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=headers, - ) - - if extra_body is not None: - data = {**data, **extra_body} - - headers, signed_json_body = provider_config.sign_request( - headers=headers, - optional_params={ - **optional_params, - **_aws_signing_overrides(optional_params, litellm_params), - }, - request_data=data, - api_base=api_base, - api_key=api_key, - stream=stream, - fake_stream=fake_stream, - model=model, - ) - - ## LOGGING - logging_obj.pre_call( - input=messages, - api_key=api_key, - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": headers, - }, - ) - - # Check if stream was converted for WebSearch interception - # This is set by the async_pre_request_hook in WebSearchInterceptionLogger - if litellm_params.get("_websearch_interception_converted_stream", False): - logging_obj.model_call_details["websearch_interception_converted_stream"] = True - - if acompletion is True: - if stream is True: - data = self._add_stream_param_to_request_body( - data=data, - provider_config=provider_config, + def sign_and_log( + transformed: dict[str, object], # mutable-ok: async_completion takes dict + ) -> tuple[dict[str, object], dict[str, object], bytes | None]: # mutable-ok: async_completion takes dict + data: Final = {**transformed, **extra_body} if extra_body is not None else transformed + signed: Final = cast( # cast-ok: sign_request is declared as a bare dict + "tuple[dict[str, object], bytes | None]", + provider_config.sign_request( + headers=request_headers, + optional_params={ + **optional_params, + **_aws_signing_overrides(optional_params, litellm_params), + }, + request_data=data, + api_base=api_base, + api_key=api_key, + stream=stream, fake_stream=fake_stream, - ) + model=model, + ), + ) + logging_obj.pre_call( + input=messages, + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": signed[0], + }, + ) + if litellm_params.get("_websearch_interception_converted_stream", False): + logging_obj.model_call_details["websearch_interception_converted_stream"] = True + return data, signed[0], signed[1] + + def dispatch_async( + data: dict[str, object], # mutable-ok: async_completion takes dict + signed_headers: dict[str, object], # mutable-ok: async_completion takes dict + signed_json_body: bytes | None, + ): + async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None + if stream is True: return self.acompletion_stream_function( model=model, messages=messages, api_base=api_base, - headers=headers, + headers=signed_headers, custom_llm_provider=custom_llm_provider, provider_config=provider_config, timeout=timeout, logging_obj=logging_obj, - data=data, + data=self._add_stream_param_to_request_body( + data=data, + provider_config=provider_config, + fake_stream=fake_stream, + ), fake_stream=fake_stream, - client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), + client=async_client, litellm_params=litellm_params, json_mode=json_mode, optional_params=optional_params, signed_json_body=signed_json_body, ) + return self.async_completion( + custom_llm_provider=custom_llm_provider, + provider_config=provider_config, + api_base=api_base, + headers=signed_headers, + data=data, + timeout=timeout, + model=model, + model_response=model_response, + logging_obj=logging_obj, + api_key=api_key, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + client=async_client, + json_mode=json_mode, + signed_json_body=signed_json_body, + shared_session=shared_session, + ) - else: - return self.async_completion( - custom_llm_provider=custom_llm_provider, - provider_config=provider_config, - api_base=api_base, - headers=headers, - data=data, - timeout=timeout, - model=model, - model_response=model_response, - logging_obj=logging_obj, - api_key=api_key, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - encoding=encoding, - client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), - json_mode=json_mode, - signed_json_body=signed_json_body, - shared_session=shared_session, + if acompletion is True and provider_config.uses_async_transform_request: + + async def transform_then_dispatch(): + transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict + "dict[str, object]", + await provider_config.async_transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=request_headers, + ), ) + return await dispatch_async(*sign_and_log(transformed)) + + return transform_then_dispatch() + + data, signed_headers, signed_json_body = sign_and_log( + provider_config.transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=request_headers, + ) + ) + + if acompletion is True: + return dispatch_async(data, signed_headers, signed_json_body) if stream is True: data = self._add_stream_param_to_request_body( @@ -641,7 +669,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, messages=messages, @@ -651,7 +679,7 @@ class BaseLLMHTTPHandler: completion_stream, headers = self.make_sync_call( provider_config=provider_config, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, original_data=data, @@ -684,7 +712,7 @@ class BaseLLMHTTPHandler: sync_httpx_client=sync_httpx_client, provider_config=provider_config, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, timeout=timeout, diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index c64fc583edc..f65b0876202 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( create_anthropic_image_param, select_anthropic_content_block_type_for_file, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload @@ -421,6 +422,21 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body) return self._transform_request_openai(model, messages, optional_params, stream, extra_body) + @property + def uses_async_transform_request(self) -> bool: + return True + + async def async_transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: BaseConfig signature + optional_params: dict[str, object], # mutable-ok: BaseConfig signature + litellm_params: dict[str, object], # mutable-ok: BaseConfig signature + headers: dict[str, object], # mutable-ok: BaseConfig signature + ) -> dict[str, object]: # mutable-ok: BaseConfig signature + inlined_messages: Final = await async_inline_remote_media(messages) if _is_claude_model(model) else messages + return self.transform_request(model, inlined_messages, optional_params, litellm_params, headers) + def _transform_request_openai( self, model: str, diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index e2d62be6a69..c9480f07150 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, response_schema_prompt, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.files import ( @@ -1348,13 +1349,15 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) - if _openai_messages_may_need_sync_gcs_metadata_fetch(messages): + inlined_messages: Final = await async_inline_remote_media(messages) if custom_llm_provider == "gemini" else messages + + if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): # _transform_request_body may issue a sync httpx.get (up to 5s timeout) # via _get_gcs_object_content_type to fetch GCS object metadata. Run the # whole sync transformation on a worker thread so it does not block the # async event loop. return await asyncify(_transform_request_body)( - messages=messages, + messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, @@ -1363,7 +1366,7 @@ async def async_transform_request_body( ) return _transform_request_body( - messages=messages, + messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, diff --git a/litellm/main.py b/litellm/main.py index 2929790f2bd..9970c203e03 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4474,7 +4474,7 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client = _dispatch_client_http(ctx) + injected_client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4486,11 +4486,11 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR shared_session: Final = ctx.shared_session stream: Final = ctx.stream timeout: Final = ctx.timeout + client: Final = ( + injected_client if injected_client is not None else (HTTPHandler(timeout=timeout) if stream is False else None) + ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible try: - client = ( - HTTPHandler(timeout=timeout) if stream is False else None - ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible response: Final = base_llm_http_handler.completion( model=model, messages=messages, diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index 62c95cb100b..aa8ba168ecc 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -7,9 +7,12 @@ # 4. Added proper cleanup in fixtures # 5. Added worker-specific isolation for parallel execution +import base64 import importlib import os from pathlib import Path +from types import SimpleNamespace +import httpx import pytest import asyncio @@ -595,3 +598,35 @@ def pytest_sessionfinish(session, exitstatus): _close_handler_if_needed(getattr(litellm, "aclient", None)) _close_handler_if_needed(getattr(litellm, "client", None)) _run_coroutine_if_needed(close_litellm_async_clients()) + + +ONE_PIXEL_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +@pytest.fixture +def async_only_image_fetch(monkeypatch): + from litellm.litellm_core_utils.prompt_templates import image_handling + + fetch = SimpleNamespace( + fetched=[], + base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(), + data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(), + ) + + def forbid_sync_fetch(client, url, **kwargs): + raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}") + + async def serve_png(client, url, **kwargs): + fetch.fetched.append(url) + return httpx.Response( + 200, + content=ONE_PIXEL_PNG, + headers={"content-type": "image/png"}, + request=httpx.Request("GET", url), + ) + + monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(image_handling, "async_safe_get", serve_png) + return fetch diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index 893472d63ae..5b6d403fa21 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -1,3 +1,5 @@ +import copy +import uuid from unittest.mock import patch import pytest @@ -8,6 +10,7 @@ from litellm import constants from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, + async_inline_remote_media, convert_url_to_base64, ) @@ -268,3 +271,62 @@ def test_image_size_limit_disabled(monkeypatch): assert "Image URL download is disabled" in str(excinfo.value) assert "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0" in str(excinfo.value) + + +async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf" + messages = [ + {"role": "system", "content": "be terse"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this?"}, + {"type": "image_url", "image_url": {"url": image_url, "detail": "low"}}, + {"type": "image_url", "image_url": image_url}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + {"type": "file", "file": {"file_id": pdf_url}}, + {"type": "file", "file": {"file_id": image_url, "format": "image/png"}}, + ], + }, + ] + snapshot = copy.deepcopy(messages) + + inlined = await async_inline_remote_media(messages) + + data_url = async_only_image_fetch.data_url + assert inlined[0] == {"role": "system", "content": "be terse"} + assert inlined[1]["content"] == [ + {"type": "text", "text": "what is this?"}, + {"type": "image_url", "image_url": {"url": data_url, "detail": "low"}}, + {"type": "image_url", "image_url": data_url}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + {"type": "file", "file": {"format": "application/pdf", "file_data": data_url}}, + {"type": "file", "file": {"format": "image/png", "file_data": data_url}}, + ] + assert sorted(async_only_image_fetch.fetched) == sorted([image_url, pdf_url]) + assert messages == snapshot + + +async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch): + messages = [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}], + }, + ] + + assert await async_inline_remote_media(messages) is messages + assert async_only_image_fetch.fetched == [] + + +async def test_async_inline_remote_media_raises_image_fetch_error_when_the_fetch_fails(monkeypatch): + async def serve_404(client, url, **kwargs): + return Response(404, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve_404) + url = f"http://img.example/{uuid.uuid4()}.png" + + with pytest.raises(litellm.ImageFetchError, match="Status code: 404"): + await async_inline_remote_media([{"role": "user", "content": [{"type": "image_url", "image_url": url}]}]) diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 0d7573a2536..e808a087e3d 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1,15 +1,19 @@ import asyncio import json +import uuid from unittest.mock import patch +import httpx import pytest # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. +import litellm from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler def test_get_supported_params_thinking(): @@ -714,3 +718,48 @@ def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking assert result["thinking"] == {"type": "adaptive"} assert result["output_config"] == {"effort": "high"} + + +async def test_bedrock_invoke_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/invoke/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] diff --git a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py new file mode 100644 index 00000000000..9b491d305c7 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -0,0 +1,51 @@ +import uuid + +import httpx + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + +async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/mantle/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 023e2d8843f..e86e671b939 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -17,7 +17,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, ) -from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, @@ -30,7 +30,7 @@ from litellm.llms.azure.videos.transformation import AzureVideoConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import TranscriptionResponse +from litellm.types.utils import ModelResponse, TranscriptionResponse _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" @@ -3186,3 +3186,105 @@ async def test_async_container_list_handler_transforms_success_response(): assert [container.id for container in response.data] == ["cntr_a"] assert response.has_more is True + + +class _TransformRecordingConfig(BaseConfig): + def __init__(self, transform_async: bool): + self.transform_async = transform_async + self.transform_calls = [] + + @property + def uses_async_transform_request(self) -> bool: + return self.transform_async + + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, non_default_params, optional_params, model, drop_params): + return optional_params + + def validate_environment( + self, headers, model, messages, optional_params, litellm_params, api_key=None, api_base=None + ): + return {} + + def transform_request(self, model, messages, optional_params, litellm_params, headers): + self.transform_calls.append("sync") + return {"transformed_by": "sync"} + + async def async_transform_request(self, model, messages, optional_params, litellm_params, headers): + self.transform_calls.append("async") + return {"transformed_by": "async"} + + def transform_response( + self, + model, + raw_response, + model_response, + logging_obj, + request_data, + messages, + optional_params, + litellm_params, + encoding, + api_key=None, + json_mode=None, + ): + model_response.choices[0].message.content = raw_response.json()["transformed_by"] + return model_response + + def get_error_class(self, error_message, status_code, headers): + return BaseLLMException(status_code=status_code, message=error_message, headers=headers) + + +def _start_async_completion(config): + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=captured["body"]) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + pending = BaseLLMHTTPHandler().completion( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + api_base="https://llm.example/v1/chat", + custom_llm_provider="openai", + model_response=ModelResponse(), + encoding=None, + logging_obj=Mock(dynamic_success_callbacks=None, model_call_details={}), + optional_params={}, + timeout=10.0, + litellm_params={}, + acompletion=True, + client=client, + provider_config=config, + ) + return pending, captured + + +async def test_completion_awaits_async_transform_request_when_config_opts_in(): + config = _TransformRecordingConfig(transform_async=True) + + pending, captured = _start_async_completion(config) + assert config.transform_calls == [] + + response = await pending + + assert config.transform_calls == ["async"] + assert captured["body"] == {"transformed_by": "async"} + assert response.choices[0].message.content == "async" + + +async def test_completion_keeps_sync_transform_request_before_returning_by_default(): + config = _TransformRecordingConfig(transform_async=False) + + pending, captured = _start_async_completion(config) + assert config.transform_calls == ["sync"] + + response = await pending + + assert config.transform_calls == ["sync"] + assert captured["body"] == {"transformed_by": "sync"} + assert response.choices[0].message.content == "sync" diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py index 25a961c3413..5687a319f06 100644 --- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -6,6 +6,7 @@ Tests tool calling request/response transformations and chat completions import asyncio import os import copy +import uuid import json from typing import Any, Dict, List @@ -945,3 +946,48 @@ class TestSnowflakeChatCompletion: assert len(chunks_received) > 0 content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content) + + +async def test_snowflake_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="snowflake/claude-sonnet-4-6", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + api_key="fake-jwt", + account_id="FAKE-ACCOUNT", + api_base=FAKE_API_BASE, + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py index fad310fc5c0..ec445342523 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py @@ -1,6 +1,10 @@ +import uuid +import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.gemini import transformation from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, @@ -338,3 +342,41 @@ def test_map_function_enterprise_web_search_snake_case(): assert len(result) == 1 assert "enterpriseWebSearch" in result[0] + + +async def test_gemini_ai_studio_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "Green"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + api_key="fake-gemini-key", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] From e355203014326462c3e70e6b6c8d37fcc0abe656 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:11:41 -0700 Subject: [PATCH 06/49] fix(gemini): keep Files API refs out of the async remote media walker Gemini AI Studio file URIs under generativelanguage.googleapis.com/v1beta/files/ answer 403 when fetched and must pass through as file_data.file_uri, which the sync transform already did. Give the shared walker a skip_url_prefixes parameter, pass that prefix from the Gemini body builder, and skip it in the AI Studio message transform for both image_url and file parts --- .../prompt_templates/image_handling.py | 12 +++-- litellm/llms/gemini/chat/transformation.py | 14 ++++-- .../llms/vertex_ai/gemini/transformation.py | 9 +++- .../litellm_core_utils/test_image_handling.py | 28 ++++++++++++ .../vertex_ai/gemini/test_transformation.py | 44 +++++++++++++++++++ 5 files changed, 99 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 1890d4eb682..5ae8d224556 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -199,13 +199,18 @@ def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]: return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one +def _inline_part(part: object, data_urls: Mapping[str, str]) -> object: + remote: Final = _parse_remote_part(part) + data_url: Final = data_urls.get(remote.url) if remote is not None else None + return _inline(remote, data_url) if remote is not None and data_url is not None else part + + def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> AllMessageValues: parts: Final = _content_parts(message) if not parts: return message inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks - _inline(remote, data_urls[remote.url]) if (remote := _parse_remote_part(part)) is not None else part - for part in parts + _inline_part(part, data_urls) for part in parts ] inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined @@ -213,13 +218,14 @@ def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> async def async_inline_remote_media( messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues] + skip_url_prefixes: tuple[str, ...] = (), ) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues] remote_urls: Final = tuple( dict.fromkeys( remote.url for message in messages for part in _content_parts(message) - if (remote := _parse_remote_part(part)) is not None + if (remote := _parse_remote_part(part)) is not None and not remote.url.startswith(skip_url_prefixes) ) ) if not remote_urls: diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 1a67b33665b..42c9ef13730 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionFileObject from litellm.types.llms.vertex_ai import ContentType, PartType from litellm.utils import supports_reasoning -from ...vertex_ai.gemini.transformation import _gemini_convert_messages_with_history +from ...vertex_ai.gemini.transformation import GEMINI_FILES_API_URI_PREFIX, _gemini_convert_messages_with_history from ...vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -127,7 +127,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): if element.get("type") == "image_url": img_element = cast(ChatCompletionImageObject, element) # cast-ok: runtime type tag checked _image_url, format, detail = _image_url_fields(img_element) - if _image_url and "https://" in _image_url: + if ( + _image_url + and "https://" in _image_url + and not _image_url.startswith(GEMINI_FILES_API_URI_PREFIX) + ): image_obj = convert_to_anthropic_image_obj(_image_url, format=format) converted_image_url = convert_generic_image_chunk_to_openai_image_obj(image_obj) if detail is not None: @@ -147,7 +151,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): llm_provider="gemini", ) file_id = _file_field.get("file_id") - if file_id and ("http://" in file_id or "https://" in file_id): + if ( + file_id + and ("http://" in file_id or "https://" in file_id) + and not file_id.startswith(GEMINI_FILES_API_URI_PREFIX) + ): # Convert HTTP/HTTPS file URL to base64 data try: base64_data = convert_url_to_base64(file_id) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index c9480f07150..6d100143e52 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -69,6 +69,7 @@ _GCS_METADATA_VERTEX_BASE: Any | None = None # Shared sync client for GCS JSON API metadata reads so proxy/SSL settings # from litellm's HTTP stack apply (see Greptile review on PR #27278). _GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None +GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/" _GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = { "image/jpg": "image/jpeg", } @@ -557,7 +558,7 @@ def _process_gemini_media( file_data = FileDataType(mime_type=mime_type, file_uri=image_url) part: PartType = {"file_data": file_data} return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata) - elif image_url.startswith("https://generativelanguage.googleapis.com/v1beta/files/"): + elif image_url.startswith(GEMINI_FILES_API_URI_PREFIX): # Gemini Files API URIs — the file is already uploaded to Google's # servers; pass the URI through as file_data without fetching it. # These URLs return 403 when accessed directly, so we must not try @@ -1349,7 +1350,11 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) - inlined_messages: Final = await async_inline_remote_media(messages) if custom_llm_provider == "gemini" else messages + inlined_messages: Final = ( + await async_inline_remote_media(messages, skip_url_prefixes=(GEMINI_FILES_API_URI_PREFIX,)) + if custom_llm_provider == "gemini" + else messages + ) if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): # _transform_request_body may issue a sync httpx.get (up to 5s timeout) diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index 5b6d403fa21..bcf4c7cb6ff 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -308,6 +308,34 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o assert messages == snapshot +async def test_async_inline_remote_media_leaves_skipped_url_prefixes_untouched(async_only_image_fetch): + skipped_prefix = "https://generativelanguage.googleapis.com/v1beta/files/" + skipped_file = f"{skipped_prefix}{uuid.uuid4().hex}" + skipped_image = f"{skipped_prefix}{uuid.uuid4().hex}" + fetched_image = f"https://img.example/{uuid.uuid4()}.png" + messages = [ + { + "role": "user", + "content": [ + {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}}, + {"type": "image_url", "image_url": {"url": skipped_image}}, + {"type": "image_url", "image_url": fetched_image}, + ], + } + ] + snapshot = copy.deepcopy(messages) + + inlined = await async_inline_remote_media(messages, skip_url_prefixes=(skipped_prefix,)) + + assert inlined[0]["content"] == [ + {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}}, + {"type": "image_url", "image_url": {"url": skipped_image}}, + {"type": "image_url", "image_url": async_only_image_fetch.data_url}, + ] + assert async_only_image_fetch.fetched == [fetched_image] + assert messages == snapshot + + async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch): messages = [ {"role": "user", "content": [{"type": "text", "text": "hi"}]}, diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py index ec445342523..d31254746d4 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py @@ -1,4 +1,5 @@ +import json import uuid import httpx import pytest @@ -380,3 +381,46 @@ async def test_gemini_ai_studio_async_completion_inlines_remote_images_off_the_e assert async_only_image_fetch.fetched == [image_url] assert image_url not in captured["body"] assert async_only_image_fetch.base64_png in captured["body"] + + +async def test_gemini_ai_studio_async_completion_passes_files_api_uris_through_unfetched(async_only_image_fetch): + files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + files_api_image = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "A report"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize these"}, + {"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}}, + {"type": "image_url", "image_url": {"url": files_api_image, "format": "image/png"}}, + ], + } + ], + api_key="fake-gemini-key", + client=client, + ) + + assert response.choices[0].message.content == "A report" + assert async_only_image_fetch.fetched == [] + file_parts = [part["file_data"] for part in captured["body"]["contents"][0]["parts"] if "file_data" in part] + assert file_parts == [ + {"mime_type": "application/pdf", "file_uri": files_api_pdf}, + {"mime_type": "image/png", "file_uri": files_api_image}, + ] From 08bb7de868f0d740c937585835333ac226faccad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:28:17 -0700 Subject: [PATCH 07/49] fix(image_handling): cap in-flight remote media fetches per request --- .../prompt_templates/image_handling.py | 9 +++++++- .../litellm_core_utils/test_image_handling.py | 22 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 5ae8d224556..4beaabc8b24 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 +MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20 in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY) @@ -216,6 +217,11 @@ def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined +async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str: + async with in_flight: + return await async_convert_url_to_base64(url) + + async def async_inline_remote_media( messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues] skip_url_prefixes: tuple[str, ...] = (), @@ -230,6 +236,7 @@ async def async_inline_remote_media( ) if not remote_urls: return messages - data_urls: Final = await asyncio.gather(*(async_convert_url_to_base64(url) for url in remote_urls)) + in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) + data_urls: Final = await asyncio.gather(*(_fetch_data_url(url, in_flight) for url in remote_urls)) inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True))) return [_inline_message(message, inlined) for message in messages] # mutable-ok: transform_request takes a list diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index bcf4c7cb6ff..ae9016f2f9e 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -1,3 +1,4 @@ +import asyncio import copy import uuid from unittest.mock import patch @@ -9,6 +10,7 @@ import litellm from litellm import constants from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.image_handling import ( + MAX_CONCURRENT_REMOTE_MEDIA_FETCHES, async_convert_url_to_base64, async_inline_remote_media, convert_url_to_base64, @@ -336,6 +338,26 @@ async def test_async_inline_remote_media_leaves_skipped_url_prefixes_untouched(a assert messages == snapshot +async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch): + in_flight = {"now": 0, "peak": 0} + + async def serve_png_slowly(client, url, **kwargs): + in_flight["now"] += 1 + in_flight["peak"] = max(in_flight["peak"], in_flight["now"]) + await asyncio.sleep(0.01) + in_flight["now"] -= 1 + return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve_png_slowly) + urls = [f"https://img.example/{uuid.uuid4()}.png" for _ in range(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES + 5)] + messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}} for url in urls]}] + + inlined = await async_inline_remote_media(messages) + + assert in_flight["peak"] == MAX_CONCURRENT_REMOTE_MEDIA_FETCHES + assert all(part["image_url"]["url"].startswith("data:image/png;base64,") for part in inlined[0]["content"]) + + async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch): messages = [ {"role": "user", "content": [{"type": "text", "text": "hi"}]}, From 567915aeee82953c4127456654f3a63ec9862970 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:51:14 -0700 Subject: [PATCH 08/49] fix(image_handling): inline url-sourced Anthropic document and image blocks in the async walker --- .../prompt_templates/image_handling.py | 51 +++++++++++++++---- .../litellm_core_utils/test_image_handling.py | 11 ++++ ...ations_anthropic_claude3_transformation.py | 47 +++++++++++++++++ ...test_bedrock_chat_mantle_transformation.py | 48 +++++++++++++++++ 4 files changed, 147 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 4beaabc8b24..55c01b3b1cb 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -157,18 +157,41 @@ def _remote_url(candidate: object) -> str | None: return candidate if isinstance(candidate, str) and candidate.startswith(_REMOTE_URL_PREFIXES) else None -def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | None: +_ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"}) + + +@dataclass(frozen=True, slots=True) +class _RemoteSource: + part: Mapping[str, object] + url: str + + +def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None: + if fields.get("type") != "image_url": + return None + image_url: Final = fields.get("image_url") + image_url_fields: Final = _as_mapping(image_url) + url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url) + return _RemoteImage(fields, image_url_fields, url) if url is not None else None + + +def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | None: + file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None + url: Final = _remote_url(file.get("file_id")) if file is not None else None + return _RemoteFile(fields, file, url) if file is not None and url is not None else None + + +def _parse_remote_source(fields: Mapping[str, object]) -> _RemoteSource | None: + source: Final = _as_mapping(fields.get("source")) if fields.get("type") in _ANTHROPIC_MEDIA_BLOCK_TYPES else None + url: Final = _remote_url(source.get("url")) if source is not None and source.get("type") == "url" else None + return _RemoteSource(fields, url) if url is not None else None + + +def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None: fields: Final = _as_mapping(part) if fields is None: return None - if fields.get("type") == "image_url": - image_url: Final = fields.get("image_url") - image_url_fields: Final = _as_mapping(image_url) - url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url) - return _RemoteImage(fields, image_url_fields, url) if url is not None else None - file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None - file_url: Final = _remote_url(file.get("file_id")) if file is not None else None - return _RemoteFile(fields, file, file_url) if file is not None and file_url is not None else None + return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields) _PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"}) @@ -187,12 +210,20 @@ def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mappin return {**kept, **_inferred_format(file, url), "file_data": data_url} # mutable-ok: json-serialized part -def _inline(remote: _RemoteImage | _RemoteFile, data_url: str) -> Mapping[str, object]: +def _base64_source(url: str, data_url: str) -> Mapping[str, str]: + fetched_media_type, data = data_url.removeprefix("data:").split(";base64,", 1) + media_type: Final = "application/pdf" if url.lower().endswith(".pdf") else fetched_media_type + return {"type": "base64", "media_type": media_type, "data": data} # mutable-ok: json-serialized message part + + +def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) -> Mapping[str, object]: match remote: case _RemoteImage(part, image_url, _): return {**part, "image_url": _inlined_image_url(image_url, data_url)} # mutable-ok: json-serialized part case _RemoteFile(part, file, url): return {**part, "file": _inlined_file(file, url, data_url)} # mutable-ok: json-serialized message part + case _RemoteSource(part, url): + return {**part, "source": _base64_source(url, data_url)} # mutable-ok: json-serialized message part def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]: diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index ae9016f2f9e..106bf90213b 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -289,6 +289,9 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, {"type": "file", "file": {"file_id": pdf_url}}, {"type": "file", "file": {"file_id": image_url, "format": "image/png"}}, + {"type": "document", "source": {"type": "url", "url": pdf_url}, "title": "the doc"}, + {"type": "image", "source": {"type": "url", "url": image_url}}, + {"type": "document", "source": {"type": "file", "file_id": "file_abc"}}, ], }, ] @@ -297,6 +300,7 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o inlined = await async_inline_remote_media(messages) data_url = async_only_image_fetch.data_url + base64_png = async_only_image_fetch.base64_png assert inlined[0] == {"role": "system", "content": "be terse"} assert inlined[1]["content"] == [ {"type": "text", "text": "what is this?"}, @@ -305,6 +309,13 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, {"type": "file", "file": {"format": "application/pdf", "file_data": data_url}}, {"type": "file", "file": {"format": "image/png", "file_data": data_url}}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": base64_png}, + "title": "the doc", + }, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": base64_png}}, + {"type": "document", "source": {"type": "file", "file_id": "file_abc"}}, ] assert sorted(async_only_image_fetch.fetched) == sorted([image_url, pdf_url]) assert messages == snapshot diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index e808a087e3d..cbf160c451f 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -763,3 +763,50 @@ async def test_bedrock_invoke_claude_async_completion_inlines_remote_images_off_ assert async_only_image_fetch.fetched == [image_url] assert image_url not in captured["body"] assert async_only_image_fetch.base64_png in captured["body"] + + +async def test_bedrock_invoke_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch): + pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "A lease"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/invoke/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this document?"}, + {"type": "document", "source": {"type": "url", "url": pdf_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "A lease" + assert async_only_image_fetch.fetched == [pdf_url] + assert { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png}, + } in captured["body"]["messages"][0]["content"] diff --git a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py index 9b491d305c7..a8448f5fa7a 100644 --- a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -1,3 +1,4 @@ +import json import uuid import httpx @@ -49,3 +50,50 @@ async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_ assert async_only_image_fetch.fetched == [image_url] assert image_url not in captured["body"] assert async_only_image_fetch.base64_png in captured["body"] + + +async def test_bedrock_mantle_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch): + pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "A lease"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/mantle/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this document?"}, + {"type": "document", "source": {"type": "url", "url": pdf_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "A lease" + assert async_only_image_fetch.fetched == [pdf_url] + assert { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png}, + } in captured["body"]["messages"][0]["content"] From 001b531ae960aad3fa38447508d4ce0d05c118c4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:11:27 -0700 Subject: [PATCH 09/49] fix(cost): bill realtime reasoning tokens nested in text_tokens once OpenAI and Azure realtime usage reports output_tokens == text_tokens + audio_tokens with reasoning_tokens counted inside text_tokens, so generic_cost_per_token billed the reasoning share twice. When the output token details sum past completion_tokens, the nested reasoning overlap is now subtracted from text_tokens before pricing; shapes where text_tokens already excludes reasoning are unchanged. --- .../litellm_core_utils/llm_cost_calc/utils.py | 19 ++++++- .../llm_cost_calc/test_llm_cost_calc_utils.py | 51 +++++++++++++++++++ tests/test_litellm/test_cost_calculator.py | 42 +++++++++++++++ 3 files changed, 111 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 8e24302b440..1b8980abd35 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -852,6 +852,17 @@ class CompletionTokensDetailsResult(TypedDict): video_tokens: int +def _text_tokens_without_nested_reasoning( + completion_tokens: int, + text_tokens: int, + reasoning_tokens: int, + other_modality_tokens: int, +) -> int: + reported_total: Final = text_tokens + reasoning_tokens + other_modality_tokens + nested_reasoning_tokens: Final = min(reasoning_tokens, text_tokens, max(reported_total - completion_tokens, 0)) + return text_tokens - nested_reasoning_tokens + + def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult: audio_tokens: Final = ( cast( @@ -860,7 +871,7 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu ) or 0 ) - text_tokens: Final = ( + reported_text_tokens: Final = ( cast( int | None, getattr(usage.completion_tokens_details, "text_tokens", None), @@ -882,6 +893,12 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu or 0 ) video_tokens: Final = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0)) + text_tokens: Final = _text_tokens_without_nested_reasoning( + completion_tokens=usage.completion_tokens, + text_tokens=reported_text_tokens, + reasoning_tokens=reasoning_tokens, + other_modality_tokens=audio_tokens + image_tokens + video_tokens, + ) return CompletionTokensDetailsResult( audio_tokens=audio_tokens, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 0f8084643ea..f7ca2a89048 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -4723,3 +4723,54 @@ def test_route_image_generation_cost_falls_back_to_requested_size(monkeypatch, r ) assert cost == expected_cost + + +def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_local_model_cost_map): + """ + Realtime usage (OpenAI and Azure) reports output_tokens == text_tokens + audio_tokens with + reasoning_tokens already counted inside text_tokens, so reasoning must not be billed on top. + """ + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=346, + completion_tokens=29, + total_tokens=375, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=152, image_tokens=194, audio_tokens=0, cached_tokens=128 + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=29, audio_tokens=0, reasoning_tokens=19 + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert completion_cost == pytest.approx(29 * info["output_cost_per_token"]) + assert completion_cost - breakdown.reasoning_cost == pytest.approx(10 * info["output_cost_per_token"]) + assert prompt_cost == pytest.approx( + 24 * info["input_cost_per_token"] + + 128 * info["cache_read_input_token_cost"] + + 194 * info["input_cost_per_image_token"] + ) + + +def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens(_local_model_cost_map): + """Providers whose text_tokens exclude reasoning (text + reasoning == completion) stay billed in full.""" + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=100, + completion_tokens=44, + total_tokens=144, + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=25, audio_tokens=0, reasoning_tokens=19 + ), + ) + + _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert completion_cost == pytest.approx(44 * info["output_cost_per_token"]) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 2046695f151..203f0a9e840 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4492,3 +4492,45 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m assert prompt_cost == pytest.approx(1000 * 5e-6) assert completion_cost == pytest.approx(500 * 2.5e-5) + + +def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once(_local_model_cost_map): + """Realtime response.done nests reasoning_tokens inside text_tokens, so they are billed once.""" + results: OpenAIRealtimeStreamList = [ + {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}}, + { + "type": "response.done", + "response": { + "usage": { + "total_tokens": 260, + "input_tokens": 237, + "output_tokens": 23, + "input_token_details": { + "text_tokens": 43, + "audio_tokens": 0, + "image_tokens": 194, + "cached_tokens": 0, + "cached_tokens_details": {"text_tokens": 0, "audio_tokens": 0, "image_tokens": 0}, + }, + "output_token_details": {"text_tokens": 23, "audio_tokens": 0, "reasoning_tokens": 18}, + } + }, + }, + ] + combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results, + ) + + total_cost = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=combined_usage_object, + custom_llm_provider="azure", + litellm_model_name="azure/gpt-realtime-2.1-mini", + ) + + info = litellm.get_model_info(model="azure/gpt-realtime-2.1-mini", custom_llm_provider="azure") + expected = ( + 43 * info["input_cost_per_token"] + 194 * info["input_cost_per_image_token"] + 23 * info["output_cost_per_token"] + ) + assert total_cost == pytest.approx(expected) + assert total_cost == pytest.approx(0.0002362) From 5a7e938f371e787d1e60af1eddea70bc10978813 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:40:46 -0700 Subject: [PATCH 10/49] test(cost): type the cost map fixture parameter on the new tests --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 4 ++-- tests/test_litellm/test_cost_calculator.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index f7ca2a89048..19464dfd2a0 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -4725,7 +4725,7 @@ def test_route_image_generation_cost_falls_back_to_requested_size(monkeypatch, r assert cost == expected_cost -def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_local_model_cost_map): +def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_local_model_cost_map: None) -> None: """ Realtime usage (OpenAI and Azure) reports output_tokens == text_tokens + audio_tokens with reasoning_tokens already counted inside text_tokens, so reasoning must not be billed on top. @@ -4757,7 +4757,7 @@ def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_loca ) -def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens(_local_model_cost_map): +def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens(_local_model_cost_map: None) -> None: """Providers whose text_tokens exclude reasoning (text + reasoning == completion) stay billed in full.""" model = "gpt-realtime-2.1-mini" diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 203f0a9e840..3eb051982e4 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4494,7 +4494,7 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m assert completion_cost == pytest.approx(500 * 2.5e-5) -def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once(_local_model_cost_map): +def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once(_local_model_cost_map: None) -> None: """Realtime response.done nests reasoning_tokens inside text_tokens, so they are billed once.""" results: OpenAIRealtimeStreamList = [ {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}}, From 92e449d4d0c486f24c7a09bfce988a6be5b1dd70 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:35:21 -0700 Subject: [PATCH 11/49] test(cost): cover reasoning nested in text_tokens beside audio output --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 19464dfd2a0..29874c5cca7 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -4774,3 +4774,22 @@ def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tok info = litellm.get_model_info(model=model, custom_llm_provider="openai") assert completion_cost == pytest.approx(44 * info["output_cost_per_token"]) + + +def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output(_local_model_cost_map: None) -> None: + """Audio-output realtime usage nests reasoning inside text_tokens next to audio_tokens; text is billed net of it.""" + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=120, + completion_tokens=100, + total_tokens=220, + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=30, audio_tokens=70, reasoning_tokens=20), + ) + + _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert breakdown.reasoning_cost == pytest.approx(20 * info["output_cost_per_token"]) + assert completion_cost == pytest.approx(30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"]) From f09992b73082e37d7337402954387def48fc7420 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:54:40 -0700 Subject: [PATCH 12/49] fix(image_handling): keep DNS, signing, and Vertex Gemini fetches off the event loop The SSRF check in async_safe_get resolved DNS on the event loop and a blocked address was retried three times; validate_url now runs in a thread and an SSRFError fails the fetch on the first attempt in both fetchers. The shared HTTP handler signed the request and ran pre_call logging on the loop after an async transform; both now run in a thread. Vertex AI Gemini still fetched http:// images and https images without an inferrable mime type with the sync converter inside its async body builder; the walker takes a should_inline predicate and Vertex AI inlines exactly those URLs, leaving https images with a known mime type and Files API refs to Google. When one download fails the other in-flight downloads for that request are now cancelled instead of finishing in the background --- .../prompt_templates/image_handling.py | 73 ++++++++-- litellm/litellm_core_utils/url_utils.py | 3 +- litellm/llms/custom_httpx/llm_http_handler.py | 2 +- .../llms/vertex_ai/gemini/transformation.py | 26 +++- .../litellm_core_utils/test_image_handling.py | 135 ++++++++++++++++-- .../litellm_core_utils/test_url_utils.py | 35 +++++ .../custom_httpx/test_llm_http_handler.py | 28 +++- .../vertex_ai/gemini/test_transformation.py | 51 +++++++ 8 files changed, 318 insertions(+), 35 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 55c01b3b1cb..dcad7776b58 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -4,7 +4,7 @@ Helper functions to handle images passed in messages import asyncio import base64 -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from types import MappingProxyType from typing import Final @@ -15,7 +15,7 @@ import litellm from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB -from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 @@ -99,6 +99,8 @@ async def async_convert_url_to_base64(url: str) -> str: return _process_image_response(response, url) except litellm.ImageFetchError: raise + except SSRFError as e: + raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e except Exception: pass raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}") @@ -125,6 +127,8 @@ def convert_url_to_base64(url: str) -> str: return _process_image_response(response, url) except litellm.ImageFetchError: raise + except SSRFError as e: + raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e except Exception as e: verbose_logger.exception(e) raise litellm.ImageFetchError( @@ -163,9 +167,23 @@ _ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"}) @dataclass(frozen=True, slots=True) class _RemoteSource: part: Mapping[str, object] + source: Mapping[str, object] url: str +@dataclass(frozen=True, slots=True) +class RemoteMedia: + url: str + fields: Mapping[str, object] + + +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def inline_every_remote_url(_media: RemoteMedia) -> bool: + return True + + def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None: if fields.get("type") != "image_url": return None @@ -184,7 +202,7 @@ def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | None: def _parse_remote_source(fields: Mapping[str, object]) -> _RemoteSource | None: source: Final = _as_mapping(fields.get("source")) if fields.get("type") in _ANTHROPIC_MEDIA_BLOCK_TYPES else None url: Final = _remote_url(source.get("url")) if source is not None and source.get("type") == "url" else None - return _RemoteSource(fields, url) if url is not None else None + return _RemoteSource(fields, source, url) if source is not None and url is not None else None def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None: @@ -194,6 +212,16 @@ def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSour return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields) +def _remote_media(remote: _RemoteImage | _RemoteFile | _RemoteSource) -> RemoteMedia: + match remote: + case _RemoteImage(_, image_url, url): + return RemoteMedia(url, image_url if image_url is not None else _NO_FIELDS) + case _RemoteFile(_, file, url): + return RemoteMedia(url, file) + case _RemoteSource(_, source, url): + return RemoteMedia(url, source) + + _PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"}) @@ -222,7 +250,7 @@ def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) - return {**part, "image_url": _inlined_image_url(image_url, data_url)} # mutable-ok: json-serialized part case _RemoteFile(part, file, url): return {**part, "file": _inlined_file(file, url, data_url)} # mutable-ok: json-serialized message part - case _RemoteSource(part, url): + case _RemoteSource(part, _, url): return {**part, "source": _base64_source(url, data_url)} # mutable-ok: json-serialized message part @@ -231,18 +259,22 @@ def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]: return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one -def _inline_part(part: object, data_urls: Mapping[str, str]) -> object: +def _inline_part(part: object, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]) -> object: remote: Final = _parse_remote_part(part) - data_url: Final = data_urls.get(remote.url) if remote is not None else None - return _inline(remote, data_url) if remote is not None and data_url is not None else part + if remote is None or not should_inline(_remote_media(remote)): + return part + data_url: Final = data_urls.get(remote.url) + return _inline(remote, data_url) if data_url is not None else part -def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> AllMessageValues: +def _inline_message( + message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool] +) -> AllMessageValues: parts: Final = _content_parts(message) if not parts: return message inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks - _inline_part(part, data_urls) for part in parts + _inline_part(part, data_urls, should_inline) for part in parts ] inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined @@ -253,21 +285,34 @@ async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str: return await async_convert_url_to_base64(url) +async def _fetch_data_urls(remote_urls: tuple[str, ...]) -> tuple[str, ...]: + in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) + fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight)) for url in remote_urls) + try: + return tuple(await asyncio.gather(*fetches)) + except BaseException: + for fetch in fetches: + fetch.cancel() + await asyncio.gather(*fetches, return_exceptions=True) + raise + + async def async_inline_remote_media( messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues] - skip_url_prefixes: tuple[str, ...] = (), + should_inline: Callable[[RemoteMedia], bool] = inline_every_remote_url, ) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues] remote_urls: Final = tuple( dict.fromkeys( remote.url for message in messages for part in _content_parts(message) - if (remote := _parse_remote_part(part)) is not None and not remote.url.startswith(skip_url_prefixes) + if (remote := _parse_remote_part(part)) is not None and should_inline(_remote_media(remote)) ) ) if not remote_urls: return messages - in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) - data_urls: Final = await asyncio.gather(*(_fetch_data_url(url, in_flight) for url in remote_urls)) + data_urls: Final = await _fetch_data_urls(remote_urls) inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True))) - return [_inline_message(message, inlined) for message in messages] # mutable-ok: transform_request takes a list + return [ # mutable-ok: transform_request takes a list + _inline_message(message, inlined, should_inline) for message in messages + ] diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 1e43117933d..fa070a648f5 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -19,6 +19,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): check but still resolve DNS and still rewrite HTTP to the resolved IP. """ +import asyncio import socket from ipaddress import ip_address, ip_network from typing import Any, Final, Protocol @@ -471,7 +472,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response kwargs.pop("follow_redirects", None) headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})} for _ in range(_MAX_REDIRECTS): - validated_url, original_host = validate_url(url) + validated_url, original_host = await asyncio.to_thread(validate_url, url) response = await fetcher.get( validated_url, headers={**headers_view["headers"], "Host": original_host}, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 54c21f8d5b3..d57aee025d5 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -640,7 +640,7 @@ class BaseLLMHTTPHandler: headers=request_headers, ), ) - return await dispatch_async(*sign_and_log(transformed)) + return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed)) return transform_then_dispatch() diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 6d100143e52..13e2238fdf6 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -7,6 +7,7 @@ Why separate file? Make it easy to see how transformation works import json import os import re +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote @@ -27,7 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, response_schema_prompt, ) -from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media +from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.files import ( @@ -1309,6 +1310,23 @@ def sync_transform_request_body( ) +def _explicit_mime_type(fields: Mapping[str, object]) -> str | None: + hint: Final = fields.get("format") or fields.get("mime_type") or fields.get("content_type") + return hint if isinstance(hint, str) else None + + +def _ai_studio_inlines(media: RemoteMedia) -> bool: + return not media.url.startswith(GEMINI_FILES_API_URI_PREFIX) + + +def _vertex_inlines(media: RemoteMedia) -> bool: + if media.url.startswith(GEMINI_FILES_API_URI_PREFIX): + return False + return media.url.startswith("http://") or ( + _explicit_mime_type(media.fields) is None and _get_image_mime_type_from_url(media.url) is None + ) + + async def async_transform_request_body( gemini_api_key: str | None, messages: list[AllMessageValues], @@ -1350,10 +1368,8 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) - inlined_messages: Final = ( - await async_inline_remote_media(messages, skip_url_prefixes=(GEMINI_FILES_API_URI_PREFIX,)) - if custom_llm_provider == "gemini" - else messages + inlined_messages: Final = await async_inline_remote_media( + messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines ) if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index 106bf90213b..e7b28d0d3a2 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -1,5 +1,6 @@ import asyncio import copy +import time import uuid from unittest.mock import patch @@ -11,10 +12,12 @@ from litellm import constants from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.image_handling import ( MAX_CONCURRENT_REMOTE_MEDIA_FETCHES, + RemoteMedia, async_convert_url_to_base64, async_inline_remote_media, convert_url_to_base64, ) +from litellm.litellm_core_utils.url_utils import SSRFError @pytest.fixture(autouse=True) @@ -321,34 +324,142 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o assert messages == snapshot -async def test_async_inline_remote_media_leaves_skipped_url_prefixes_untouched(async_only_image_fetch): - skipped_prefix = "https://generativelanguage.googleapis.com/v1beta/files/" - skipped_file = f"{skipped_prefix}{uuid.uuid4().hex}" - skipped_image = f"{skipped_prefix}{uuid.uuid4().hex}" - fetched_image = f"https://img.example/{uuid.uuid4()}.png" +async def test_async_inline_remote_media_inlines_only_the_parts_the_predicate_accepts(async_only_image_fetch): + files_api_prefix = "https://generativelanguage.googleapis.com/v1beta/files/" + files_api_pdf = f"{files_api_prefix}{uuid.uuid4().hex}" + hinted_image = f"https://img.example/{uuid.uuid4()}.png" + plain_image = f"https://img.example/{uuid.uuid4()}.png" + hinted_document = f"https://docs.example/{uuid.uuid4()}.pdf" + seen = [] + + def inline_unhinted_outside_files_api(media: RemoteMedia) -> bool: + seen.append(media) + return not media.url.startswith(files_api_prefix) and "format" not in media.fields + messages = [ { "role": "user", "content": [ - {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}}, - {"type": "image_url", "image_url": {"url": skipped_image}}, - {"type": "image_url", "image_url": fetched_image}, + {"type": "file", "file": {"file_id": files_api_pdf}}, + {"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": plain_image}}, + {"type": "image_url", "image_url": plain_image}, + {"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}}, ], } ] snapshot = copy.deepcopy(messages) - inlined = await async_inline_remote_media(messages, skip_url_prefixes=(skipped_prefix,)) + inlined = await async_inline_remote_media(messages, should_inline=inline_unhinted_outside_files_api) assert inlined[0]["content"] == [ - {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}}, - {"type": "image_url", "image_url": {"url": skipped_image}}, + {"type": "file", "file": {"file_id": files_api_pdf}}, + {"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}}, {"type": "image_url", "image_url": async_only_image_fetch.data_url}, + {"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}}, + ] + assert async_only_image_fetch.fetched == [plain_image] + assert [(media.url, dict(media.fields)) for media in seen[:5]] == [ + (files_api_pdf, {"file_id": files_api_pdf}), + (hinted_image, {"url": hinted_image, "format": "image/png"}), + (plain_image, {"url": plain_image}), + (plain_image, {}), + (hinted_document, {"type": "url", "url": hinted_document, "format": "application/pdf"}), ] - assert async_only_image_fetch.fetched == [fetched_image] assert messages == snapshot +async def test_async_inline_remote_media_inlines_a_shared_url_only_where_the_predicate_accepts_it( + async_only_image_fetch, +): + shared = f"https://img.example/{uuid.uuid4()}.png" + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": shared, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": shared}}, + ], + } + ] + + inlined = await async_inline_remote_media(messages, should_inline=lambda media: "format" not in media.fields) + + assert inlined[0]["content"] == [ + {"type": "image_url", "image_url": {"url": shared, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}}, + ] + assert async_only_image_fetch.fetched == [shared] + + +async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fails(monkeypatch): + missing = f"http://img.example/{uuid.uuid4()}-missing.png" + slow = f"http://img.example/{uuid.uuid4()}-slow.png" + slow_fetch_outcomes = [] + + async def serve(client, url, **kwargs): + if url == missing: + return Response(404, request=Request("GET", url)) + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + slow_fetch_outcomes.append("cancelled") + raise + slow_fetch_outcomes.append("finished") + return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve) + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": missing}}, + {"type": "image_url", "image_url": {"url": slow}}, + ], + } + ] + started = time.perf_counter() + + with pytest.raises(litellm.ImageFetchError, match="Status code: 404"): + await async_inline_remote_media(messages) + + assert slow_fetch_outcomes == ["cancelled"] + assert time.perf_counter() - started < 1 + + +async def test_async_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch): + attempts = [] + + async def block(client, url, **kwargs): + attempts.append(url) + raise SSRFError("URL targets a blocked address (10.0.0.8)") + + monkeypatch.setattr(image_handling, "async_safe_get", block) + url = f"http://internal.example/{uuid.uuid4()}.png" + + with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"): + await async_convert_url_to_base64(url) + + assert attempts == [url] + + +def test_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch): + attempts = [] + + def block(client, url, **kwargs): + attempts.append(url) + raise SSRFError("URL targets a blocked address (10.0.0.8)") + + monkeypatch.setattr(image_handling, "safe_get", block) + url = f"http://internal.example/{uuid.uuid4()}.png" + + with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"): + convert_url_to_base64(url) + + assert attempts == [url] + + async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch): in_flight = {"now": 0, "peak": 0} diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index aaaa43a0dc4..fccdc1a2a0a 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -1,5 +1,9 @@ +import asyncio import socket +import threading +import time +import httpx import pytest import litellm @@ -535,3 +539,34 @@ def test_assert_same_origin_error_message_does_not_leak_hostnames(): detail = str(exc.value) assert "attacker.example.com" not in detail assert "api.internal-corp.example" not in detail + + +async def test_async_safe_get_resolves_dns_off_the_event_loop(monkeypatch): + loop_thread = threading.current_thread() + resolver_threads = [] + + def slow_getaddrinfo(host, port, *args, **kwargs): + resolver_threads.append(threading.current_thread()) + time.sleep(0.4) + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port or 443))] + + monkeypatch.setattr(url_utils.socket, "getaddrinfo", slow_getaddrinfo) + + class FakeClient: + async def get(self, url, **kwargs): + return httpx.Response(200, request=httpx.Request("GET", url)) + + ticks = [time.perf_counter()] + + async def heartbeat(): + while True: + await asyncio.sleep(0.01) + ticks.append(time.perf_counter()) + + beating = asyncio.create_task(heartbeat()) + response = await url_utils.async_safe_get(FakeClient(), "https://img.example/a.png") + beating.cancel() + + assert response.status_code == 200 + assert resolver_threads and all(thread is not loop_thread for thread in resolver_threads) + assert max(b - a for a, b in zip(ticks, ticks[1:])) < 0.2 diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index e86e671b939..f1614654ffb 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import threading import time from unittest.mock import AsyncMock, Mock, patch @@ -3192,6 +3193,7 @@ class _TransformRecordingConfig(BaseConfig): def __init__(self, transform_async: bool): self.transform_async = transform_async self.transform_calls = [] + self.sign_threads = [] @property def uses_async_transform_request(self) -> bool: @@ -3216,6 +3218,12 @@ class _TransformRecordingConfig(BaseConfig): self.transform_calls.append("async") return {"transformed_by": "async"} + def sign_request( + self, headers, optional_params, request_data, api_base, api_key=None, model=None, stream=None, fake_stream=None + ): + self.sign_threads.append(threading.current_thread()) + return headers, None + def transform_response( self, model, @@ -3237,7 +3245,7 @@ class _TransformRecordingConfig(BaseConfig): return BaseLLMException(status_code=status_code, message=error_message, headers=headers) -def _start_async_completion(config): +def _start_async_completion(config, logging_obj=None): captured = {} def handle(request): @@ -3253,7 +3261,7 @@ def _start_async_completion(config): custom_llm_provider="openai", model_response=ModelResponse(), encoding=None, - logging_obj=Mock(dynamic_success_callbacks=None, model_call_details={}), + logging_obj=logging_obj if logging_obj is not None else Mock(dynamic_success_callbacks=None, model_call_details={}), optional_params={}, timeout=10.0, litellm_params={}, @@ -3277,6 +3285,22 @@ async def test_completion_awaits_async_transform_request_when_config_opts_in(): assert response.choices[0].message.content == "async" +async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_transform(): + config = _TransformRecordingConfig(transform_async=True) + loop_thread = threading.current_thread() + pre_call_threads = [] + logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={}) + logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread()) + + pending, captured = _start_async_completion(config, logging_obj) + response = await pending + + assert response.choices[0].message.content == "async" + assert captured["body"] == {"transformed_by": "async"} + assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) + assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads) + + async def test_completion_keeps_sync_transform_request_before_returning_by_default(): config = _TransformRecordingConfig(transform_async=False) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py index d31254746d4..f135acd094f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py @@ -1,6 +1,8 @@ import json import uuid +from unittest.mock import Mock + import httpx import pytest @@ -424,3 +426,52 @@ async def test_gemini_ai_studio_async_completion_passes_files_api_uris_through_u {"mime_type": "application/pdf", "file_uri": files_api_pdf}, {"mime_type": "image/png", "file_uri": files_api_image}, ] + + +async def test_vertex_ai_async_transform_inlines_only_the_urls_gemini_cannot_fetch_itself(async_only_image_fetch): + plain_http_png = f"http://img.example/{uuid.uuid4()}.png" + extensionless_https = f"https://cdn.example/files/{uuid.uuid4().hex}" + https_png = f"https://img.example/{uuid.uuid4()}.png" + hinted_extensionless = f"https://cdn.example/files/{uuid.uuid4().hex}" + files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe these"}, + {"type": "image_url", "image_url": {"url": plain_http_png}}, + {"type": "image_url", "image_url": {"url": extensionless_https}}, + {"type": "image_url", "image_url": {"url": https_png}}, + {"type": "image_url", "image_url": {"url": hinted_extensionless, "mime_type": "image/webp"}}, + {"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}}, + ], + } + ] + + body = await transformation.async_transform_request_body( + gemini_api_key=None, + messages=messages, + api_base=None, + model="gemini-3.8-flash", + client=None, + timeout=None, + extra_headers=None, + optional_params={}, + logging_obj=Mock(), + custom_llm_provider="vertex_ai", + litellm_params={}, + vertex_project="qa-project", + vertex_location="us-central1", + vertex_auth_header=None, + ) + + inlined = {"inline_data": {"mime_type": "image/png", "data": async_only_image_fetch.base64_png}} + assert body["contents"][0]["parts"] == [ + {"text": "Describe these"}, + inlined, + inlined, + {"file_data": {"mime_type": "image/png", "file_uri": https_png}}, + {"file_data": {"mime_type": "image/webp", "file_uri": hinted_extensionless}}, + {"file_data": {"mime_type": "application/pdf", "file_uri": files_api_pdf}}, + ] + assert sorted(async_only_image_fetch.fetched) == sorted([plain_http_png, extensionless_https]) From 836c20a4b6dd931d119ac2cce92977c3c947e062 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 22:31:06 -0700 Subject: [PATCH 13/49] test(ai-gateway): drop provider fingerprint and cache identity assertions --- litellm-rust/crates/ai-gateway/src/io/tls.rs | 43 +------------------- 1 file changed, 1 insertion(+), 42 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs index 16fd11e2e79..a2562f60345 100644 --- a/litellm-rust/crates/ai-gateway/src/io/tls.rs +++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs @@ -69,26 +69,7 @@ where #[cfg(test)] mod tests { - use rustls::CipherSuite; - use rustls::NamedGroup; - use rustls::crypto::{CryptoProvider, aws_lc_rs, ring}; - - use super::{Arc, build_config, tls_config}; - - fn fingerprint(provider: &CryptoProvider) -> (Vec, Vec) { - ( - provider - .cipher_suites - .iter() - .map(|suite| suite.suite()) - .collect(), - provider - .kx_groups - .iter() - .map(|group| group.name()) - .collect(), - ) - } + use super::build_config; #[test] fn builds_a_usable_config_with_both_provider_features_enabled() { @@ -96,26 +77,4 @@ mod tests { assert!(!config.crypto_provider().cipher_suites.is_empty()); } - - #[test] - fn dials_with_ring_rather_than_aws_lc_rs() { - let config = build_config().expect("a client config"); - - assert_eq!( - fingerprint(config.crypto_provider()), - fingerprint(&ring::default_provider()) - ); - assert_ne!( - fingerprint(config.crypto_provider()), - fingerprint(&aws_lc_rs::default_provider()) - ); - } - - #[test] - fn the_trust_store_is_loaded_once_and_shared() { - let first = tls_config().expect("a client config"); - let second = tls_config().expect("a client config"); - - assert!(Arc::ptr_eq(&first, &second)); - } } From e1d900d1c29e7f5dbd5db2ca09ab3592312bb1fe Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 22:47:09 -0700 Subject: [PATCH 14/49] fix(realtime): store text_tokens without the nested reasoning share The realtime usage writer passed the provider's output_token_details through as sent, so spend logs and callbacks kept a text_tokens that still contained reasoning_tokens while every other completion_tokens_details producer stores the partitioned share. The writer now applies the same rule the cost calculator uses, moved to litellm/types/utils.py so both read one definition, and the calculator keeps it for usage objects that arrive nested from elsewhere --- .../litellm_core_utils/llm_cost_calc/utils.py | 14 +----- litellm/responses/utils.py | 25 ++++++++-- litellm/types/utils.py | 11 ++++ .../responses/test_responses_utils.py | 50 +++++++++++++++++++ tests/test_litellm/test_cost_calculator.py | 39 +++++++++++++++ 5 files changed, 123 insertions(+), 16 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 1b8980abd35..9432fefc368 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -26,6 +26,7 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, ServiceTier, Usage, + text_tokens_without_nested_reasoning, ) from litellm.utils import get_model_info @@ -852,17 +853,6 @@ class CompletionTokensDetailsResult(TypedDict): video_tokens: int -def _text_tokens_without_nested_reasoning( - completion_tokens: int, - text_tokens: int, - reasoning_tokens: int, - other_modality_tokens: int, -) -> int: - reported_total: Final = text_tokens + reasoning_tokens + other_modality_tokens - nested_reasoning_tokens: Final = min(reasoning_tokens, text_tokens, max(reported_total - completion_tokens, 0)) - return text_tokens - nested_reasoning_tokens - - def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult: audio_tokens: Final = ( cast( @@ -893,7 +883,7 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu or 0 ) video_tokens: Final = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0)) - text_tokens: Final = _text_tokens_without_nested_reasoning( + text_tokens: Final = text_tokens_without_nested_reasoning( completion_tokens=usage.completion_tokens, text_tokens=reported_text_tokens, reasoning_tokens=reasoning_tokens, diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 39675faf735..28f590dcb66 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -25,9 +25,15 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, SpecialEnums, Usage, + text_tokens_without_nested_reasoning, ) +def _output_token_detail(details: object, field: str) -> int | None: + value: Final = getattr(details, field, None) + return value if isinstance(value, int) else None + + def _is_object_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a list is a Sequence of anything return isinstance(value, list) @@ -1134,11 +1140,22 @@ class ResponseAPILoggingUtils: response_api_usage, "output_tokens_details", None ) if output_tokens_details: + reasoning_tokens: Final = _output_token_detail(output_tokens_details, "reasoning_tokens") + image_tokens: Final = _output_token_detail(output_tokens_details, "image_tokens") + audio_tokens: Final = _output_token_detail(output_tokens_details, "audio_tokens") + reported_text_tokens: Final = _output_token_detail(output_tokens_details, "text_tokens") completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None), - image_tokens=getattr(output_tokens_details, "image_tokens", None), - text_tokens=getattr(output_tokens_details, "text_tokens", None), - audio_tokens=getattr(output_tokens_details, "audio_tokens", None), + reasoning_tokens=reasoning_tokens, + image_tokens=image_tokens, + text_tokens=None + if reported_text_tokens is None + else text_tokens_without_nested_reasoning( + completion_tokens=completion_tokens, + text_tokens=reported_text_tokens, + reasoning_tokens=reasoning_tokens or 0, + other_modality_tokens=(audio_tokens or 0) + (image_tokens or 0), + ), + audio_tokens=audio_tokens, ) extra_usage_fields: Final = { diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5052cd6ef48..9b24f1d837b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1622,6 +1622,17 @@ class Choices(SafeAttributeModel, OpenAIObject): setattr(self, key, value) +def text_tokens_without_nested_reasoning( + completion_tokens: int, + text_tokens: int, + reasoning_tokens: int, + other_modality_tokens: int, +) -> int: + reported_total: Final = text_tokens + reasoning_tokens + other_modality_tokens + nested_reasoning_tokens: Final = min(reasoning_tokens, text_tokens, max(reported_total - completion_tokens, 0)) + return text_tokens - nested_reasoning_tokens + + class CompletionTokensDetailsWrapper(CompletionTokensDetails): # wrapper for older openai versions text_tokens: int | None = None """Text tokens generated by the model.""" diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index cb6efa21036..9d9eefdceb3 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -440,6 +440,56 @@ class TestResponseAPILoggingUtils: assert result.completion_tokens_details.text_tokens == 20 assert result.completion_tokens_details.audio_tokens is None + def test_transform_realtime_usage_partitions_reasoning_out_of_text_tokens(self): + """Realtime nests reasoning_tokens inside text_tokens; the stored text share excludes them.""" + usage = { + "input_tokens": 237, + "output_tokens": 70, + "total_tokens": 307, + "input_token_details": {"text_tokens": 43, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 70, "audio_tokens": 0, "reasoning_tokens": 52}, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.completion_tokens == 70 + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 18 + assert result.completion_tokens_details.reasoning_tokens == 52 + assert result.completion_tokens_details.audio_tokens == 0 + + def test_transform_realtime_usage_partitions_reasoning_beside_audio_output(self): + """Audio output stays as reported; only the text share sheds the nested reasoning tokens.""" + usage = { + "input_tokens": 100, + "output_tokens": 70, + "total_tokens": 170, + "input_token_details": {"text_tokens": 100, "audio_tokens": 0, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 39, "audio_tokens": 31, "reasoning_tokens": 23}, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 16 + assert result.completion_tokens_details.audio_tokens == 31 + assert result.completion_tokens_details.reasoning_tokens == 23 + + def test_transform_response_api_usage_keeps_partitioned_text_tokens(self): + """A provider already reporting text_tokens beside reasoning_tokens is stored as sent.""" + usage = { + "input_tokens": 10, + "output_tokens": 20, + "total_tokens": 30, + "output_tokens_details": {"text_tokens": 12, "reasoning_tokens": 5}, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 12 + assert result.completion_tokens_details.reasoning_tokens == 5 + def test_transform_response_api_usage_carries_extra_provider_fields(self): """Non-standard usage fields (e.g. xAI tool details) must survive chat normalization.""" details = {"web_search_calls": 2, "x_search_calls": 0} diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 3eb051982e4..1f14fdc2c67 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4534,3 +4534,42 @@ def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_o ) assert total_cost == pytest.approx(expected) assert total_cost == pytest.approx(0.0002362) + + +def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> None: + """The combined usage that lands in spend logs keeps reasoning out of text_tokens for every turn.""" + results: OpenAIRealtimeStreamList = [ + {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}}, + { + "type": "response.done", + "response": { + "usage": { + "total_tokens": 307, + "input_tokens": 237, + "output_tokens": 70, + "input_token_details": {"text_tokens": 43, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 70, "audio_tokens": 0, "reasoning_tokens": 52}, + } + }, + }, + { + "type": "response.done", + "response": { + "usage": { + "total_tokens": 363, + "input_tokens": 300, + "output_tokens": 63, + "input_token_details": {"text_tokens": 106, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 63, "audio_tokens": 0, "reasoning_tokens": 43}, + } + }, + }, + ] + + combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) + + assert combined.completion_tokens == 133 + assert combined.completion_tokens_details is not None + assert combined.completion_tokens_details.reasoning_tokens == 95 + assert combined.completion_tokens_details.text_tokens == 38 + assert combined.completion_tokens_details.audio_tokens == 0 From 53178486a0c919481bd40da62bfa4ff98c5c2e5a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 22:59:22 -0700 Subject: [PATCH 15/49] style(tests): wrap realtime cost test lines to the 120-column limit --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 8 +++++-- tests/test_litellm/test_cost_calculator.py | 22 +++++++++++++++---- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index e7a92a96f34..e8eae73df3f 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -4800,7 +4800,9 @@ def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_loca ) -def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens(_local_model_cost_map: None) -> None: +def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens( + _local_model_cost_map: None, +) -> None: """Providers whose text_tokens exclude reasoning (text + reasoning == completion) stay billed in full.""" model = "gpt-realtime-2.1-mini" @@ -4835,4 +4837,6 @@ def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output( info = litellm.get_model_info(model=model, custom_llm_provider="openai") assert breakdown.reasoning_cost == pytest.approx(20 * info["output_cost_per_token"]) - assert completion_cost == pytest.approx(30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"]) + assert completion_cost == pytest.approx( + 30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"] + ) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 1f14fdc2c67..57246ec08a5 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4494,7 +4494,9 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m assert completion_cost == pytest.approx(500 * 2.5e-5) -def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once(_local_model_cost_map: None) -> None: +def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once( + _local_model_cost_map: None, +) -> None: """Realtime response.done nests reasoning_tokens inside text_tokens, so they are billed once.""" results: OpenAIRealtimeStreamList = [ {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}}, @@ -4530,7 +4532,9 @@ def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_o info = litellm.get_model_info(model="azure/gpt-realtime-2.1-mini", custom_llm_provider="azure") expected = ( - 43 * info["input_cost_per_token"] + 194 * info["input_cost_per_image_token"] + 23 * info["output_cost_per_token"] + 43 * info["input_cost_per_token"] + + 194 * info["input_cost_per_image_token"] + + 23 * info["output_cost_per_token"] ) assert total_cost == pytest.approx(expected) assert total_cost == pytest.approx(0.0002362) @@ -4547,7 +4551,12 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> "total_tokens": 307, "input_tokens": 237, "output_tokens": 70, - "input_token_details": {"text_tokens": 43, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0}, + "input_token_details": { + "text_tokens": 43, + "audio_tokens": 0, + "image_tokens": 194, + "cached_tokens": 0, + }, "output_token_details": {"text_tokens": 70, "audio_tokens": 0, "reasoning_tokens": 52}, } }, @@ -4559,7 +4568,12 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> "total_tokens": 363, "input_tokens": 300, "output_tokens": 63, - "input_token_details": {"text_tokens": 106, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0}, + "input_token_details": { + "text_tokens": 106, + "audio_tokens": 0, + "image_tokens": 194, + "cached_tokens": 0, + }, "output_token_details": {"text_tokens": 63, "audio_tokens": 0, "reasoning_tokens": 43}, } }, From 0ee3bec046bf0f15eebdddc1d52610d389ca220f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:20:21 -0700 Subject: [PATCH 16/49] fix(image_edit): await an async transform hook and hide the URL policy verdict from callers The image edit handler now awaits BaseImageEditConfig.async_transform_image_edit_request, and Black Forest Labs overrides it so URL images and masks download through async_safe_get instead of the blocking safe_get on the event loop. Rejected image fetches raise a fixed policy message with the user_url_allowed_hosts hint rather than echoing the resolver's verdict (resolved IP, DNS failure) back to the caller. The test fixture also fails any request-path call of the sync convert_url_to_base64 so a regression cannot pass unnoticed. --- .../prompt_templates/image_handling.py | 12 +- .../base_llm/image_edit/transformation.py | 19 ++ .../image_edit/transformation.py | 102 ++++++----- litellm/llms/custom_httpx/llm_http_handler.py | 2 +- tests/test_litellm/conftest.py | 10 +- .../litellm_core_utils/test_image_handling.py | 65 ++++--- .../test_bfl_image_edit_transformation.py | 78 ++++++++- .../custom_httpx/test_llm_http_handler.py | 163 +++++++++++++++++- 8 files changed, 379 insertions(+), 72 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index dcad7776b58..4e7dfdb4017 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -78,6 +78,14 @@ def _process_image_response(response: Response, url: str) -> str: return result +def _url_policy_rejection(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": + verbose_logger.warning("Image fetch of %s rejected by the URL policy: %s", url, verdict) + return litellm.ImageFetchError( + "Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; " + f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}" + ) + + async def async_convert_url_to_base64(url: str) -> str: if url.startswith("data:") and ";base64," in url: return url @@ -100,7 +108,7 @@ async def async_convert_url_to_base64(url: str) -> str: except litellm.ImageFetchError: raise except SSRFError as e: - raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e + raise _url_policy_rejection(url, e) from e except Exception: pass raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}") @@ -128,7 +136,7 @@ def convert_url_to_base64(url: str) -> str: except litellm.ImageFetchError: raise except SSRFError as e: - raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e + raise _url_policy_rejection(url, e) from e except Exception as e: verbose_logger.exception(e) raise litellm.ImageFetchError( diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index 9a25d3294e0..4faf0aaaf30 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/litellm/llms/base_llm/image_edit/transformation.py @@ -1,5 +1,6 @@ import types from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx @@ -102,6 +103,24 @@ class BaseImageEditConfig(ABC): ) -> tuple[dict, RequestFiles]: pass + async def async_transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: FileTypes | None, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[dict, RequestFiles]: + return self.transform_image_edit_request( + model=model, + prompt=prompt, + image=image, + image_edit_optional_request_params=dict(image_edit_optional_request_params), + litellm_params=litellm_params, + headers=dict(headers), + ) + def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict: """ Last pass on the request dict after ``transform_image_edit_request``, using the diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 013053e5bd5..c6de9f9ba65 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -9,6 +9,7 @@ API Reference: https://docs.bfl.ai/ import base64 import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -16,7 +17,7 @@ from httpx._types import RequestFiles import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH -from litellm.litellm_core_utils.url_utils import safe_get +from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -37,6 +38,22 @@ else: LiteLLMLoggingObj = Any +_BFL_REQUEST_PARAMS: Final = ( + "seed", + "output_format", + "safety_tolerance", + "prompt_upsampling", + "aspect_ratio", + "steps", + "guidance", + "grow_mask", + "top", + "bottom", + "left", + "right", +) + + class BlackForestLabsImageEditConfig(BaseImageEditConfig): """ Configuration for Black Forest Labs image editing. @@ -85,34 +102,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): BFL-specific params are passed through directly. """ optional_params: Final[dict[str, Any]] = {} - - # Pass through BFL-specific params - bfl_params: Final = [ - "seed", - "output_format", - "safety_tolerance", - "prompt_upsampling", - # Kontext-specific - "aspect_ratio", - # Fill/Inpaint-specific - "steps", - "guidance", - "grow_mask", - # Expand-specific - "top", - "bottom", - "left", - "right", - ] - - # Convert TypedDict to regular dict for access - params_dict: Final = dict(image_edit_optional_params) - - for param in bfl_params: - if param in params_dict: - value = params_dict[param] - if value is not None: - optional_params[param] = value + params: Final[Mapping[str, object]] = image_edit_optional_params + for param in _BFL_REQUEST_PARAMS: + if (value := params.get(param)) is not None: + optional_params[param] = value # Set default output format if "output_format" not in optional_params: @@ -251,23 +244,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): "input_image": b64_image, } - # Add optional params (only BFL-recognized parameters) - bfl_request_params: Final = [ - "seed", - "output_format", - "safety_tolerance", - "prompt_upsampling", - "aspect_ratio", - "steps", - "guidance", - "grow_mask", - "top", - "bottom", - "left", - "right", - ] for key, value in image_edit_optional_request_params.items(): - if key in bfl_request_params and value is not None: + if key in _BFL_REQUEST_PARAMS and value is not None: request_body[key] = value # Handle mask if provided (for inpainting) @@ -277,7 +255,39 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8") # BFL uses JSON, not multipart - return empty files - return request_body, [] + return request_body, () + + async def async_transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: FileTypes | None, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[dict, RequestFiles]: + downloaded_image: Final = await self._fetch_remote_image(image) + downloaded_mask: Final = await self._fetch_remote_image(image_edit_optional_request_params.get("mask")) + return self.transform_image_edit_request( + model=model, + prompt=prompt, + image=image if downloaded_image is None else downloaded_image, + image_edit_optional_request_params=( + dict(image_edit_optional_request_params) + if downloaded_mask is None + else {**image_edit_optional_request_params, "mask": downloaded_mask} + ), + litellm_params=litellm_params, + headers=dict(headers), + ) + + async def _fetch_remote_image(self, image: object) -> bytes | None: + candidate: Final = image[0] if isinstance(image, list) and image else image + if not isinstance(candidate, str) or not candidate.startswith(("http://", "https://")): + return None + response: Final = await async_safe_get(litellm.module_level_aclient, candidate, timeout=60.0) + response.raise_for_status() + return response.content def transform_image_edit_response( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d57aee025d5..6a17841faa2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6787,7 +6787,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - data, files = image_edit_provider_config.transform_image_edit_request( + data, files = await image_edit_provider_config.async_transform_image_edit_request( model=model, image=image, prompt=prompt, diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index aa8ba168ecc..a4f32df46ae 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -607,7 +607,8 @@ ONE_PIXEL_PNG = base64.b64decode( @pytest.fixture def async_only_image_fetch(monkeypatch): - from litellm.litellm_core_utils.prompt_templates import image_handling + from litellm.litellm_core_utils.prompt_templates import factory, image_handling + from litellm.llms.gemini.chat import transformation as gemini_chat_transformation fetch = SimpleNamespace( fetched=[], @@ -627,6 +628,13 @@ def async_only_image_fetch(monkeypatch): request=httpx.Request("GET", url), ) + def forbid_sync_convert(url, *args, **kwargs): + if url.startswith(("http://", "https://")): + raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}") + return url + monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch) monkeypatch.setattr(image_handling, "async_safe_get", serve_png) + for module in (image_handling, factory, gemini_chat_transformation): + monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert) return fetch diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index e7b28d0d3a2..d8ccc9251dd 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -428,36 +428,61 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail assert time.perf_counter() - started < 1 -async def test_async_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch): +_SSRF_VERDICTS = ( + "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " + "add the host to `user_url_allowed_hosts` in general_settings.", + "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known", + "No addresses found for 'internal.example'", +) + + +def _assert_one_verdict_free_message(messages, url): + assert len(set(messages)) == 1 + assert "10.0.0.8" not in messages[0] + assert "DNS" not in messages[0] + assert "No addresses" not in messages[0] + assert "user_url_allowed_hosts" in messages[0] + assert url in messages[0] + + +async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): attempts = [] - - async def block(client, url, **kwargs): - attempts.append(url) - raise SSRFError("URL targets a blocked address (10.0.0.8)") - - monkeypatch.setattr(image_handling, "async_safe_get", block) + messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"): - await async_convert_url_to_base64(url) + for verdict in _SSRF_VERDICTS: - assert attempts == [url] + async def block(client, fetched_url, verdict=verdict, **kwargs): + attempts.append(fetched_url) + raise SSRFError(verdict) + + monkeypatch.setattr(image_handling, "async_safe_get", block) + with pytest.raises(litellm.ImageFetchError) as raised: + await async_convert_url_to_base64(url) + messages.append(raised.value.message) + + assert attempts == [url] * len(_SSRF_VERDICTS) + _assert_one_verdict_free_message(messages, url) -def test_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch): +def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): attempts = [] - - def block(client, url, **kwargs): - attempts.append(url) - raise SSRFError("URL targets a blocked address (10.0.0.8)") - - monkeypatch.setattr(image_handling, "safe_get", block) + messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"): - convert_url_to_base64(url) + for verdict in _SSRF_VERDICTS: - assert attempts == [url] + def block(client, fetched_url, verdict=verdict, **kwargs): + attempts.append(fetched_url) + raise SSRFError(verdict) + + monkeypatch.setattr(image_handling, "safe_get", block) + with pytest.raises(litellm.ImageFetchError) as raised: + convert_url_to_base64(url) + messages.append(raised.value.message) + + assert attempts == [url] * len(_SSRF_VERDICTS) + _assert_one_verdict_free_message(messages, url) async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch): diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py index ec243b7058d..d9d6e813d86 100644 --- a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py +++ b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py @@ -16,6 +16,9 @@ import httpx import pytest +from litellm.llms.black_forest_labs.image_edit import ( + transformation as bfl_transformation, +) from litellm.llms.black_forest_labs.image_edit.transformation import ( BlackForestLabsImageEditConfig, ) @@ -186,7 +189,7 @@ class TestBlackForestLabsImageEditTransformation: assert data["output_format"] == "jpeg" # BFL uses JSON, not multipart - files should be empty - assert files == [] + assert files == () def test_transform_image_edit_request_with_mask(self): """Test request transformation with mask for inpainting.""" @@ -299,3 +302,76 @@ class TestBlackForestLabsImageEditTransformation: def test_use_multipart_form_data_returns_false(self): """Test that use_multipart_form_data returns False for BFL.""" assert self.config.use_multipart_form_data() is False + + +async def test_async_transform_image_edit_request_downloads_url_images_with_the_async_fetcher(monkeypatch): + served = b"png-bytes-from-cdn" + fetched = [] + + def forbid_sync_fetch(client, url, **kwargs): + raise AssertionError(f"sync image fetch ran on the event loop: {url}") + + async def serve(client, url, **kwargs): + fetched.append((url, kwargs.get("timeout"))) + return httpx.Response(200, content=served, request=httpx.Request("GET", url)) + + monkeypatch.setattr(bfl_transformation, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(bfl_transformation, "async_safe_get", serve) + + data, files = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image="https://cdn.example/photo.png", + image_edit_optional_request_params={"mask": "https://cdn.example/mask.png", "seed": 7}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert base64.b64decode(data["input_image"]) == served + assert base64.b64decode(data["mask"]) == served + assert data["seed"] == 7 + assert files == () + assert fetched == [("https://cdn.example/photo.png", 60.0), ("https://cdn.example/mask.png", 60.0)] + + +async def test_async_transform_image_edit_request_never_fetches_for_local_images(monkeypatch): + def refuse(*args, **kwargs): + raise AssertionError("no network fetch expected for local image bytes") + + monkeypatch.setattr(bfl_transformation, "safe_get", refuse) + monkeypatch.setattr(bfl_transformation, "async_safe_get", refuse) + + data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image=[BytesIO(b"first"), BytesIO(b"other")], + image_edit_optional_request_params={"mask": b"mask-bytes"}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert base64.b64decode(data["input_image"]) == b"first" + assert base64.b64decode(data["mask"]) == b"mask-bytes" + + +async def test_async_transform_image_edit_request_downloads_only_the_first_url_of_a_list(monkeypatch): + fetched = [] + + async def serve(client, url, **kwargs): + fetched.append(url) + return httpx.Response(200, content=b"first-bytes", request=httpx.Request("GET", url)) + + monkeypatch.setattr(bfl_transformation, "safe_get", lambda *args, **kwargs: pytest.fail("sync fetch ran")) + monkeypatch.setattr(bfl_transformation, "async_safe_get", serve) + + data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image=["https://cdn.example/a.png", "https://cdn.example/b.png"], + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert fetched == ["https://cdn.example/a.png"] + assert base64.b64decode(data["input_image"]) == b"first-bytes" diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index f1614654ffb..3afd57bbf34 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -19,6 +19,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, @@ -31,7 +32,7 @@ from litellm.llms.azure.videos.transformation import AzureVideoConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import ModelResponse, TranscriptionResponse +from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" @@ -3244,6 +3245,11 @@ class _TransformRecordingConfig(BaseConfig): def get_error_class(self, error_message, status_code, headers): return BaseLLMException(status_code=status_code, message=error_message, headers=headers) + def get_model_response_iterator(self, streaming_response, sync_stream, json_mode=False): + return litellm.OpenAIGPTConfig().get_model_response_iterator( + streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode + ) + def _start_async_completion(config, logging_obj=None): captured = {} @@ -3312,3 +3318,158 @@ async def test_completion_keeps_sync_transform_request_before_returning_by_defau assert config.transform_calls == ["sync"] assert captured["body"] == {"transformed_by": "sync"} assert response.choices[0].message.content == "sync" + + +def _sse_echoing_transformed_by(request): + transformed_by = json.loads(request.content)["transformed_by"] + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "stub-model", + "choices": [{"index": 0, "delta": {"content": transformed_by}, "finish_reason": None}], + } + return httpx.Response( + 200, + content=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode(), + headers={"content-type": "text/event-stream"}, + request=request, + ) + + +def _streaming_logging_obj(): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="async-transform-stream", + function_id="f", + ) + logging_obj.update_environment_variables( + model="stub-model", user="", optional_params={}, litellm_params={}, custom_llm_provider="openai" + ) + return logging_obj + + +async def test_completion_streams_after_the_async_transform_request(): + config = _TransformRecordingConfig(transform_async=True) + loop_thread = threading.current_thread() + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(_sse_echoing_transformed_by)) + + stream = await BaseLLMHTTPHandler().completion( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + api_base="https://llm.example/v1/chat", + custom_llm_provider="openai", + model_response=ModelResponse(), + encoding=None, + logging_obj=_streaming_logging_obj(), + optional_params={}, + timeout=10.0, + litellm_params={}, + acompletion=True, + stream=True, + client=client, + provider_config=config, + ) + collected = [chunk async for chunk in stream] + + assert config.transform_calls == ["async"] + assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) + assert "".join(chunk.choices[0].delta.content or "" for chunk in collected) == "async" + + +class _ImageEditRecordingConfig(BaseImageEditConfig): + def __init__(self): + self.transform_calls = [] + + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, image_edit_optional_params, model, drop_params): + return dict(image_edit_optional_params) + + def validate_environment(self, headers, model, api_key=None, litellm_params=None, api_base=None): + return {} + + def get_complete_url(self, model, api_base, litellm_params): + return "https://images.example/v1/edits" + + def use_multipart_form_data(self): + return False + + def transform_image_edit_request( + self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers + ): + self.transform_calls.append("sync") + return {"transformed_by": "sync"}, [] + + async def async_transform_image_edit_request( + self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers + ): + self.transform_calls.append("async") + return {"transformed_by": "async"}, [] + + def transform_image_edit_response(self, model, raw_response, logging_obj): + return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["transformed_by"])]) + + +def _echo_json_transport(captured): + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=captured["body"]) + + return httpx.MockTransport(handle) + + +async def test_async_image_edit_handler_awaits_the_async_transform(): + config = _ImageEditRecordingConfig() + captured = {} + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_echo_json_transport(captured)) + + response = await BaseLLMHTTPHandler().async_image_edit_handler( + model="edit-model", + image=b"raw-image", + prompt="add a hat", + image_edit_provider_config=config, + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert config.transform_calls == ["async"] + assert captured["body"] == {"transformed_by": "async"} + assert response.data[0].b64_json == "async" + + +def test_image_edit_handler_keeps_the_sync_transform(): + config = _ImageEditRecordingConfig() + captured = {} + client = HTTPHandler() + client.client = httpx.Client(transport=_echo_json_transport(captured)) + + response = BaseLLMHTTPHandler().image_edit_handler( + model="edit-model", + image=b"raw-image", + prompt="add a hat", + image_edit_provider_config=config, + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert config.transform_calls == ["sync"] + assert captured["body"] == {"transformed_by": "sync"} + assert response.data[0].b64_json == "sync" From 3920cf4dfe458b6d4982e2a64c162de21defa2a0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:45:37 -0700 Subject: [PATCH 17/49] fix(image_handling): tell callers when the image host did not resolve instead of blaming the URL policy validate_url raises HostResolutionError, a SSRFError subclass, for the two DNS outcomes (lookup failed, no addresses). The image fetch helper maps that to a "host could not be resolved" message and keeps the user_url_allowed_hosts hint for the policy verdicts it can actually fix. --- .../prompt_templates/image_handling.py | 14 +++-- litellm/litellm_core_utils/url_utils.py | 8 ++- .../litellm_core_utils/test_image_handling.py | 56 +++++++++++-------- .../litellm_core_utils/test_url_utils.py | 19 ++++++- 4 files changed, 64 insertions(+), 33 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 4e7dfdb4017..99efbd13320 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -15,7 +15,7 @@ import litellm from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB -from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError, async_safe_get, safe_get from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 @@ -78,8 +78,12 @@ def _process_image_response(response: Response, url: str) -> str: return result -def _url_policy_rejection(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": - verbose_logger.warning("Image fetch of %s rejected by the URL policy: %s", url, verdict) +def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": + verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict) + if isinstance(verdict, HostResolutionError): + return litellm.ImageFetchError( + f"Error: Unable to fetch image from URL. The image host could not be resolved. url={url}" + ) return litellm.ImageFetchError( "Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; " f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}" @@ -108,7 +112,7 @@ async def async_convert_url_to_base64(url: str) -> str: except litellm.ImageFetchError: raise except SSRFError as e: - raise _url_policy_rejection(url, e) from e + raise _rejected_image_fetch(url, e) from e except Exception: pass raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}") @@ -136,7 +140,7 @@ def convert_url_to_base64(url: str) -> str: except litellm.ImageFetchError: raise except SSRFError as e: - raise _url_policy_rejection(url, e) from e + raise _rejected_image_fetch(url, e) from e except Exception as e: verbose_logger.exception(e) raise litellm.ImageFetchError( diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index fa070a648f5..f4388d5fd12 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -93,6 +93,10 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" +class HostResolutionError(SSRFError): + pass + + def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. @@ -324,10 +328,10 @@ def validate_url(url: str) -> tuple[str, str]: try: addrinfo: Final = socket.getaddrinfo(hostname, effective_port, proto=socket.IPPROTO_TCP) except socket.gaierror as e: - raise SSRFError(f"DNS resolution failed for '{hostname}': {e}") + raise HostResolutionError(f"DNS resolution failed for '{hostname}': {e}") if not addrinfo: - raise SSRFError(f"No addresses found for '{hostname}'") + raise HostResolutionError(f"No addresses found for '{hostname}'") if not is_allowlisted: for addrinfo_entry in addrinfo: diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index d8ccc9251dd..3a909298aba 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -17,7 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, convert_url_to_base64, ) -from litellm.litellm_core_utils.url_utils import SSRFError +from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError @pytest.fixture(autouse=True) @@ -115,9 +115,7 @@ class StreamingLargeImageClient: request=Request("GET", url), ) # Mock the iter_bytes method to return our generator - response.iter_bytes = lambda chunk_size=8192: generate_chunks( - size_bytes, chunk_size - ) + response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size) return response @@ -215,9 +213,7 @@ def test_streaming_download_handles_petabyte_file(monkeypatch): """ # Simulate a 1 petabyte file (1,000,000 GB) # Without streaming protection, this would cause OOM or hang indefinitely - client = StreamingLargeImageClient( - size_mb=1_000_000_000, include_content_length=False - ) + client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False) monkeypatch.setattr(litellm, "module_level_client", client) with pytest.raises(litellm.ImageFetchError) as excinfo: @@ -429,20 +425,32 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail _SSRF_VERDICTS = ( - "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " - "add the host to `user_url_allowed_hosts` in general_settings.", - "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known", - "No addresses found for 'internal.example'", + ( + SSRFError( + "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " + "add the host to `user_url_allowed_hosts` in general_settings." + ), + "The proxy's URL policy rejected this host; an admin can allow it with `user_url_allowed_hosts`", + ), + ( + HostResolutionError( + "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known" + ), + "The image host could not be resolved", + ), + (HostResolutionError("No addresses found for 'internal.example'"), "The image host could not be resolved"), ) -def _assert_one_verdict_free_message(messages, url): - assert len(set(messages)) == 1 - assert "10.0.0.8" not in messages[0] - assert "DNS" not in messages[0] - assert "No addresses" not in messages[0] - assert "user_url_allowed_hosts" in messages[0] - assert url in messages[0] +def _assert_verdict_free_messages(messages, url): + for message, (verdict, expected_guidance) in zip(messages, _SSRF_VERDICTS, strict=True): + assert expected_guidance in message + assert url in message + assert "10.0.0.8" not in message + assert "DNS" not in message + assert "No addresses" not in message + if isinstance(verdict, HostResolutionError): + assert "user_url_allowed_hosts" not in message async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): @@ -450,11 +458,11 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - for verdict in _SSRF_VERDICTS: + for verdict, _ in _SSRF_VERDICTS: async def block(client, fetched_url, verdict=verdict, **kwargs): attempts.append(fetched_url) - raise SSRFError(verdict) + raise verdict monkeypatch.setattr(image_handling, "async_safe_get", block) with pytest.raises(litellm.ImageFetchError) as raised: @@ -462,7 +470,7 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r messages.append(raised.value.message) assert attempts == [url] * len(_SSRF_VERDICTS) - _assert_one_verdict_free_message(messages, url) + _assert_verdict_free_messages(messages, url) def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): @@ -470,11 +478,11 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - for verdict in _SSRF_VERDICTS: + for verdict, _ in _SSRF_VERDICTS: def block(client, fetched_url, verdict=verdict, **kwargs): attempts.append(fetched_url) - raise SSRFError(verdict) + raise verdict monkeypatch.setattr(image_handling, "safe_get", block) with pytest.raises(litellm.ImageFetchError) as raised: @@ -482,7 +490,7 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp messages.append(raised.value.message) assert attempts == [url] * len(_SSRF_VERDICTS) - _assert_one_verdict_free_message(messages, url) + _assert_verdict_free_messages(messages, url) async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch): diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index fccdc1a2a0a..151c7ceb5bd 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.litellm_core_utils import url_utils from litellm.litellm_core_utils.url_utils import ( + HostResolutionError, SSRFError, _is_blocked_ip, assert_same_origin, @@ -161,10 +162,24 @@ class TestValidateUrl: assert "/path" in rewritten assert "key=value" in rewritten - def test_dns_failure_raises(self, mock_dns_failure): - with pytest.raises(SSRFError, match="DNS resolution failed"): + def test_dns_failure_raises_a_host_resolution_error(self, mock_dns_failure): + with pytest.raises(HostResolutionError, match="DNS resolution failed"): validate_url("http://this-domain-does-not-exist-xyz123.invalid/test") + def test_empty_resolution_raises_a_host_resolution_error(self, monkeypatch): + monkeypatch.setattr(url_utils.socket, "getaddrinfo", lambda *args, **kwargs: []) + with pytest.raises(HostResolutionError, match="No addresses found"): + validate_url("http://this-domain-resolves-to-nothing.invalid/test") + + def test_blocked_address_is_not_a_host_resolution_error(self, monkeypatch): + def fake(host, port, *a, **kw): + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.8", port or 80))] + + monkeypatch.setattr(url_utils.socket, "getaddrinfo", fake) + with pytest.raises(SSRFError, match="blocked address") as raised: + validate_url("http://internal.example/test") + assert not isinstance(raised.value, HostResolutionError) + def test_blocks_localhost_hostname(self, monkeypatch): def fake(host, port, *a, **kw): return [ From de2ba3fab1daab25ba2517ea1a93cf7120f996f5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:50:01 -0700 Subject: [PATCH 18/49] fix(make check): lint the test tree on tests-only changes like CI does CI's required lint job runs ruff with ruff-tests.toml over tests/ and the test-quality budget gate, but scripts/pre_commit_lint.sh only triggered make lint on litellm/ files, so a tests-only commit passed make check with a no-op note and then failed CI (a duplicate test name, ruff F811, did exactly that). When tests/ Python files are in scope and no litellm/ files are, run ruff with ruff-tests.toml over the changed test files and make lint-test-quality, with the matching partial-staging warning, summary line, and no-op condition. --- scripts/pre_commit_lint.sh | 25 +++++- tests/test_litellm/test_pre_commit_lint.py | 98 ++++++++++++++++++++-- 2 files changed, 114 insertions(+), 9 deletions(-) diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index d38a0eee3de..ad1b4e793ab 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -12,6 +12,8 @@ # - litellm/ Python -> `make lint` (test-linting.yml's lint job) # - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) # + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) +# - tests/ Python -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's +# test-tree ruff and test-quality budget steps) # - dashboard -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) # - proxy/types -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (check-ui-api-types.yml) # @@ -88,15 +90,18 @@ existing_files() { litellm_py_pattern='^litellm/.*\.py$' e2e_py_pattern='^tests/e2e/.*\.py$' +tests_py_pattern='^tests/.*\.py$' spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$' ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' -# CI's lint job (test-linting.yml) only inspects litellm/, so a tests-only or -# scripts-only commit can't turn it red; scope the trigger there to skip the slow -# make lint when it couldn't catch anything. +# CI's lint job (test-linting.yml) is make lint: litellm/ checks plus two test-tree +# steps (ruff over ruff-tests.toml and the test-quality budget). Trigger the slow +# make lint on litellm/ files only; a tests-only commit runs just those two steps. litellm_py_files=$(scope_match "$litellm_py_pattern") e2e_py_files=$(scope_match "$e2e_py_pattern") +tests_py_changed=$(scope_match "$tests_py_pattern") +tests_py_files=$(printf '%s\n' "$tests_py_changed" | existing_files) # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it. fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files) # check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types @@ -136,6 +141,7 @@ if [ -n "$staged" ]; then } warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files" warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files" + warn_skipped "test-tree lint (ruff-tests.toml + test-quality budget)" "$tests_py_pattern" "$tests_py_changed" warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed" warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files" fi @@ -225,6 +231,16 @@ if [ -n "$e2e_py_files" ]; then || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make check." >&2; status=1; } fi +if [ -n "$tests_py_changed" ] && [ -z "$litellm_py_files" ]; then + if [ -n "$tests_py_files" ]; then + echo "check: linting the test tree (ruff check --config ruff-tests.toml, scoped tests files)" + printf '%s\n' "$tests_py_files" | xargs uv run --no-sync ruff check --config ruff-tests.toml \ + || { echo "✗ Test-tree ruff failed. Fix the errors above, then re-run make check." >&2; status=1; } + fi + echo "check: checking the test-quality budget (make lint-test-quality)" + make lint-test-quality || { echo "✗ Test-quality budget failed. Fix the errors above, then re-run make check." >&2; status=1; } +fi + dashboard_checks() { echo "check: linting dashboard (prettier + eslint + lint budgets)" if [ ! -d ui/litellm-dashboard/node_modules ]; then @@ -313,10 +329,11 @@ summary_item() { echo "check: summary" summary_item "Python lint (make lint)" "$litellm_py_files" "no litellm/ Python files in scope" summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_files" "no tests/e2e Python files in scope" +summary_item "test-tree lint (ruff-tests.toml + test-quality budget)" "$tests_py_changed" "no tests/ Python files in scope" summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope" summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope" -if [ -z "$litellm_py_files$e2e_py_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then +if [ -z "$litellm_py_files$e2e_py_files$tests_py_changed$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2 printf '%s\n' "$scope" | sed 's/^/ /' >&2 echo " A pass here is a no-op, not a lint verdict." >&2 diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index aa2260e89ea..22ae0e1bfaa 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -40,6 +40,9 @@ case "$*" in sleep 60 fi ;; + lint-test-quality) + [ "${STUB_FAIL:-}" = "test-quality" ] && exit 1 + ;; esac exit 0 """ @@ -69,6 +72,10 @@ case "$*" in *orjson*) [ -n "${STUB_BARRIER_DIR:-}" ] && barrier_sync genapi "python dashboard" ;; + "run --no-sync ruff check --config ruff-tests.toml"*) + [ -n "${STUB_ARGS_DIR:-}" ] && echo "$*" >> "$STUB_ARGS_DIR/ruff_tests.args" + [ "${STUB_FAIL:-}" = "tests-ruff" ] && exit 1 + ;; esac exit 0 """ @@ -153,6 +160,13 @@ def _set_base_ref(repo: Path) -> None: ) +def _stage_file(repo: Path, relative: str, body: str) -> None: + path = repo / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(body) + subprocess.run(["git", "add", relative], cwd=repo, check=True) + + def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") @@ -405,6 +419,7 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No assert "ran: dashboard lint (prettier + eslint + lint budgets)" in proc.stdout assert "ran: dashboard API-type sync (npm run gen:api)" in proc.stdout assert "skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope)" in proc.stdout + assert "skipped: test-tree lint (ruff-tests.toml + test-quality budget) (no tests/ Python files in scope)" in proc.stdout assert "check: PASS" in proc.stdout assert "check: FAIL" not in proc.stdout @@ -412,20 +427,93 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty_log(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") - tests_dir = repo / "tests" / "test_litellm" - tests_dir.mkdir(parents=True) - (tests_dir / "test_x.py").write_text("def test_x() -> None: ...\n") - subprocess.run(["git", "add", "tests"], cwd=repo, check=True) + _stage_file(repo, "scripts/tool.py", "def main() -> None: ...\n") proc = _run(repo, bin_dir, {}) assert proc.returncode == 0, proc.stdout + proc.stderr assert "no gating lint check matches the files in scope, so nothing ran" in proc.stdout - assert "tests/test_litellm/test_x.py" in proc.stdout + assert "scripts/tool.py" in proc.stdout assert "a no-op, not a lint verdict" in proc.stdout assert "check: PASS" in proc.stdout assert "linting Python" not in proc.stdout log = (repo / ".git" / "pre_commit_lint.log").read_text() assert "check: summary" in log assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log + assert "skipped: test-tree lint (ruff-tests.toml + test-quality budget) (no tests/ Python files in scope)" in log + + +def test_tests_only_change_runs_test_tree_ruff_on_staged_files_and_the_quality_gate(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _stage_file(repo, "tests/test_b.py", "def test_b() -> None: ...\n") + _stage_file(repo, "tests/fixtures/data.json", "{}\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + ruff_args = (args_dir / "ruff_tests.args").read_text().splitlines() + assert ruff_args == ["run --no-sync ruff check --config ruff-tests.toml tests/test_a.py tests/test_b.py"] + assert "ran: test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout + assert "no gating lint check matches" not in proc.stdout + assert "linting Python" not in proc.stdout + assert "check: PASS" in proc.stdout + + +@pytest.mark.parametrize( + ("fail", "message"), + [ + ("tests-ruff", "Test-tree ruff failed"), + ("test-quality", "Test-quality budget failed"), + ], +) +def test_a_failing_test_tree_check_fails_a_tests_only_run(tmp_path: Path, fail: str, message: str) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + proc = _run(repo, bin_dir, {"STUB_FAIL": fail}) + assert proc.returncode == 1 + assert message in proc.stdout + proc.stderr + assert "check: FAIL" in proc.stdout + + +def test_tests_changed_alongside_litellm_files_defer_to_make_lint(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "litellm/foo.py", "x = 2\n") + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "test-quality"}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "linting Python" in proc.stdout + assert not (args_dir / "ruff_tests.args").exists() + assert "ran: test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout + + +def test_deleted_test_file_still_runs_the_quality_gate_without_feeding_ruff_the_missing_file(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + subprocess.run(["git", "rm", "-q", "tests/test_a.py"], cwd=repo, check=True) + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "test-quality"}) + assert proc.returncode == 1 + assert "Test-quality budget failed" in proc.stdout + proc.stderr + assert not (args_dir / "ruff_tests.args").exists() + + +def test_partial_staging_warns_when_test_files_are_left_unstaged(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _commit_all(repo, "base") + _stage_file(repo, "notes.md", "hi\n") + (repo / "tests" / "test_a.py").write_text("def test_a() -> None:\n assert True\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "SKIPPED test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout + assert "tests/test_a.py" in proc.stdout + assert "make lint-test-quality" not in proc.stdout def test_run_queues_through_the_machine_wide_gate_slot_lock(tmp_path: Path) -> None: From d22962248c1c659279e7389e6c2e1d1b640b400a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:07:25 -0700 Subject: [PATCH 19/49] fix(check): run CI's whole-tree test ruff and widen the test-tree trigger The scoped xargs list missed a ruff-tests.toml rule change and skipped ruff on deletions, so the block now runs test-linting.yml's exact command over tests/. ruff-tests.toml, test-quality-budget.json, and scripts/check_test_quality.py trigger the block too, and it sits after the background launches so the dashboard and gen:api jobs overlap it. --- scripts/pre_commit_lint.sh | 36 ++++---- tests/test_litellm/test_pre_commit_lint.py | 102 ++++++++++++++++----- 2 files changed, 96 insertions(+), 42 deletions(-) diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index ad1b4e793ab..ab84d2d518a 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -12,7 +12,8 @@ # - litellm/ Python -> `make lint` (test-linting.yml's lint job) # - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) # + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) -# - tests/ Python -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's +# - tests/ Python, ruff-tests.toml, test-quality-budget.json, scripts/check_test_quality.py +# -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's # test-tree ruff and test-quality budget steps) # - dashboard -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) # - proxy/types -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (check-ui-api-types.yml) @@ -90,18 +91,17 @@ existing_files() { litellm_py_pattern='^litellm/.*\.py$' e2e_py_pattern='^tests/e2e/.*\.py$' -tests_py_pattern='^tests/.*\.py$' +test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|test-quality-budget\.json|scripts/check_test_quality\.py)$' spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$' ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' # CI's lint job (test-linting.yml) is make lint: litellm/ checks plus two test-tree # steps (ruff over ruff-tests.toml and the test-quality budget). Trigger the slow -# make lint on litellm/ files only; a tests-only commit runs just those two steps. +# make lint on litellm/ files only; without them, the test-tree steps run on their own below. litellm_py_files=$(scope_match "$litellm_py_pattern") e2e_py_files=$(scope_match "$e2e_py_pattern") -tests_py_changed=$(scope_match "$tests_py_pattern") -tests_py_files=$(printf '%s\n' "$tests_py_changed" | existing_files) +test_tree_files=$(scope_match "$test_tree_pattern") # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it. fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files) # check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types @@ -141,7 +141,7 @@ if [ -n "$staged" ]; then } warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files" warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files" - warn_skipped "test-tree lint (ruff-tests.toml + test-quality budget)" "$tests_py_pattern" "$tests_py_changed" + warn_skipped "test-tree lint (ruff-tests.toml + test-quality budget)" "$test_tree_pattern" "$test_tree_files" warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed" warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files" fi @@ -231,16 +231,6 @@ if [ -n "$e2e_py_files" ]; then || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make check." >&2; status=1; } fi -if [ -n "$tests_py_changed" ] && [ -z "$litellm_py_files" ]; then - if [ -n "$tests_py_files" ]; then - echo "check: linting the test tree (ruff check --config ruff-tests.toml, scoped tests files)" - printf '%s\n' "$tests_py_files" | xargs uv run --no-sync ruff check --config ruff-tests.toml \ - || { echo "✗ Test-tree ruff failed. Fix the errors above, then re-run make check." >&2; status=1; } - fi - echo "check: checking the test-quality budget (make lint-test-quality)" - make lint-test-quality || { echo "✗ Test-quality budget failed. Fix the errors above, then re-run make check." >&2; status=1; } -fi - dashboard_checks() { echo "check: linting dashboard (prettier + eslint + lint budgets)" if [ ! -d ui/litellm-dashboard/node_modules ]; then @@ -304,6 +294,15 @@ if [ -n "$spec_files" ]; then set +m fi +if [ -n "$test_tree_files" ] && [ -z "$litellm_py_files" ]; then + echo "check: linting the test tree (ruff check --config ruff-tests.toml tests)" + uv run --no-sync ruff check --config ruff-tests.toml tests \ + || { echo "✗ Test-tree ruff failed. Fix the errors above, then re-run make check." >&2; status=1; } + echo "check: checking the test-quality budget (make lint-test-quality)" + make lint-test-quality \ + || { echo "✗ Test-quality budget failed. Fix the errors above, then re-run make check." >&2; status=1; } +fi + if [ -n "${python_pid:-}" ]; then wait "$python_pid" || status=1 cat "$python_log"; rm -f "$python_log" @@ -329,11 +328,12 @@ summary_item() { echo "check: summary" summary_item "Python lint (make lint)" "$litellm_py_files" "no litellm/ Python files in scope" summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_files" "no tests/e2e Python files in scope" -summary_item "test-tree lint (ruff-tests.toml + test-quality budget)" "$tests_py_changed" "no tests/ Python files in scope" +summary_item "test-tree lint (ruff-tests.toml + test-quality budget)" "$test_tree_files" \ + "no tests/ Python files or test-tree lint inputs in scope" summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope" summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope" -if [ -z "$litellm_py_files$e2e_py_files$tests_py_changed$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then +if [ -z "$litellm_py_files$e2e_py_files$test_tree_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2 printf '%s\n' "$scope" | sed 's/^/ /' >&2 echo " A pass here is a no-op, not a lint verdict." >&2 diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index 22ae0e1bfaa..e17abf2fa9f 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -11,6 +11,12 @@ import pytest ROOT = Path(__file__).resolve().parents[2] SCRIPT = ROOT / "scripts" / "pre_commit_lint.sh" +WHOLE_TREE_RUFF = "run --no-sync ruff check --config ruff-tests.toml tests" +TEST_TREE_RAN = "ran: test-tree lint (ruff-tests.toml + test-quality budget)" +TEST_TREE_SKIPPED = ( + "skipped: test-tree lint (ruff-tests.toml + test-quality budget) " + "(no tests/ Python files or test-tree lint inputs in scope)" +) BARRIER_HELPER = """barrier_sync() { touch "$STUB_BARRIER_DIR/$1.started" @@ -30,6 +36,7 @@ BARRIER_HELPER = """barrier_sync() { MAKE_STUB = """#!/bin/sh . "$STUB_BIN/barrier.sh" +[ -n "${STUB_ARGS_DIR:-}" ] && echo "$*" >> "$STUB_ARGS_DIR/make.args" case "$*" in lint) [ "${STUB_FAIL:-}" = "make-lint" ] && exit 1 @@ -419,7 +426,7 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No assert "ran: dashboard lint (prettier + eslint + lint budgets)" in proc.stdout assert "ran: dashboard API-type sync (npm run gen:api)" in proc.stdout assert "skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope)" in proc.stdout - assert "skipped: test-tree lint (ruff-tests.toml + test-quality budget) (no tests/ Python files in scope)" in proc.stdout + assert TEST_TREE_SKIPPED in proc.stdout assert "check: PASS" in proc.stdout assert "check: FAIL" not in proc.stdout @@ -438,41 +445,83 @@ def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty log = (repo / ".git" / "pre_commit_lint.log").read_text() assert "check: summary" in log assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log - assert "skipped: test-tree lint (ruff-tests.toml + test-quality budget) (no tests/ Python files in scope)" in log + assert TEST_TREE_SKIPPED in log -def test_tests_only_change_runs_test_tree_ruff_on_staged_files_and_the_quality_gate(tmp_path: Path) -> None: +def _recorded(args_dir: Path, name: str) -> list[str]: + path = args_dir / name + return path.read_text().splitlines() if path.exists() else [] + + +def test_tests_only_change_runs_the_whole_test_tree_ruff_and_the_quality_gate(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") args_dir = tmp_path / "args" args_dir.mkdir() _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") - _stage_file(repo, "tests/test_b.py", "def test_b() -> None: ...\n") _stage_file(repo, "tests/fixtures/data.json", "{}\n") proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) assert proc.returncode == 0, proc.stdout + proc.stderr - ruff_args = (args_dir / "ruff_tests.args").read_text().splitlines() - assert ruff_args == ["run --no-sync ruff check --config ruff-tests.toml tests/test_a.py tests/test_b.py"] - assert "ran: test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + assert TEST_TREE_RAN in proc.stdout assert "no gating lint check matches" not in proc.stdout assert "linting Python" not in proc.stdout assert "check: PASS" in proc.stdout @pytest.mark.parametrize( - ("fail", "message"), - [ - ("tests-ruff", "Test-tree ruff failed"), - ("test-quality", "Test-quality budget failed"), - ], + "changed", + ["ruff-tests.toml", "test-quality-budget.json", "scripts/check_test_quality.py", "tests/e2e/test_x.py"], ) -def test_a_failing_test_tree_check_fails_a_tests_only_run(tmp_path: Path, fail: str, message: str) -> None: +def test_test_tree_lint_inputs_trigger_the_test_tree_checks(tmp_path: Path, changed: str) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, changed, "x = 1\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert "lint-test-quality" in _recorded(args_dir, "make.args") + assert TEST_TREE_RAN in proc.stdout + + +def test_nothing_staged_tests_only_working_tree_change_runs_the_test_tree_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _commit_all(repo, "base") + _set_base_ref(repo) + args_dir = tmp_path / "args" + args_dir.mkdir() + (repo / "tests" / "test_a.py").write_text("def test_a() -> None:\n assert True\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged; scoping to the working tree's diff" in proc.stdout + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + + +def test_a_failing_test_tree_ruff_fails_the_run_and_still_runs_the_quality_gate(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "tests-ruff"}) + assert proc.returncode == 1 + assert "Test-tree ruff failed" in proc.stdout + proc.stderr + assert "check: FAIL" in proc.stdout + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + + +def test_a_failing_quality_gate_fails_a_tests_only_run(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") - proc = _run(repo, bin_dir, {"STUB_FAIL": fail}) + proc = _run(repo, bin_dir, {"STUB_FAIL": "test-quality"}) assert proc.returncode == 1 - assert message in proc.stdout + proc.stderr + assert "Test-quality budget failed" in proc.stdout + proc.stderr assert "check: FAIL" in proc.stdout @@ -486,34 +535,39 @@ def test_tests_changed_alongside_litellm_files_defer_to_make_lint(tmp_path: Path proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "test-quality"}) assert proc.returncode == 0, proc.stdout + proc.stderr assert "linting Python" in proc.stdout - assert not (args_dir / "ruff_tests.args").exists() - assert "ran: test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout + assert _recorded(args_dir, "ruff_tests.args") == [] + assert _recorded(args_dir, "make.args") == ["lint"] + assert TEST_TREE_RAN in proc.stdout -def test_deleted_test_file_still_runs_the_quality_gate_without_feeding_ruff_the_missing_file(tmp_path: Path) -> None: +def test_deleted_test_file_still_runs_the_test_tree_checks(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") _commit_all(repo, "base") args_dir = tmp_path / "args" args_dir.mkdir() subprocess.run(["git", "rm", "-q", "tests/test_a.py"], cwd=repo, check=True) - proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "test-quality"}) - assert proc.returncode == 1 - assert "Test-quality budget failed" in proc.stdout + proc.stderr - assert not (args_dir / "ruff_tests.args").exists() + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + assert TEST_TREE_RAN in proc.stdout def test_partial_staging_warns_when_test_files_are_left_unstaged(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() _stage_file(repo, "notes.md", "hi\n") (repo / "tests" / "test_a.py").write_text("def test_a() -> None:\n assert True\n") - proc = _run(repo, bin_dir, {}) + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) assert proc.returncode == 0, proc.stdout + proc.stderr assert "SKIPPED test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout assert "tests/test_a.py" in proc.stdout - assert "make lint-test-quality" not in proc.stdout + assert _recorded(args_dir, "ruff_tests.args") == [] + assert _recorded(args_dir, "make.args") == [] def test_run_queues_through_the_machine_wide_gate_slot_lock(tmp_path: Path) -> None: From 36b0d80d3a1459f116f07e45b205e953e21c8978 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:09:00 -0700 Subject: [PATCH 20/49] test(cost): pin the nested-reasoning helper's clamps The strip in text_tokens_without_nested_reasoning is capped at the reasoning share, the reported text, and the over-sum past completion_tokens. Dropping the caps to a bare over-sum passed every existing test, so this pins each cap with a parametrized helper test plus one billing test where text over-reports past the reasoning share and only the nested share may be netted out --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 23 ++++++++++++++ tests/test_litellm/types/test_types_utils.py | 31 ++++++++++++++++++- 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index e8eae73df3f..c9fe53651cd 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -4821,6 +4821,29 @@ def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tok assert completion_cost == pytest.approx(44 * info["output_cost_per_token"]) +def test_generic_cost_per_token_strips_only_the_reasoning_share_when_text_over_reports( + _local_model_cost_map: None, +) -> None: + """Text over-reported past the reasoning share keeps its extra tokens billed; only the nested reasoning is netted out.""" + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=120, + completion_tokens=100, + total_tokens=220, + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=100, audio_tokens=70, reasoning_tokens=10), + ) + + _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert breakdown.reasoning_cost == pytest.approx(10 * info["output_cost_per_token"]) + assert completion_cost == pytest.approx( + 100 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"] + ) + + def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output(_local_model_cost_map: None) -> None: """Audio-output realtime usage nests reasoning inside text_tokens next to audio_tokens; text is billed net of it.""" diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index 554604ab200..5f44ba1773e 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -3,7 +3,7 @@ from typing import Final import pytest -from litellm.types.utils import HiddenParams, all_litellm_params +from litellm.types.utils import HiddenParams, all_litellm_params, text_tokens_without_nested_reasoning def test_rust_is_a_known_litellm_param(): @@ -768,3 +768,32 @@ def test_image_response_keeps_background(): response = ImageResponse(created=1, data=[{"b64_json": "aGk="}], background="transparent", output_format="png") assert response.background == "transparent" assert response.model_dump()["background"] == "transparent" + + +@pytest.mark.parametrize( + ("completion_tokens", "text_tokens", "reasoning_tokens", "other_modality_tokens", "expected_text_tokens"), + ( + pytest.param(50, 30, 20, 0, 30, id="details_sum_to_completion_is_a_no_op"), + pytest.param(34, 30, 24, 0, 10, id="strip_is_capped_at_the_over_sum"), + pytest.param(100, 100, 10, 70, 90, id="only_the_reasoning_share_is_stripped_when_text_over_reports_further"), + pytest.param(10, 5, 20, 0, 0, id="text_never_goes_negative_when_reasoning_exceeds_it"), + ), +) +def test_text_tokens_without_nested_reasoning_clamps( + completion_tokens: int, + text_tokens: int, + reasoning_tokens: int, + other_modality_tokens: int, + expected_text_tokens: int, +) -> None: + """The strip never exceeds the reasoning share, the reported text, or the over-sum past completion_tokens.""" + + assert ( + text_tokens_without_nested_reasoning( + completion_tokens=completion_tokens, + text_tokens=text_tokens, + reasoning_tokens=reasoning_tokens, + other_modality_tokens=other_modality_tokens, + ) + == expected_text_tokens + ) From 827554954d305b7c8c24f524c961cf7a27e3a029 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:16:14 -0700 Subject: [PATCH 21/49] fix(image_handling): answer every SSRF rejection with one message so error text cannot probe internal hostnames --- .../prompt_templates/image_handling.py | 10 ++--- litellm/litellm_core_utils/url_utils.py | 8 +--- .../litellm_core_utils/test_image_handling.py | 41 ++++++++----------- .../litellm_core_utils/test_url_utils.py | 19 +-------- 4 files changed, 24 insertions(+), 54 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 99efbd13320..24f3b8bca7f 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -15,7 +15,7 @@ import litellm from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB -from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError, async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 @@ -80,13 +80,9 @@ def _process_image_response(response: Response, url: str) -> str: def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict) - if isinstance(verdict, HostResolutionError): - return litellm.ImageFetchError( - f"Error: Unable to fetch image from URL. The image host could not be resolved. url={url}" - ) return litellm.ImageFetchError( - "Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; " - f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}" + "Error: Unable to fetch image from URL. The proxy could not resolve this host or its URL policy rejected it; " + f"an admin can check the proxy log and `user_url_allowed_hosts` in general_settings. url={url}" ) diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index f4388d5fd12..fa070a648f5 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -93,10 +93,6 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" -class HostResolutionError(SSRFError): - pass - - def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. @@ -328,10 +324,10 @@ def validate_url(url: str) -> tuple[str, str]: try: addrinfo: Final = socket.getaddrinfo(hostname, effective_port, proto=socket.IPPROTO_TCP) except socket.gaierror as e: - raise HostResolutionError(f"DNS resolution failed for '{hostname}': {e}") + raise SSRFError(f"DNS resolution failed for '{hostname}': {e}") if not addrinfo: - raise HostResolutionError(f"No addresses found for '{hostname}'") + raise SSRFError(f"No addresses found for '{hostname}'") if not is_allowlisted: for addrinfo_entry in addrinfo: diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index 3a909298aba..8fa4bd6c14d 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -17,7 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, convert_url_to_base64, ) -from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError +from litellm.litellm_core_utils.url_utils import SSRFError @pytest.fixture(autouse=True) @@ -425,32 +425,25 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail _SSRF_VERDICTS = ( - ( - SSRFError( - "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " - "add the host to `user_url_allowed_hosts` in general_settings." - ), - "The proxy's URL policy rejected this host; an admin can allow it with `user_url_allowed_hosts`", + SSRFError( + "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " + "add the host to `user_url_allowed_hosts` in general_settings." ), - ( - HostResolutionError( - "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known" - ), - "The image host could not be resolved", - ), - (HostResolutionError("No addresses found for 'internal.example'"), "The image host could not be resolved"), + SSRFError("DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known"), + SSRFError("No addresses found for 'internal.example'"), ) def _assert_verdict_free_messages(messages, url): - for message, (verdict, expected_guidance) in zip(messages, _SSRF_VERDICTS, strict=True): - assert expected_guidance in message - assert url in message - assert "10.0.0.8" not in message - assert "DNS" not in message - assert "No addresses" not in message - if isinstance(verdict, HostResolutionError): - assert "user_url_allowed_hosts" not in message + assert len(messages) == len(_SSRF_VERDICTS) + assert len(set(messages)) == 1, "a caller must not be able to tell a blocked host from one that does not resolve" + message = messages[0] + assert "The proxy could not resolve this host or its URL policy rejected it" in message + assert "user_url_allowed_hosts" in message + assert url in message + assert "10.0.0.8" not in message + assert "DNS" not in message + assert "No addresses" not in message async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): @@ -458,7 +451,7 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - for verdict, _ in _SSRF_VERDICTS: + for verdict in _SSRF_VERDICTS: async def block(client, fetched_url, verdict=verdict, **kwargs): attempts.append(fetched_url) @@ -478,7 +471,7 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - for verdict, _ in _SSRF_VERDICTS: + for verdict in _SSRF_VERDICTS: def block(client, fetched_url, verdict=verdict, **kwargs): attempts.append(fetched_url) diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index 151c7ceb5bd..fccdc1a2a0a 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -9,7 +9,6 @@ import pytest import litellm from litellm.litellm_core_utils import url_utils from litellm.litellm_core_utils.url_utils import ( - HostResolutionError, SSRFError, _is_blocked_ip, assert_same_origin, @@ -162,24 +161,10 @@ class TestValidateUrl: assert "/path" in rewritten assert "key=value" in rewritten - def test_dns_failure_raises_a_host_resolution_error(self, mock_dns_failure): - with pytest.raises(HostResolutionError, match="DNS resolution failed"): + def test_dns_failure_raises(self, mock_dns_failure): + with pytest.raises(SSRFError, match="DNS resolution failed"): validate_url("http://this-domain-does-not-exist-xyz123.invalid/test") - def test_empty_resolution_raises_a_host_resolution_error(self, monkeypatch): - monkeypatch.setattr(url_utils.socket, "getaddrinfo", lambda *args, **kwargs: []) - with pytest.raises(HostResolutionError, match="No addresses found"): - validate_url("http://this-domain-resolves-to-nothing.invalid/test") - - def test_blocked_address_is_not_a_host_resolution_error(self, monkeypatch): - def fake(host, port, *a, **kw): - return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.8", port or 80))] - - monkeypatch.setattr(url_utils.socket, "getaddrinfo", fake) - with pytest.raises(SSRFError, match="blocked address") as raised: - validate_url("http://internal.example/test") - assert not isinstance(raised.value, HostResolutionError) - def test_blocks_localhost_hostname(self, monkeypatch): def fake(host, port, *a, **kw): return [ From e3366dddf44c7450a856ed472c2966798413ef36 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:26:16 -0700 Subject: [PATCH 22/49] fix(check): trigger the test-tree checks on the gate script and drop the scope comment --- scripts/pre_commit_lint.sh | 8 +++----- tests/test_litellm/test_pre_commit_lint.py | 8 +++++++- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index ab84d2d518a..f245803408c 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -12,7 +12,8 @@ # - litellm/ Python -> `make lint` (test-linting.yml's lint job) # - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) # + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) -# - tests/ Python, ruff-tests.toml, test-quality-budget.json, scripts/check_test_quality.py +# - tests/ Python, ruff-tests.toml, test-quality-budget.json, scripts/check_test_quality.py, +# scripts/test_quality_gate.py # -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's # test-tree ruff and test-quality budget steps) # - dashboard -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) @@ -91,14 +92,11 @@ existing_files() { litellm_py_pattern='^litellm/.*\.py$' e2e_py_pattern='^tests/e2e/.*\.py$' -test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|test-quality-budget\.json|scripts/check_test_quality\.py)$' +test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|test-quality-budget\.json|scripts/(check_test_quality|test_quality_gate)\.py)$' spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$' ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' -# CI's lint job (test-linting.yml) is make lint: litellm/ checks plus two test-tree -# steps (ruff over ruff-tests.toml and the test-quality budget). Trigger the slow -# make lint on litellm/ files only; without them, the test-tree steps run on their own below. litellm_py_files=$(scope_match "$litellm_py_pattern") e2e_py_files=$(scope_match "$e2e_py_pattern") test_tree_files=$(scope_match "$test_tree_pattern") diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index e17abf2fa9f..b84cb8aa657 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -472,7 +472,13 @@ def test_tests_only_change_runs_the_whole_test_tree_ruff_and_the_quality_gate(tm @pytest.mark.parametrize( "changed", - ["ruff-tests.toml", "test-quality-budget.json", "scripts/check_test_quality.py", "tests/e2e/test_x.py"], + [ + "ruff-tests.toml", + "test-quality-budget.json", + "scripts/check_test_quality.py", + "scripts/test_quality_gate.py", + "tests/e2e/test_x.py", + ], ) def test_test_tree_lint_inputs_trigger_the_test_tree_checks(tmp_path: Path, changed: str) -> None: repo, bin_dir = _sandbox(tmp_path) From 2d2b5dabf27405fd413dfbfd71aef5a304775a75 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 01:44:16 -0700 Subject: [PATCH 23/49] fix(test-quality-gate): tear the base worktree down on SIGTERM and SIGHUP --- scripts/test_quality_gate.py | 22 ++++-- tests/test_litellm/test_test_quality_gate.py | 77 ++++++++++++++++++++ 2 files changed, 92 insertions(+), 7 deletions(-) diff --git a/scripts/test_quality_gate.py b/scripts/test_quality_gate.py index 7d34b194f1c..4f79c9488b1 100644 --- a/scripts/test_quality_gate.py +++ b/scripts/test_quality_gate.py @@ -34,13 +34,14 @@ import argparse import json import re import shutil +import signal import subprocess import sys import tempfile from collections import Counter from collections.abc import Mapping, Sequence from pathlib import Path -from types import MappingProxyType +from types import FrameType, MappingProxyType from typing import Final, NamedTuple REPO_ROOT: Final = Path(__file__).resolve().parent.parent @@ -48,6 +49,7 @@ CHECKER: Final = REPO_ROOT / "scripts" / "check_test_quality.py" BUDGET_PATH: Final = REPO_ROOT / "test-quality-budget.json" TARGET: Final = "tests" DEFAULT_BASE: Final = "origin/litellm_internal_staging" +TERMINATION_SIGNALS: Final = (signal.SIGTERM, signal.SIGHUP) _HUNK: Final = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@", re.MULTILINE) _FILE_HEADER: Final = re.compile(r"^\+\+\+ b/(.+)$", re.MULTILINE) @@ -116,22 +118,28 @@ def count_by_rule(violations: Sequence[Violation]) -> Mapping[str, int]: return MappingProxyType(dict(Counter(v.code for v in violations))) -def base_counts(ref: str) -> Mapping[str, int]: +def _exit_on_termination(signum: int, _frame: FrameType | None) -> None: + raise SystemExit(128 + signum) + + +def base_counts(ref: str, repo_root: Path = REPO_ROOT, checker: Path = CHECKER) -> Mapping[str, int]: """Rule counts at `ref`, measured with the *current* rule logic rather than whatever the checker looked like at that commit.""" + for termination in TERMINATION_SIGNALS: + signal.signal(termination, _exit_on_termination) parent: Final = Path(tempfile.mkdtemp(prefix="tq_base_")) worktree: Final = parent / "wt" try: - _run(["git", "worktree", "add", "--detach", str(worktree), ref]) + _run(["git", "worktree", "add", "--detach", str(worktree), ref], cwd=repo_root) (worktree / "scripts").mkdir(parents=True, exist_ok=True) - checker: Final = worktree / "scripts" / "check_test_quality.py" - shutil.copy(CHECKER, checker) - return count_by_rule(_check(worktree, checker)) + base_checker: Final = worktree / "scripts" / "check_test_quality.py" + shutil.copy(checker, base_checker) + return count_by_rule(_check(worktree, base_checker)) finally: # Teardown must never raise, or it masks the real error when the body failed. subprocess.run( ["git", "worktree", "remove", "--force", str(worktree)], - cwd=REPO_ROOT, capture_output=True, text=True, + cwd=repo_root, capture_output=True, text=True, ) shutil.rmtree(parent, ignore_errors=True) diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/test_litellm/test_test_quality_gate.py index 8cce6bc735a..0664153f671 100644 --- a/tests/test_litellm/test_test_quality_gate.py +++ b/tests/test_litellm/test_test_quality_gate.py @@ -9,7 +9,13 @@ file:line. """ import importlib.util +import os +import signal +import subprocess import sys +import time +from collections.abc import Callable +from contextlib import suppress from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[2] @@ -23,6 +29,15 @@ _spec.loader.exec_module(gate) _BUDGET = {"TQ001": {"limit": 10}, "TQ003": {"limit": 5}} +_SCAN_BASE = ( + "import importlib.util, pathlib, sys\n" + "spec = importlib.util.spec_from_file_location('test_quality_gate', sys.argv[1])\n" + "gate = importlib.util.module_from_spec(spec)\n" + "sys.modules[spec.name] = gate\n" + "spec.loader.exec_module(gate)\n" + "gate.base_counts('HEAD', repo_root=pathlib.Path(sys.argv[2]), checker=pathlib.Path(sys.argv[3]))\n" +) + def test_a_rule_within_its_limit_is_not_a_breach(): assert gate.evaluate({"TQ001": 10}, {"TQ001": 10}, _BUDGET) == () @@ -146,3 +161,65 @@ def test_the_shipped_budget_covers_every_rule_the_checker_can_emit(): budget = json.loads((_REPO_ROOT / "test-quality-budget.json").read_text()) assert set(budget) == {"TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ008"} assert all(spec["limit"] >= 0 for spec in budget.values()) + + +def _git(cwd: Path, *args: str) -> str: + proc = subprocess.run(["git", *args], cwd=cwd, capture_output=True, text=True) + assert proc.returncode == 0, proc.stderr + return proc.stdout.strip() + + +def _committed_repo(tmp_path: Path) -> Path: + repo = tmp_path / "repo" + (repo / "tests").mkdir(parents=True) + (repo / "tests" / "test_seed.py").write_text("def test_seed():\n assert True\n") + _git(repo, "init", "-q", "-b", "main") + _git(repo, "config", "user.email", "gate@example.com") + _git(repo, "config", "user.name", "gate") + _git(repo, "config", "commit.gpgsign", "false") + _git(repo, "add", "-A") + _git(repo, "commit", "-q", "-m", "seed") + return repo + + +def _wait_until(predicate: Callable[[], bool], timeout_seconds: float) -> bool: + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.05) + return predicate() + + +def _reap(process: subprocess.Popen[bytes]) -> None: + with suppress(subprocess.TimeoutExpired): + process.wait(timeout=10) + if process.poll() is None: + process.kill() + process.wait(timeout=10) + + +def _registered_worktrees(repo: Path) -> int: + listing = _git(repo, "worktree", "list", "--porcelain") + return sum(line.startswith("worktree ") for line in listing.splitlines()) + + +def test_a_terminated_base_scan_still_removes_its_worktree(tmp_path: Path) -> None: + repo = _committed_repo(tmp_path) + scanning = tmp_path / "scanning" + slow_checker = tmp_path / "slow_checker.py" + slow_checker.write_text(f"import pathlib, time\npathlib.Path({str(scanning)!r}).touch()\ntime.sleep(30)\n") + temp_dir = tmp_path / "tmp" + temp_dir.mkdir() + scan = subprocess.Popen( + [sys.executable, "-c", _SCAN_BASE, str(_MODULE_PATH), str(repo), str(slow_checker)], + env={**os.environ, "TMPDIR": str(temp_dir)}, + ) + try: + assert _wait_until(scanning.exists, 30), "the base scan never reached the checker" + scan.send_signal(signal.SIGTERM) + assert scan.wait(timeout=30) == 128 + signal.SIGTERM + finally: + _reap(scan) + assert _registered_worktrees(repo) == 1 + assert list(temp_dir.iterdir()) == [] From 952f082e3ee80eb6bec0855a752ee3c101652744 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:07:14 -0700 Subject: [PATCH 24/49] fix(test-quality-gate): keep a termination signal the parent already ignores ignored The SIGTERM/SIGHUP teardown handlers were installed unconditionally, so a base scan started under nohup (SIGHUP inherited as SIG_IGN) would start dying on hangups it was told to ignore. Install them only where the disposition is still the default, and cover the ignored case with a regression test that hangs up a scan started with SIGHUP ignored and expects it to finish. --- scripts/test_quality_gate.py | 9 +++- tests/test_litellm/test_test_quality_gate.py | 56 ++++++++++++++++---- 2 files changed, 53 insertions(+), 12 deletions(-) diff --git a/scripts/test_quality_gate.py b/scripts/test_quality_gate.py index 4f79c9488b1..c3316915b5e 100644 --- a/scripts/test_quality_gate.py +++ b/scripts/test_quality_gate.py @@ -122,11 +122,16 @@ def _exit_on_termination(signum: int, _frame: FrameType | None) -> None: raise SystemExit(128 + signum) +def _install_termination_handlers() -> None: + for termination in TERMINATION_SIGNALS: + if signal.getsignal(termination) == signal.SIG_DFL: + signal.signal(termination, _exit_on_termination) + + def base_counts(ref: str, repo_root: Path = REPO_ROOT, checker: Path = CHECKER) -> Mapping[str, int]: """Rule counts at `ref`, measured with the *current* rule logic rather than whatever the checker looked like at that commit.""" - for termination in TERMINATION_SIGNALS: - signal.signal(termination, _exit_on_termination) + _install_termination_handlers() parent: Final = Path(tempfile.mkdtemp(prefix="tq_base_")) worktree: Final = parent / "wt" try: diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/test_litellm/test_test_quality_gate.py index 0664153f671..b3511eabbfb 100644 --- a/tests/test_litellm/test_test_quality_gate.py +++ b/tests/test_litellm/test_test_quality_gate.py @@ -17,6 +17,7 @@ import time from collections.abc import Callable from contextlib import suppress from pathlib import Path +from typing import NamedTuple _REPO_ROOT = Path(__file__).resolve().parents[2] _MODULE_PATH = _REPO_ROOT / "scripts" / "test_quality_gate.py" @@ -37,6 +38,7 @@ _SCAN_BASE = ( "spec.loader.exec_module(gate)\n" "gate.base_counts('HEAD', repo_root=pathlib.Path(sys.argv[2]), checker=pathlib.Path(sys.argv[3]))\n" ) +_SCAN_BASE_WITH_SIGHUP_IGNORED = "import signal\nsignal.signal(signal.SIGHUP, signal.SIG_IGN)\n" + _SCAN_BASE def test_a_rule_within_its_limit_is_not_a_breach(): @@ -204,22 +206,56 @@ def _registered_worktrees(repo: Path) -> int: return sum(line.startswith("worktree ") for line in listing.splitlines()) -def test_a_terminated_base_scan_still_removes_its_worktree(tmp_path: Path) -> None: +class _StalledScan(NamedTuple): + process: subprocess.Popen[bytes] + repo: Path + release: Path + temp_dir: Path + + +def _base_scan_stalled_in_its_checker(tmp_path: Path, driver: str) -> _StalledScan: repo = _committed_repo(tmp_path) scanning = tmp_path / "scanning" + release = tmp_path / "release" slow_checker = tmp_path / "slow_checker.py" - slow_checker.write_text(f"import pathlib, time\npathlib.Path({str(scanning)!r}).touch()\ntime.sleep(30)\n") + slow_checker.write_text( + "import pathlib, time\n" + f"pathlib.Path({str(scanning)!r}).touch()\n" + f"while not pathlib.Path({str(release)!r}).exists():\n" + " time.sleep(0.05)\n" + ) temp_dir = tmp_path / "tmp" temp_dir.mkdir() scan = subprocess.Popen( - [sys.executable, "-c", _SCAN_BASE, str(_MODULE_PATH), str(repo), str(slow_checker)], + [sys.executable, "-c", driver, str(_MODULE_PATH), str(repo), str(slow_checker)], env={**os.environ, "TMPDIR": str(temp_dir)}, ) - try: - assert _wait_until(scanning.exists, 30), "the base scan never reached the checker" - scan.send_signal(signal.SIGTERM) - assert scan.wait(timeout=30) == 128 + signal.SIGTERM - finally: + if not _wait_until(scanning.exists, 30): _reap(scan) - assert _registered_worktrees(repo) == 1 - assert list(temp_dir.iterdir()) == [] + raise AssertionError("the base scan never reached the checker") + return _StalledScan(scan, repo, release, temp_dir) + + +def test_a_terminated_base_scan_still_removes_its_worktree(tmp_path: Path) -> None: + stalled = _base_scan_stalled_in_its_checker(tmp_path, _SCAN_BASE) + try: + stalled.process.send_signal(signal.SIGTERM) + assert stalled.process.wait(timeout=30) == 128 + signal.SIGTERM + finally: + _reap(stalled.process) + assert _registered_worktrees(stalled.repo) == 1 + assert list(stalled.temp_dir.iterdir()) == [] + + +def test_a_base_scan_keeps_ignoring_the_hangup_its_parent_ignored(tmp_path: Path) -> None: + stalled = _base_scan_stalled_in_its_checker(tmp_path, _SCAN_BASE_WITH_SIGHUP_IGNORED) + try: + stalled.process.send_signal(signal.SIGHUP) + time.sleep(1) + assert stalled.process.poll() is None, "a hangup the parent ignored killed the scan" + stalled.release.touch() + assert stalled.process.wait(timeout=30) == 0 + finally: + _reap(stalled.process) + assert _registered_worktrees(stalled.repo) == 1 + assert list(stalled.temp_dir.iterdir()) == [] From 292f926fcd5b4882e3cda1e28c6fa9211a7860d4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:56:17 -0700 Subject: [PATCH 25/49] fix(oci): stream Cohere tool-calling answers once OCI Cohere restates the whole assistant text on the chunk that carries the tool calls and again on the terminal chunk that carries chatHistory. Only the terminal restatement was dropped, so a tool-calling turn streamed the text twice. Treat a toolCalls-bearing chunk as a restatement too and drop its text once deltas were already emitted. Resolves LIT-6819 --- litellm/llms/oci/chat/cohere.py | 52 ++++------------- .../oci/chat/test_oci_cohere_tool_calls.py | 57 +++++++++++++++++++ 2 files changed, 69 insertions(+), 40 deletions(-) diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 6e9bb83b0a0..1b494ebad47 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -290,19 +290,14 @@ def handle_cohere_stream_chunk( ) -> ModelResponseStream: """Parse a single Cohere SSE chunk into a LiteLLM ModelResponseStream. - ``prior_tool_calls_emitted`` lets the caller signal whether tool calls - were already emitted in earlier chunks of the same stream. When set, the - terminal consolidation chunk's tool calls are suppressed (they would - duplicate prior deltas); otherwise they are passed through so a stream - that delivers tool calls only on the terminal chunk doesn't silently - drop them. - - ``prior_text_emitted`` plays the analogous role for the ``text`` field: - when set, the terminal consolidation chunk's ``text`` is suppressed - (it would re-emit the full assembled response on top of prior deltas); - when unset (e.g. a degenerate stream that delivers the entire response - in a single SSE event carrying both ``chatHistory`` and ``finishReason``), - the text is passed through so the response content isn't silently lost. + OCI Cohere streams the answer as single-token ``text`` deltas, then restates + the whole assembled ``text`` on every chunk that carries ``toolCalls`` or + ``chatHistory`` (the tool-calls event and the terminal event). Once the + caller reports that earlier chunks already emitted text + (``prior_text_emitted``), those restatements are dropped so the client does + not see the answer twice; a stream whose only text lives on such a chunk + keeps it. ``prior_tool_calls_emitted`` plays the same role for the tool + calls the terminal ``chatHistory`` chunk repeats. """ try: typed_chunk: Final = CohereStreamChunk.model_validate(dict_chunk) @@ -315,33 +310,10 @@ def handle_cohere_stream_chunk( if typed_chunk.index is None: typed_chunk.index = 0 - # OCI Cohere's terminal SSE event re-sends the full assembled response in - # `text` alongside a populated `chatHistory` and a non-null `finishReason`. - # Emitting that text would concatenate the whole response onto the - # already-streamed deltas. We require both signals to be present so that a - # future API change which adds `chatHistory` to intermediate chunks (or a - # rare early-populated case) doesn't silently drop legitimate token deltas. - is_terminal_consolidation: Final = typed_chunk.chatHistory is not None and typed_chunk.finishReason is not None - # On non-terminal text-free chunks (e.g. tool-call-only or keep-alive - # chunks) emit ``content=None`` rather than ``content=""`` so downstream - # stream-mergers that distinguish "no text in this delta" from "an - # explicitly empty text delta" behave correctly. - # - # We only suppress the terminal chunk's ``text`` when the caller has - # confirmed that text deltas were already emitted earlier — otherwise - # (e.g. a degenerate stream that delivers the whole response in a - # single SSE event), passing it through is the only chance to surface it. - text: Final[str | None] = None if (is_terminal_consolidation and prior_text_emitted) else typed_chunk.text - - # Tool calls on the terminal consolidation chunk (whether from - # `typed_chunk.toolCalls` or from `chatHistory`) typically restate what - # was already streamed in intermediate chunks. Re-emitting them would - # mint fresh `uuid4` IDs and cause downstream consumers to execute each - # tool call twice. We only suppress when the caller has confirmed that - # tool calls were already emitted earlier — otherwise (e.g. a short - # response that delivers tool calls exclusively on the terminal chunk), - # passing them through is the only chance to surface them. - cohere_tool_calls = None if (is_terminal_consolidation and prior_tool_calls_emitted) else typed_chunk.toolCalls + restates_text: Final = typed_chunk.chatHistory is not None or typed_chunk.toolCalls is not None + restates_tool_calls: Final = typed_chunk.chatHistory is not None + text: Final[str | None] = None if (restates_text and prior_text_emitted) else typed_chunk.text + cohere_tool_calls: Final = None if (restates_tool_calls and prior_tool_calls_emitted) else typed_chunk.toolCalls tool_calls: list[dict[str, object]] | None = None if cohere_tool_calls: diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py index 5dd44d72d68..97b2fd88cb9 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py @@ -445,6 +445,63 @@ class TestOCICohereToolCalls: assert result.choices[0].index == 0 assert result.choices[0].finish_reason == "stop" # COMPLETE is mapped to stop + _TOOL_TURN_TEXT = "I will use the tool to find out the weather in Paris." + _TOOL_TURN_DELTAS = ["I", " will", " use", " the", " tool", " to", " find", " out", " the", " weather", " in", " Paris", "."] + _TOOL_TURN_CALLS = [{"name": "get_weather", "parameters": {"city": "Paris"}}] + _TOOL_TURN_HISTORY = [ + {"role": "USER", "message": "Briefly say what you will do, then find out the weather in Paris using the tool."}, + {"role": "CHATBOT", "message": _TOOL_TURN_TEXT, "toolCalls": _TOOL_TURN_CALLS}, + ] + _TOOL_TURN_TERMINAL_TOGETHER = [ + { + "apiFormat": "COHERE", + "text": _TOOL_TURN_TEXT, + "chatHistory": _TOOL_TURN_HISTORY, + "finishReason": "COMPLETE", + "toolCalls": _TOOL_TURN_CALLS, + }, + ] + _TOOL_TURN_TERMINAL_SPLIT = [ + {"apiFormat": "COHERE", "text": _TOOL_TURN_TEXT, "chatHistory": _TOOL_TURN_HISTORY, "toolCalls": _TOOL_TURN_CALLS}, + {"apiFormat": "COHERE", "finishReason": "COMPLETE"}, + ] + + @staticmethod + def _drain_cohere_stream(events): + wrapper = OCIStreamWrapper(completion_stream=MagicMock(), model="cohere.command-a-03-2025", logging_obj=MagicMock()) + chunks = [wrapper.chunk_creator(f"data: {json.dumps(event)}") for event in events] + content = "".join(chunk.choices[0].delta.content or "" for chunk in chunks) + tool_calls = [call for chunk in chunks for call in (chunk.choices[0].delta.tool_calls or [])] + finish_reasons = [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] + return content, tool_calls, finish_reasons + + @pytest.mark.parametrize("terminal_events", [_TOOL_TURN_TERMINAL_TOGETHER, _TOOL_TURN_TERMINAL_SPLIT]) + def test_cohere_tool_turn_streams_the_answer_once(self, terminal_events): + """OCI restates the whole answer on the tool-calls chunk and again on the terminal chunk; + the client must read it exactly once, with one tool call and one finish reason.""" + deltas = [{"apiFormat": "COHERE", "text": token} for token in self._TOOL_TURN_DELTAS] + tool_calls_event = {"apiFormat": "COHERE", "text": self._TOOL_TURN_TEXT, "toolCalls": self._TOOL_TURN_CALLS} + + content, tool_calls, finish_reasons = self._drain_cohere_stream([*deltas, tool_calls_event, *terminal_events]) + + assert content == self._TOOL_TURN_TEXT + assert [(call["function"]["name"], call["function"]["arguments"]) for call in tool_calls] == [ + ("get_weather", '{"city": "Paris"}') + ] + assert finish_reasons == ["stop"] + + def test_cohere_tool_turn_without_preamble_deltas_keeps_the_only_text(self): + """When the tool-calls chunk carries the only copy of the text, dropping it would lose the answer.""" + tool_calls_event = {"apiFormat": "COHERE", "text": self._TOOL_TURN_TEXT, "toolCalls": self._TOOL_TURN_CALLS} + + content, tool_calls, finish_reasons = self._drain_cohere_stream( + [tool_calls_event, *self._TOOL_TURN_TERMINAL_TOGETHER] + ) + + assert content == self._TOOL_TURN_TEXT + assert len(tool_calls) == 1 + assert finish_reasons == ["stop"] + def test_cohere_parameter_mapping_excludes_tool_choice(self): """Test that tool_choice is excluded from Cohere parameter mapping""" config = OCIChatConfig() From 43bb55d84934762a775bef3c32f6ed5ad8bc5ce6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:38:23 -0700 Subject: [PATCH 26/49] style(oci): wrap the Cohere tool-turn test fixtures to 120 columns --- .../llms/oci/chat/test_oci_cohere_tool_calls.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py index 97b2fd88cb9..729a2d25f41 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py @@ -446,7 +446,9 @@ class TestOCICohereToolCalls: assert result.choices[0].finish_reason == "stop" # COMPLETE is mapped to stop _TOOL_TURN_TEXT = "I will use the tool to find out the weather in Paris." - _TOOL_TURN_DELTAS = ["I", " will", " use", " the", " tool", " to", " find", " out", " the", " weather", " in", " Paris", "."] + _TOOL_TURN_DELTAS = [ + "I", " will", " use", " the", " tool", " to", " find", " out", " the", " weather", " in", " Paris", ".", + ] _TOOL_TURN_CALLS = [{"name": "get_weather", "parameters": {"city": "Paris"}}] _TOOL_TURN_HISTORY = [ {"role": "USER", "message": "Briefly say what you will do, then find out the weather in Paris using the tool."}, @@ -462,13 +464,20 @@ class TestOCICohereToolCalls: }, ] _TOOL_TURN_TERMINAL_SPLIT = [ - {"apiFormat": "COHERE", "text": _TOOL_TURN_TEXT, "chatHistory": _TOOL_TURN_HISTORY, "toolCalls": _TOOL_TURN_CALLS}, + { + "apiFormat": "COHERE", + "text": _TOOL_TURN_TEXT, + "chatHistory": _TOOL_TURN_HISTORY, + "toolCalls": _TOOL_TURN_CALLS, + }, {"apiFormat": "COHERE", "finishReason": "COMPLETE"}, ] @staticmethod def _drain_cohere_stream(events): - wrapper = OCIStreamWrapper(completion_stream=MagicMock(), model="cohere.command-a-03-2025", logging_obj=MagicMock()) + wrapper = OCIStreamWrapper( + completion_stream=MagicMock(), model="cohere.command-a-03-2025", logging_obj=MagicMock() + ) chunks = [wrapper.chunk_creator(f"data: {json.dumps(event)}") for event in events] content = "".join(chunk.choices[0].delta.content or "" for chunk in chunks) tool_calls = [call for chunk in chunks for call in (chunk.choices[0].delta.tool_calls or [])] From 347ea8f7ca08b2c5ce8fde4430f867b47ace4489 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:54:50 -0700 Subject: [PATCH 27/49] fix: keep provider id and metadata on Responses API bridged chat completions --- .../transformation.py | 23 +++- ...responses_transformation_transformation.py | 106 ++++++++++++++++++ 2 files changed, 126 insertions(+), 3 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index d70f947469a..331cd081fcd 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -33,7 +33,7 @@ from litellm.responses.sse_output_recovery import ( record_output_item_chunk, record_output_text_chunk, ) -from litellm.responses.utils import normalize_responses_api_stream_options +from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_responses_api_stream_options from litellm.types.llms.openai import ( REASONING_EFFORT, ChatCompletionAnnotation, @@ -54,7 +54,7 @@ if TYPE_CHECKING: ) from pydantic import BaseModel - from litellm import LiteLLMLoggingObj, ModelResponse + from litellm import LiteLLMLoggingObj from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.openai import ( ALL_RESPONSES_API_TOOL_PARAMS, @@ -69,6 +69,15 @@ if TYPE_CHECKING: from litellm.types.utils import Choices +_CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage")) + + +def _upstream_response_id(response_id: str | None) -> str | None: + if response_id is None: + return None + return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id) + + class _ReasoningSummaryText(TypedDict): type: str text: str @@ -904,6 +913,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_response.usage), ) + model_response.id = _upstream_response_id(raw_response.id) or raw_response.id + provider_extras: Final = raw_response.model_extra.items() if raw_response.model_extra else () + for key, value in provider_extras: + if key not in _CHAT_COMPLETION_FIELDS and value is not None: + setattr(model_response, key, value) + # Preserve hidden params from the ResponsesAPIResponse, especially the headers # which contain important provider information like x-request-id raw_response_hidden_params: Final = getattr(raw_response, "_hidden_params", {}) @@ -1359,14 +1374,16 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if event_type == "response.created": # Initial response creation event verbose_logger.debug("Chat provider: response.created -> %s", parsed_chunk) + created_response: Final = parsed_chunk.get("response") return ModelResponseStream( + id=_upstream_response_id(created_response.get("id")) if created_response else None, choices=[ StreamingChoices( index=0, delta=Delta(content=""), finish_reason=None, ) - ] + ], ) elif event_type == "response.output_item.added": # New output item added diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 6590718878d..97b18382a68 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3952,3 +3952,109 @@ def test_convert_chat_completion_messages_to_responses_api_tool_result_with_tool function_call_output = next(item for item in response if item.get("type") == "function_call_output") assert function_call_output["output"] == [{"type": "input_text", "text": "1 tool found"}] + + +def _litellm_encoded_response_id(upstream_id: str) -> str: + """The id a Responses API call hands back through LiteLLM: the provider's id wrapped with the + deployment it came from, the way ``/v1/responses`` clients see it.""" + import base64 + + tagged = f"litellm:custom_llm_provider:azure;model_id:deployment-1;response_id:{upstream_id}" + return "resp_" + base64.b64encode(tagged.encode()).decode() + + +def test_transform_response_keeps_upstream_id_and_provider_extras(): + """A chat completion bridged through the Responses API must answer with the provider's own + response id, decoded out of the deployment-tagged id LiteLLM wraps around it, and every + top-level field the provider adds beyond the Responses schema (Azure's ``content_filters``, + ``service_tier``), the way the native chat path passes unknown top-level fields through, + instead of a locally minted ``chatcmpl-`` id and nothing else.""" + from unittest.mock import Mock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse, Usage + + content_filters = [ + {"blocked": False, "source_type": "prompt", "content_filter_results": {"hate": {"filtered": False}}} + ] + raw_response = ResponsesAPIResponse.model_validate( + { + "id": _litellm_encoded_response_id("resp_azure_123"), + "created_at": 1734366691, + "object": "response", + "model": "gpt-5.6", + "status": "completed", + "output": [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_weather", + "arguments": '{"city": "Seattle"}', + "status": "completed", + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "service_tier": "default", + "content_filters": content_filters, + "max_tool_calls": None, + } + ) + model_response = ModelResponse( + id="chatcmpl-local", + created=1734366691, + model=None, + object="chat.completion", + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ) + + result = LiteLLMResponsesTransformationHandler().transform_response( + model="gpt-5.6", + raw_response=raw_response, + model_response=model_response, + logging_obj=Mock(), + request_data={"model": "gpt-5.6"}, + messages=[{"role": "user", "content": "What is the weather in Seattle?"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + dumped = result.model_dump() + + assert dumped["id"] == "resp_azure_123" + assert dumped["object"] == "chat.completion" + assert dumped["service_tier"] == "default" + assert dumped["content_filters"] == content_filters + assert "max_tool_calls" not in dumped, "a null provider field must not appear as a null top-level key" + assert "output" not in dumped and "status" not in dumped, ( + "Responses schema fields must not leak into the chat response" + ) + assert dumped["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "lookup_weather" + + +def test_streaming_chunks_carry_the_upstream_response_id(): + """Every streamed chunk of a bridged chat completion must carry the provider's response id + from ``response.created`` rather than a locally minted ``chatcmpl-`` id, so a client can + correlate the stream with the provider's request the same way the non-streaming path does.""" + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + encoded_id = _litellm_encoded_response_id("resp_azure_stream") + events = [ + {"type": "response.created", "response": {"id": encoded_id, "output": []}}, + {"type": "response.output_text.delta", "delta": "Hel"}, + {"type": "response.completed", "response": {"id": encoded_id, "output": [{"type": "message"}]}}, + ] + + ids = [iterator.chunk_parser(event).id for event in events] + + assert ids == ["resp_azure_stream"] * len(events), f"streamed chunks did not carry the upstream id: {ids}" From 089b4b8ff458681bcbfa9741e85952e3c6aa0ece Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:59:24 -0700 Subject: [PATCH 28/49] test: drop docstrings from the bridge id regression tests --- ..._litellm_responses_transformation_transformation.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 97b18382a68..1599449382b 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3955,8 +3955,6 @@ def test_convert_chat_completion_messages_to_responses_api_tool_result_with_tool def _litellm_encoded_response_id(upstream_id: str) -> str: - """The id a Responses API call hands back through LiteLLM: the provider's id wrapped with the - deployment it came from, the way ``/v1/responses`` clients see it.""" import base64 tagged = f"litellm:custom_llm_provider:azure;model_id:deployment-1;response_id:{upstream_id}" @@ -3964,11 +3962,6 @@ def _litellm_encoded_response_id(upstream_id: str) -> str: def test_transform_response_keeps_upstream_id_and_provider_extras(): - """A chat completion bridged through the Responses API must answer with the provider's own - response id, decoded out of the deployment-tagged id LiteLLM wraps around it, and every - top-level field the provider adds beyond the Responses schema (Azure's ``content_filters``, - ``service_tier``), the way the native chat path passes unknown top-level fields through, - instead of a locally minted ``chatcmpl-`` id and nothing else.""" from unittest.mock import Mock from litellm.completion_extras.litellm_responses_transformation.transformation import ( @@ -4040,9 +4033,6 @@ def test_transform_response_keeps_upstream_id_and_provider_extras(): def test_streaming_chunks_carry_the_upstream_response_id(): - """Every streamed chunk of a bridged chat completion must carry the provider's response id - from ``response.created`` rather than a locally minted ``chatcmpl-`` id, so a client can - correlate the stream with the provider's request the same way the non-streaming path does.""" from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, ) From 64cbe6d0aa332bb6295dd37d6eb2066629c3af07 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:21:40 -0700 Subject: [PATCH 29/49] fix(bridge): carry provider metadata on streamed chats and keep served ids in spend logs The terminal chunk of a bridged streaming chat now carries the same provider fields the non-streaming response does (service_tier, content_filters), so streaming clients see them on the final chunk The passthrough drops Responses API bookkeeping (background, top_logprobs, store, ...) by subtracting the OpenAI SDK's Response schema from the fields it copies instead of a hand-kept denylist Spend log rows for /v1/messages calls served through the Responses adapter keep the id the client was handed instead of the decoded upstream id --- .../transformation.py | 25 ++++- litellm/litellm_core_utils/litellm_logging.py | 9 +- ...responses_transformation_transformation.py | 105 +++++++++++++++++- .../test_litellm_logging.py | 33 ++++++ 4 files changed, 162 insertions(+), 10 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 331cd081fcd..87350b5479c 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -5,8 +5,11 @@ Handler for transforming /chat/completions api requests to litellm.responses req import json import os from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args +from openai.types.chat import ChatCompletion +from openai.types.responses import Response from openai.types.responses.custom_tool_param import CustomToolParam from openai.types.responses.response_input_param import ( FunctionCallOutput, @@ -43,6 +46,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParamFunctionChunk, Reasoning, ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, ResponsesAPIStreamEvents, ) from litellm.types.utils import GenericStreamingChunk, ModelResponseStream @@ -70,6 +74,19 @@ if TYPE_CHECKING: _CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage")) +_RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *ResponsesAPIResponse.model_fields)) - frozenset( + ChatCompletion.model_fields +) + + +def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]: + return MappingProxyType( + { + key: value + for key, value in (response_fields.items() if response_fields else ()) + if value is not None and key not in _CHAT_COMPLETION_FIELDS and key not in _RESPONSES_API_ONLY_FIELDS + } + ) def _upstream_response_id(response_id: str | None) -> str | None: @@ -914,10 +931,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) model_response.id = _upstream_response_id(raw_response.id) or raw_response.id - provider_extras: Final = raw_response.model_extra.items() if raw_response.model_extra else () - for key, value in provider_extras: - if key not in _CHAT_COMPLETION_FIELDS and value is not None: - setattr(model_response, key, value) + for key, value in _provider_metadata(raw_response.model_extra).items(): + setattr(model_response, key, value) # Preserve hidden params from the ResponsesAPIResponse, especially the headers # which contain important provider information like x-request-id @@ -1551,6 +1566,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): from litellm.responses.utils import ResponseAPILoggingUtils usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_data.get("usage")) + provider_metadata: Final = _provider_metadata(response_data) return ModelResponseStream( choices=[ StreamingChoices( @@ -1563,6 +1579,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) ], usage=usage, + provider_specific_fields=dict(provider_metadata) or None, # mutable-ok: field is typed dict ) else: pass diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c31c4323157..dbd58766488 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3918,11 +3918,12 @@ class Logging(LiteLLMLoggingBaseClass): LiteLLMResponsesTransformationHandler, ) + served_id: Final = _provider_response_id(result) try: - return LiteLLMResponsesTransformationHandler().transform_response( + translated: Final = LiteLLMResponsesTransformationHandler().transform_response( model=self.model, raw_response=result, - model_response=litellm.ModelResponse(id=_provider_response_id(result)), + model_response=litellm.ModelResponse(id=served_id), logging_obj=self, request_data={}, messages=[], @@ -3930,6 +3931,8 @@ class Logging(LiteLLMLoggingBaseClass): litellm_params={}, encoding=litellm.encoding, ) + translated.id = served_id or translated.id + return translated except Exception as e: verbose_logger.debug( "Responses API -> ModelResponse translation failed for " @@ -3937,7 +3940,7 @@ class Logging(LiteLLMLoggingBaseClass): "usage-only ModelResponse to keep the spend_logs row.", str(e), ) - model_response: Final = litellm.ModelResponse(id=_provider_response_id(result)) + model_response: Final = litellm.ModelResponse(id=served_id) model_response.model = self.model usage: Final = getattr(result, "usage", None) if usage is not None and ResponseAPILoggingUtils._is_response_api_usage(usage): diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 1599449382b..d4d47b145d1 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3955,10 +3955,11 @@ def test_convert_chat_completion_messages_to_responses_api_tool_result_with_tool def _litellm_encoded_response_id(upstream_id: str) -> str: - import base64 + from litellm.responses.utils import ResponsesAPIRequestUtils - tagged = f"litellm:custom_llm_provider:azure;model_id:deployment-1;response_id:{upstream_id}" - return "resp_" + base64.b64encode(tagged.encode()).decode() + return ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="azure", model_id="deployment-1", response_id=upstream_id + ) def test_transform_response_keeps_upstream_id_and_provider_extras(): @@ -3997,6 +3998,9 @@ def test_transform_response_keeps_upstream_id_and_provider_extras(): "service_tier": "default", "content_filters": content_filters, "max_tool_calls": None, + "background": False, + "top_logprobs": 0, + "store": True, } ) model_response = ModelResponse( @@ -4029,9 +4033,63 @@ def test_transform_response_keeps_upstream_id_and_provider_extras(): assert "output" not in dumped and "status" not in dumped, ( "Responses schema fields must not leak into the chat response" ) + assert not {"background", "top_logprobs", "store"} & dumped.keys(), ( + "Responses API bookkeeping must not ride along as chat metadata" + ) assert dumped["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "lookup_weather" +def test_bridged_response_is_priced_by_the_reported_service_tier(): + from unittest.mock import Mock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse + + raw_response = ResponsesAPIResponse.model_validate( + { + "id": "resp_flex", + "created_at": 1734366691, + "object": "response", + "model": "gpt-5.4", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, + "service_tier": "flex", + } + ) + + result = LiteLLMResponsesTransformationHandler().transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + pricing = litellm.model_cost["gpt-5.4"] + flex_cost = 1000 * pricing["input_cost_per_token_flex"] + 100 * pricing["output_cost_per_token_flex"] + standard_cost = 1000 * pricing["input_cost_per_token"] + 100 * pricing["output_cost_per_token"] + + cost = litellm.completion_cost(completion_response=result, custom_llm_provider="openai") + + assert cost == pytest.approx(flex_cost) + assert cost < standard_cost + + def test_streaming_chunks_carry_the_upstream_response_id(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, @@ -4048,3 +4106,44 @@ def test_streaming_chunks_carry_the_upstream_response_id(): ids = [iterator.chunk_parser(event).id for event in events] assert ids == ["resp_azure_stream"] * len(events), f"streamed chunks did not carry the upstream id: {ids}" + + +def test_streaming_final_chunk_carries_provider_metadata(): + from unittest.mock import MagicMock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + content_filters = [{"blocked": False, "source_type": "completion", "content_filter_results": {}}] + events = [ + {"type": "response.created", "response": {"id": "resp_azure_stream", "output": []}}, + {"type": "response.output_text.delta", "delta": "Hello"}, + { + "type": "response.completed", + "response": { + "id": "resp_azure_stream", + "output": [{"type": "message"}], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + "service_tier": "default", + "content_filters": content_filters, + "background": False, + }, + }, + ] + stream = CustomStreamWrapper( + completion_stream=iter([iterator.chunk_parser(event) for event in events]), + model="gpt-5.6", + custom_llm_provider="azure", + logging_obj=MagicMock(), + ) + + chunks = [chunk.model_dump() for chunk in stream] + + assert chunks[-1]["choices"][0]["finish_reason"] == "stop" + assert chunks[-1]["service_tier"] == "default" + assert chunks[-1]["content_filters"] == content_filters + assert "background" not in chunks[-1] + assert all("service_tier" not in chunk for chunk in chunks[:-1]) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 16a99713a06..b5d56f390a9 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4391,6 +4391,39 @@ def test_handle_anthropic_messages_response_logging_translates_bare_responses_ap assert result.usage.total_tokens == 18 # type: ignore[attr-defined] +def test_handle_anthropic_messages_response_logging_keeps_the_served_response_id(): + from openai.types.responses import ResponseOutputMessage, ResponseOutputText + + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + served_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-1", response_id="resp_upstream" + ) + logging_obj = _anthropic_messages_logging_obj() + result = logging_obj._handle_anthropic_messages_response_logging( + result=ResponsesAPIResponse( + id=served_id, + created_at=1700000000, + output=[ + ResponseOutputMessage( + id="msg-1", + type="message", + role="assistant", + status="completed", + content=[ResponseOutputText(annotations=[], text="hi", type="output_text")], + ) + ], + usage=ResponseAPIUsage(input_tokens=2, output_tokens=1, total_tokens=3), + service_tier="flex", + ) + ) + + assert isinstance(result, ModelResponse) + assert result.id == served_id, "the spend log row must keep the id the caller was served" + assert result.service_tier == "flex" + + def test_handle_anthropic_messages_response_logging_passes_model_response_through(): """Anthropic-native path already yields a ModelResponse; it must be returned unchanged.""" logging_obj = _anthropic_messages_logging_obj() From 8da43835a6be47f58a45ce2967c43f50867401c9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:26:09 -0700 Subject: [PATCH 30/49] test(ci): guard against two workflow jobs publishing one check-run name A ruleset's required status check names a check run and GitHub matches it by that name alone, so two jobs publishing the same name leave the gate unable to say which job proved it. The new code-quality check reads every workflow, expands matrix values and local reusable-workflow calls the way Actions does, and fails when one name has more than one job behind it. --- .github/workflows/test-code-quality.yml | 6 + .../check_workflow_job_name_collisions.py | 179 +++++++++++++++ .../test_workflow_job_name_collisions.py | 211 ++++++++++++++++++ 3 files changed, 396 insertions(+) create mode 100644 tests/code_coverage_tests/check_workflow_job_name_collisions.py create mode 100644 tests/code_coverage_tests/test_workflow_job_name_collisions.py diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 36fd656c231..e6d2264fbf0 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -74,6 +74,12 @@ jobs: - name: check_workflow_startup_safety run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py + - name: check_workflow_job_name_collisions + run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_job_name_collisions.py + + - name: test_workflow_job_name_collisions + run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py + - name: test_e2e_changed_gate run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py new file mode 100644 index 00000000000..42e11a3e81d --- /dev/null +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -0,0 +1,179 @@ +"""Catch two workflow jobs that publish check runs under the same name. + +A ruleset's required status check names a check run and GitHub matches it by that +name alone. When two jobs publish the same name the required context stops +mapping to the job that proves it: the commit carries two check runs under one +name and nothing says which one the ruleset required. Both being green hides the +clash completely, so the context quietly stops meaning what the ruleset intended. + +`.github/workflows/auto-close-duplicates.yml` shipped a job id `test` while +`.github/workflows/test-mcp.yml` already published the required `test` context, +and commit ed5761daef4ae17152446d182c860630c38b7268 carried both check runs. +This invariant has to be enforced here because CI cannot enforce it on itself. + +A job publishes its `name:` when it sets one and its job id otherwise. Two shapes +expand that further. A `${{ matrix.key }}` reference publishes one check run per +value the matrix supplies, so two shard lists that overlap collide even though +their templates read differently. A job calling a local reusable workflow +publishes one check run per job of the callee, named ` / `, which +is why a caller's name never collides with a plain job that happens to match it. +A reference nothing resolves stays in the string, so two jobs carrying the same +unresolved template still compare equal and their collision is still caught. +""" + +import itertools +import operator +import re +import sys +from collections.abc import Iterator, Mapping, Sequence +from pathlib import Path +from typing import Final + +import yaml +from pydantic import BaseModel, Field, ValidationError + +REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent +WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows" +MATRIX_REF: Final = re.compile(r"\$\{\{\s*matrix\.(?P[\w-]+)\s*\}\}") +LOCAL_CALL_PREFIX: Final = "./" + + +class CheckRunNameCollision(Exception): + pass + + +class Job(BaseModel): + name: str | None = None + uses: str | None = None + strategy: Mapping[str, object] = Field(default_factory=dict) + + +class Workflow(BaseModel): + jobs: Mapping[str, Job] = Field(default_factory=dict) + + +def parse(text: str) -> tuple[Workflow, object] | None: + """The workflow plus its raw `on:` value, or None when the file is not a workflow.""" + parsed: Final = yaml.safe_load(text) + if not isinstance(parsed, dict): + return None + try: + return Workflow.model_validate(parsed), parsed.get(True, parsed.get("on")) + except ValidationError: + return None + + +def events(raw_on: object) -> frozenset[str]: + if isinstance(raw_on, Mapping): + return frozenset(str(key) for key in raw_on) + if isinstance(raw_on, str): + return frozenset({raw_on}) + if isinstance(raw_on, Sequence): + return frozenset(str(event) for event in raw_on) + return frozenset() + + +def publishes_check_runs(raw_on: object) -> bool: + """A `workflow_call`-only workflow posts its check runs through callers, never itself.""" + return events(raw_on) != frozenset({"workflow_call"}) + + +def matrix_values(job: Job, key: str) -> tuple[str, ...]: + matrix: Final = job.strategy.get("matrix") + if not isinstance(matrix, Mapping): + return () + listed: Final = matrix.get(key) + rows: Final = matrix.get("include") + from_list: Final = ( + tuple(str(value) for value in listed if isinstance(value, (str, int, float))) + if isinstance(listed, Sequence) and not isinstance(listed, str) + else () + ) + from_rows: Final = ( + tuple(str(row[key]) for row in rows if isinstance(row, Mapping) and isinstance(row.get(key), (str, int, float))) + if isinstance(rows, Sequence) and not isinstance(rows, str) + else () + ) + return tuple(dict.fromkeys(from_list + from_rows)) + + +def substituted(template: str, values: Mapping[str, str]) -> str: + return MATRIX_REF.sub(lambda ref: values.get(ref.group("key"), ref.group(0)), template) + + +def expand(template: str, job: Job) -> tuple[str, ...]: + keys: Final = tuple(dict.fromkeys(ref.group("key") for ref in MATRIX_REF.finditer(template))) + candidates: Final = tuple((key, matrix_values(job, key)) for key in keys) + resolvable: Final = tuple((key, values) for key, values in candidates if values) + if not resolvable: + return (template,) + return tuple( + substituted(template, dict(zip((key for key, _ in resolvable), combination))) + for combination in itertools.product(*(values for _, values in resolvable)) + ) + + +def callee_path(job: Job) -> str | None: + if job.uses is None or not job.uses.startswith(LOCAL_CALL_PREFIX): + return None + return job.uses[len(LOCAL_CALL_PREFIX) :].split("@")[0] + + +def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow]) -> tuple[str, ...]: + prefixes: Final = expand(job.name or job_id, job) + callee: Final = workflows.get(callee_path(job) or "") + if callee is None: + return prefixes + suffixes: Final = tuple( + name + for callee_id, callee_job in callee.jobs.items() + for name in expand(callee_job.name or callee_id, callee_job) + ) + return tuple(f"{prefix} / {suffix}" for prefix in prefixes for suffix in suffixes) + + +def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]: + parsed: Final = {rel: entry for rel, text in sources.items() if (entry := parse(text)) is not None} + workflows: Final = {rel: workflow for rel, (workflow, _) in parsed.items()} + for rel, (workflow, raw_on) in parsed.items(): + if not publishes_check_runs(raw_on): + continue + for job_id, job in workflow.jobs.items(): + for name in job_names(job_id, job, workflows): + yield name, f"{rel} job `{job_id}`" + + +def owners_by_name(sources: Mapping[str, str]) -> Iterator[tuple[str, tuple[str, ...]]]: + for name, pairs in itertools.groupby(sorted(published(sources)), key=operator.itemgetter(0)): + yield name, tuple(owner for _, owner in pairs) + + +def collisions(sources: Mapping[str, str]) -> tuple[str, ...]: + return tuple( + f"`{name}` is published by {len(owners)} jobs: {', '.join(owners)}. A required status check matching " + f"that name cannot say which job proves it; give one of them a distinct `name:` or job id." + for name, owners in owners_by_name(sources) + if len(owners) > 1 + ) + + +def workflow_sources() -> Mapping[str, str]: + """Repo-relative posix paths to text, the keys `uses: ./...` resolves against.""" + return {path.relative_to(REPO_ROOT).as_posix(): path.read_text() for path in sorted(WORKFLOWS_DIR.glob("*.y*ml"))} + + +def main() -> None: + sources: Final = workflow_sources() + found: Final = collisions(sources) + if found: + raise CheckRunNameCollision("Check-run names are not unique:\n - " + "\n - ".join(found)) + + print(f"Check-run names are unique across {len(sources)} workflows") + + +if __name__ == "__main__": + try: + main() + except CheckRunNameCollision as exc: + print(f"ERROR: {exc}", file=sys.stderr) + sys.exit(1) diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py new file mode 100644 index 00000000000..35c1480ecd8 --- /dev/null +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -0,0 +1,211 @@ +import sys +from pathlib import Path +from typing import Final + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from check_workflow_job_name_collisions import callee_path, collisions, parse, workflow_sources # noqa: E402 + + +def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None: + assert collisions(workflow_sources()) == () + + +def test_every_workflow_in_the_repo_parses_into_jobs() -> None: + unparsed: Final = tuple( + rel for rel, source in workflow_sources().items() if (entry := parse(source)) is None or not entry[0].jobs + ) + + assert unparsed == () + + +def test_every_local_reusable_call_in_the_repo_resolves_to_a_workflow() -> None: + sources: Final = workflow_sources() + parsed: Final = tuple(entry for text in sources.values() if (entry := parse(text)) is not None) + unresolved: Final = tuple( + job.uses + for workflow, _ in parsed + for job in workflow.jobs.values() + if (callee := callee_path(job)) is not None and callee not in sources + ) + + assert unresolved == () + + +def test_two_jobs_falling_back_to_the_same_job_id_collide() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`test` is published by 2 jobs" in found[0] + assert "a.yml job `test`" in found[0] and "b.yml job `test`" in found[0] + + +def test_an_explicit_name_overrides_the_job_id_and_clears_the_collision() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n test:\n name: Sweep tests\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + + +def test_an_explicit_name_matching_another_job_id_collides() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n sweep:\n name: test\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`test` is published by 2 jobs" in found[0] + + +def test_two_callers_of_one_reusable_workflow_collide_on_a_shared_matrix_value() -> None: + base: Final = "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n" + caller: Final = ( + "on: pull_request\n" + "jobs:\n" + " {job}:\n" + " name: ${{{{ matrix.shard }}}}\n" + " uses: ./.github/workflows/base.yml\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: {shard}\n" + ) + sources: Final = { + ".github/workflows/base.yml": base, + "unit.yml": caller.format(job="unit", shard="proxy-auth"), + "proxy-db.yml": caller.format(job="proxy-db", shard="proxy-auth"), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`proxy-auth / Run tests` is published by 2 jobs" in found[0] + + +def test_distinct_matrix_values_through_one_reusable_workflow_do_not_collide() -> None: + base: Final = "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n" + caller: Final = ( + "on: pull_request\n" + "jobs:\n" + " {job}:\n" + " name: ${{{{ matrix.shard }}}}\n" + " uses: ./.github/workflows/base.yml\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: {shard}\n" + ) + sources: Final = { + ".github/workflows/base.yml": base, + "unit.yml": caller.format(job="unit", shard="proxy-auth"), + "proxy-db.yml": caller.format(job="proxy-db", shard="budgets"), + } + + assert collisions(sources) == () + + +def test_a_reusable_caller_does_not_collide_with_a_plain_job_of_the_same_name() -> None: + sources: Final = { + ".github/workflows/base.yml": ( + "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n" + ), + "unit.yml": ( + "on: pull_request\n" + "jobs:\n" + " unit:\n" + " name: ${{ matrix.shard }}\n" + " uses: ./.github/workflows/base.yml\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: proxy-behavior\n" + ), + "postgres.yml": ( + "on: pull_request\n" + "jobs:\n" + " postgres:\n" + " name: ${{ matrix.shard }}\n" + " runs-on: ubuntu-latest\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: proxy-behavior\n" + ), + } + + assert collisions(sources) == () + + +def test_a_workflow_call_only_workflow_publishes_nothing_of_its_own() -> None: + sources: Final = { + "base.yml": "on:\n workflow_call:\njobs:\n run:\n runs-on: ubuntu-latest\n", + "other.yml": "on:\n workflow_call:\njobs:\n run:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + + +def test_a_workflow_call_workflow_that_also_runs_on_pull_request_still_publishes() -> None: + sources: Final = { + "base.yml": "on:\n workflow_call:\n pull_request:\njobs:\n run:\n runs-on: ubuntu-latest\n", + "other.yml": "on: pull_request\njobs:\n run:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`run` is published by 2 jobs" in found[0] + + +def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\n" + "jobs:\n" + " build:\n" + " name: Analyze (${{ matrix.language }})\n" + " runs-on: ubuntu-latest\n" + " strategy:\n" + " matrix:\n" + " language: [python, go]\n" + ), + "b.yml": "on: pull_request\njobs:\n go:\n name: Analyze (go)\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`Analyze (go)` is published by 2 jobs" in found[0] + + +def test_two_jobs_sharing_one_unresolvable_template_still_collide() -> None: + template: Final = ( + "on: pull_request\njobs:\n {job}:\n name: ${{{{ matrix.shard }}}}\n runs-on: ubuntu-latest\n" + ) + sources: Final = { + "a.yml": template.format(job="one"), + "b.yml": template.format(job="two"), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "is published by 2 jobs" in found[0] + + +def test_a_file_that_is_not_a_workflow_is_ignored() -> None: + sources: Final = { + "notes.yml": "just a string\n", + "a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () From 44b7a76439f1896736e61679b4cd0b5e66dd297a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:26:38 -0700 Subject: [PATCH 31/49] fix(ci): stop the auto-close duplicates job colliding with the required test check Its job id was `test`, the same check-run name test-mcp.yml publishes and guard-internal-staging requires, so a commit touching the sweep carried two check runs called `test`. Renaming it to `sweep-tests` leaves the required context with exactly one job behind it. --- .github/workflows/auto-close-duplicates.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/auto-close-duplicates.yml b/.github/workflows/auto-close-duplicates.yml index d8256917805..9a362b4c85c 100644 --- a/.github/workflows/auto-close-duplicates.yml +++ b/.github/workflows/auto-close-duplicates.yml @@ -22,7 +22,7 @@ on: permissions: {} jobs: - test: + sweep-tests: if: github.event_name == 'pull_request' runs-on: ubuntu-latest timeout-minutes: 5 From baee7d8175a1e46085a477f45257dc8b9006b4c1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:47:23 -0700 Subject: [PATCH 32/49] fix: evaluate job name expressions in the check-run collision guard The guard only substituted a bare `${{ matrix.key }}`, so any name built from a larger expression stayed in the string as its own template. `_test-unit-base.yml` names its job with a ternary over `format()`, which meant every shard published an opaque name and 23 of the 33 required contexts, all of them ` / Run tests`, were invisible to the very check meant to protect them. Job names are now evaluated per matrix combination over the pieces a name can hold: string literals, `matrix.`, `format()`, `==` and `!=`, and the ` && || ` idiom. All 33 required contexts now resolve, and nothing in the repo leaves an expression unresolved. An expression the evaluator does not understand still falls back to its verbatim template, so two jobs sharing one stays a collision. The test also drops its `sys.path.insert`, which the test-quality budget counts under TQ003; pytest already puts the file's own directory on the path. --- .../check_workflow_job_name_collisions.py | 109 +++++++++++++++--- .../test_workflow_job_name_collisions.py | 63 +++++++++- 2 files changed, 152 insertions(+), 20 deletions(-) diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py index 42e11a3e81d..849a906c3b8 100644 --- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -12,13 +12,17 @@ and commit ed5761daef4ae17152446d182c860630c38b7268 carried both check runs. This invariant has to be enforced here because CI cannot enforce it on itself. A job publishes its `name:` when it sets one and its job id otherwise. Two shapes -expand that further. A `${{ matrix.key }}` reference publishes one check run per -value the matrix supplies, so two shard lists that overlap collide even though -their templates read differently. A job calling a local reusable workflow -publishes one check run per job of the callee, named ` / `, which -is why a caller's name never collides with a plain job that happens to match it. -A reference nothing resolves stays in the string, so two jobs carrying the same -unresolved template still compare equal and their collision is still caught. +expand that further. A name carrying `${{ ... }}` publishes one check run per +matrix combination it reads, so two shard lists that overlap collide even though +their templates read differently. Each expression is evaluated per combination +over the pieces a job name can hold: string literals, `matrix.`, `format()`, +`==` and `!=`, and the ` && || ` idiom, which is how the shards +reach their real ` / Run tests` names rather than staying opaque. A job +calling a local reusable workflow publishes one check run per job of the callee, +named ` / `, which is why a caller's name never collides with a +plain job that happens to match it. An expression nothing resolves stays in the +string, so two jobs carrying the same unresolved template still compare equal and +their collision is still caught. """ import itertools @@ -27,6 +31,7 @@ import re import sys from collections.abc import Iterator, Mapping, Sequence from pathlib import Path +from types import MappingProxyType from typing import Final import yaml @@ -34,7 +39,13 @@ from pydantic import BaseModel, Field, ValidationError REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows" -MATRIX_REF: Final = re.compile(r"\$\{\{\s*matrix\.(?P[\w-]+)\s*\}\}") +EXPRESSION: Final = re.compile(r"\$\{\{(?P.*?)\}\}", re.DOTALL) +MATRIX_KEY: Final = re.compile(r"matrix\.(?P[\w-]+)") +MATRIX_REF: Final = re.compile(r"^matrix\.(?P[\w-]+)$") +LITERAL: Final = re.compile(r"^'(?P[^']*)'$") +FORMAT_CALL: Final = re.compile(r"^format\((?P.*)\)$", re.DOTALL) +COMPARISON: Final = re.compile(r"^(?P.+?)\s*(?P==|!=)\s*(?P.+)$", re.DOTALL) +NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({}) LOCAL_CALL_PREFIX: Final = "./" @@ -97,19 +108,85 @@ def matrix_values(job: Job, key: str) -> tuple[str, ...]: return tuple(dict.fromkeys(from_list + from_rows)) -def substituted(template: str, values: Mapping[str, str]) -> str: - return MATRIX_REF.sub(lambda ref: values.get(ref.group("key"), ref.group(0)), template) +def scanned(state: tuple[int, bool], char: str) -> tuple[int, bool]: + depth, quoted = state + if char == "'": + return depth, not quoted + if quoted: + return depth, quoted + return depth + int(char == "(") - int(char == ")"), quoted + + +def split_outside(text: str, token: str) -> tuple[str, ...]: + """`text` cut on every `token` that sits outside quotes and parentheses.""" + states: Final = tuple(itertools.accumulate(text, scanned, initial=(0, False))) + cuts: Final = tuple( + index + for index in range(len(text) - len(token) + 1) + if text.startswith(token, index) and states[index] == (0, False) + ) + starts: Final = (0, *(cut + len(token) for cut in cuts)) + return tuple(text[start:end] for start, end in zip(starts, (*cuts, len(text)))) + + +def value_of(text: str, values: Mapping[str, str]) -> str | None: + expression: Final = text.strip() + literal: Final = LITERAL.match(expression) + if literal is not None: + return literal.group("text") + reference: Final = MATRIX_REF.match(expression) + if reference is not None: + return values.get(reference.group("key")) + call: Final = FORMAT_CALL.match(expression) + if call is None: + return None + arguments: Final = tuple(value_of(part, values) for part in split_outside(call.group("args"), ",")) + resolved: Final = tuple(argument for argument in arguments if argument is not None) + if not resolved or len(resolved) != len(arguments): + return None + return resolved[0].format(*resolved[1:]) + + +def holds(condition: str, values: Mapping[str, str]) -> bool | None: + comparison: Final = COMPARISON.match(condition.strip()) + if comparison is None: + return None + left: Final = value_of(comparison.group("left"), values) + right: Final = value_of(comparison.group("right"), values) + if left is None or right is None: + return None + return (left == right) == (comparison.group("operator") == "==") + + +def evaluate(body: str, values: Mapping[str, str]) -> str | None: + """The single string this expression yields, or None when its shape is not understood.""" + branches: Final = tuple(split_outside(alternative, "&&") for alternative in split_outside(body, "||")) + outcomes: Final = tuple(tuple(holds(part, values) for part in branch[:-1]) for branch in branches) + if any(outcome is None for branch in outcomes for outcome in branch): + return None + taken: Final = next((branch[-1] for branch, outcome in zip(branches, outcomes) if all(outcome)), None) + return None if taken is None else value_of(taken, values) + + +def resolved_span(span: re.Match[str], values: Mapping[str, str]) -> str: + substitution: Final = evaluate(span.group("body"), values) + return span.group(0) if substitution is None else substitution + + +def rendered(template: str, values: Mapping[str, str]) -> str: + return EXPRESSION.sub(lambda span: resolved_span(span, values), template) def expand(template: str, job: Job) -> tuple[str, ...]: - keys: Final = tuple(dict.fromkeys(ref.group("key") for ref in MATRIX_REF.finditer(template))) - candidates: Final = tuple((key, matrix_values(job, key)) for key in keys) - resolvable: Final = tuple((key, values) for key, values in candidates if values) + keys: Final = tuple(dict.fromkeys(ref.group("key") for ref in MATRIX_KEY.finditer(template))) + resolvable: Final = tuple((key, values) for key in keys if (values := matrix_values(job, key))) if not resolvable: - return (template,) + return (rendered(template, NO_MATRIX),) return tuple( - substituted(template, dict(zip((key for key, _ in resolvable), combination))) - for combination in itertools.product(*(values for _, values in resolvable)) + dict.fromkeys( + rendered(template, dict(zip((key for key, _ in resolvable), combination))) + for combination in itertools.product(*(values for _, values in resolvable)) + ) ) diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py index 35c1480ecd8..96820d50040 100644 --- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -1,10 +1,30 @@ -import sys -from pathlib import Path from typing import Final -sys.path.insert(0, str(Path(__file__).resolve().parent)) +from check_workflow_job_name_collisions import callee_path, collisions, parse, published, workflow_sources -from check_workflow_job_name_collisions import callee_path, collisions, parse, workflow_sources # noqa: E402 +REUSABLE_BASE: Final = """on: + workflow_call: +jobs: + run: + name: >- + ${{ matrix.python-version == '3.12' && 'Run tests' + || format('Run tests (Python {0})', matrix.python-version) }} + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12", "3.13"] +""" + +SHARD_CALLER: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.shard }} + uses: ./.github/workflows/base.yml + strategy: + matrix: + include: + - shard: core-utils +""" def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None: @@ -209,3 +229,38 @@ def test_a_file_that_is_not_a_workflow_is_ignored() -> None: } assert collisions(sources) == () + + +def test_a_conditional_name_expands_to_the_branch_each_matrix_value_takes() -> None: + names: Final = frozenset( + name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER}) + ) + + assert names == frozenset({"core-utils / Run tests", "core-utils / Run tests (Python 3.13)"}) + + +def test_a_conditional_name_never_publishes_the_branch_its_condition_rules_out() -> None: + names: Final = frozenset( + name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER}) + ) + + assert "core-utils / Run tests (Python 3.12)" not in names + + +def test_a_conditional_reusable_name_collides_with_a_plain_job_publishing_the_same_name() -> None: + sources: Final = { + ".github/workflows/base.yml": REUSABLE_BASE, + "unit.yml": SHARD_CALLER, + "postgres.yml": ("on: pull_request\njobs:\n legacy:\n name: core-utils / Run tests\n"), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`core-utils / Run tests` is published by 2 jobs" in found[0] + + +def test_every_workflow_in_the_repo_resolves_every_expression_in_its_job_names() -> None: + unresolved: Final = tuple(f"{owner}: {name}" for name, owner in published(workflow_sources()) if "${{" in name) + + assert unresolved == () From d6a727fe0f52e8d280235d1cc98c43462643d150 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:02:32 -0700 Subject: [PATCH 33/49] fix: keep matrix include rows whole when expanding job names The guard read each matrix key's values independently and crossed them, so a job name reading two keys off one include row published pairs no job ever runs, which could fail a valid workflow on a required check It now builds the combinations GitHub builds: the listed keys crossed, each include row folded into the combinations it overwrites nothing in, and a row that fits nowhere standing on its own --- .../check_workflow_job_name_collisions.py | 73 +++++++++++++------ .../test_workflow_job_name_collisions.py | 46 ++++++++++++ 2 files changed, 95 insertions(+), 24 deletions(-) diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py index 849a906c3b8..c683de7f08b 100644 --- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -13,8 +13,9 @@ This invariant has to be enforced here because CI cannot enforce it on itself. A job publishes its `name:` when it sets one and its job id otherwise. Two shapes expand that further. A name carrying `${{ ... }}` publishes one check run per -matrix combination it reads, so two shard lists that overlap collide even though -their templates read differently. Each expression is evaluated per combination +combination the matrix produces, with each `include` row's values held together +rather than crossed with the other rows', so two shard lists that overlap collide +even though their templates read differently. Each expression is evaluated per combination over the pieces a job name can hold: string literals, `matrix.`, `format()`, `==` and `!=`, and the ` && || ` idiom, which is how the shards reach their real ` / Run tests` names rather than staying opaque. A job @@ -40,12 +41,13 @@ from pydantic import BaseModel, Field, ValidationError REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows" EXPRESSION: Final = re.compile(r"\$\{\{(?P.*?)\}\}", re.DOTALL) -MATRIX_KEY: Final = re.compile(r"matrix\.(?P[\w-]+)") MATRIX_REF: Final = re.compile(r"^matrix\.(?P[\w-]+)$") LITERAL: Final = re.compile(r"^'(?P[^']*)'$") FORMAT_CALL: Final = re.compile(r"^format\((?P.*)\)$", re.DOTALL) COMPARISON: Final = re.compile(r"^(?P.+?)\s*(?P==|!=)\s*(?P.+)$", re.DOTALL) NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({}) +SCALAR: Final = (str, int, float) +MATRIX_DIRECTIVES: Final = frozenset({"include", "exclude"}) LOCAL_CALL_PREFIX: Final = "./" @@ -89,23 +91,52 @@ def publishes_check_runs(raw_on: object) -> bool: return events(raw_on) != frozenset({"workflow_call"}) -def matrix_values(job: Job, key: str) -> tuple[str, ...]: +def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...]: + return tuple( + (str(key), tuple(str(value) for value in values if isinstance(value, SCALAR))) + for key, values in matrix.items() + if str(key) not in MATRIX_DIRECTIVES and isinstance(values, Sequence) and not isinstance(values, str) + ) + + +def include_rows(matrix: Mapping[str, object]) -> tuple[Mapping[str, str], ...]: + rows: Final = matrix.get("include") + if not isinstance(rows, Sequence) or isinstance(rows, str): + return () + return tuple( + MappingProxyType({str(key): str(value) for key, value in row.items() if isinstance(value, SCALAR)}) + for row in rows + if isinstance(row, Mapping) + ) + + +def extends(row: Mapping[str, str], combination: Mapping[str, str]) -> bool: + """GitHub folds an `include` row into a combination only where it overwrites no listed value.""" + return all(combination[key] == value for key, value in row.items() if key in combination) + + +def extended(combination: Mapping[str, str], rows: Sequence[Mapping[str, str]]) -> Mapping[str, str]: + additions: Final = {key: value for row in rows if extends(row, combination) for key, value in row.items()} + return MappingProxyType({**combination, **additions}) + + +def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...]: + """One mapping per job the matrix produces, each `include` row's values staying together.""" matrix: Final = job.strategy.get("matrix") if not isinstance(matrix, Mapping): return () - listed: Final = matrix.get(key) - rows: Final = matrix.get("include") - from_list: Final = ( - tuple(str(value) for value in listed if isinstance(value, (str, int, float))) - if isinstance(listed, Sequence) and not isinstance(listed, str) + listed: Final = listed_values(matrix) + rows: Final = include_rows(matrix) + crossed: Final = ( + tuple( + MappingProxyType(dict(zip((key for key, _ in listed), values))) + for values in itertools.product(*(values for _, values in listed)) + ) + if listed else () ) - from_rows: Final = ( - tuple(str(row[key]) for row in rows if isinstance(row, Mapping) and isinstance(row.get(key), (str, int, float))) - if isinstance(rows, Sequence) and not isinstance(rows, str) - else () - ) - return tuple(dict.fromkeys(from_list + from_rows)) + standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in crossed)) + return (*(extended(combination, rows) for combination in crossed), *standalone) def scanned(state: tuple[int, bool], char: str) -> tuple[int, bool]: @@ -178,16 +209,10 @@ def rendered(template: str, values: Mapping[str, str]) -> str: def expand(template: str, job: Job) -> tuple[str, ...]: - keys: Final = tuple(dict.fromkeys(ref.group("key") for ref in MATRIX_KEY.finditer(template))) - resolvable: Final = tuple((key, values) for key in keys if (values := matrix_values(job, key))) - if not resolvable: + combinations: Final = matrix_combinations(job) + if not combinations: return (rendered(template, NO_MATRIX),) - return tuple( - dict.fromkeys( - rendered(template, dict(zip((key for key, _ in resolvable), combination))) - for combination in itertools.product(*(values for _, values in resolvable)) - ) - ) + return tuple(dict.fromkeys(rendered(template, values) for values in combinations)) def callee_path(job: Job) -> str | None: diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py index 96820d50040..359072326bd 100644 --- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -27,6 +27,33 @@ jobs: """ +CORRELATED_ROWS: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.shard }} on ${{ matrix.test-path }} + runs-on: ubuntu-latest + strategy: + matrix: + include: + - shard: core-utils + test-path: tests/core + - shard: proxy + test-path: tests/proxy +""" + +LISTED_PLUS_ROW: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.python-version }} ${{ matrix.label }} + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12", "3.13"] + include: + - label: fast +""" + + def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None: assert collisions(workflow_sources()) == () @@ -260,6 +287,25 @@ def test_a_conditional_reusable_name_collides_with_a_plain_job_publishing_the_sa assert "`core-utils / Run tests` is published by 2 jobs" in found[0] +def test_a_name_reading_two_matrix_keys_publishes_only_the_pairs_each_include_row_holds() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": CORRELATED_ROWS})) + + assert names == frozenset({"core-utils on tests/core", "proxy on tests/proxy"}) + + +def test_a_name_reading_two_matrix_keys_never_publishes_a_pair_no_include_row_holds() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": CORRELATED_ROWS})) + + assert "core-utils on tests/proxy" not in names + assert "proxy on tests/core" not in names + + +def test_an_include_row_carrying_no_listed_key_extends_every_listed_combination() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": LISTED_PLUS_ROW})) + + assert names == frozenset({"3.12 fast", "3.13 fast"}) + + def test_every_workflow_in_the_repo_resolves_every_expression_in_its_job_names() -> None: unresolved: Final = tuple(f"{owner}: {name}" for name, owner in published(workflow_sources()) if "${{" in name) From 9f379b36b9c484ff045fadfcec213f47857670db Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:18:18 -0700 Subject: [PATCH 34/49] refactor: return the collision check's failure instead of raising it The checker raised a custom exception and caught it two lines down in the same module, which is the throw-then-catch the repo's coding guide rules out. `main` now prints the same message and returns the exit code, so the collision list stays a value the whole way out --- .../check_workflow_job_name_collisions.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py index c683de7f08b..9dc13c2ea0c 100644 --- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -51,10 +51,6 @@ MATRIX_DIRECTIVES: Final = frozenset({"include", "exclude"}) LOCAL_CALL_PREFIX: Final = "./" -class CheckRunNameCollision(Exception): - pass - - class Job(BaseModel): name: str | None = None uses: str | None = None @@ -264,18 +260,16 @@ def workflow_sources() -> Mapping[str, str]: return {path.relative_to(REPO_ROOT).as_posix(): path.read_text() for path in sorted(WORKFLOWS_DIR.glob("*.y*ml"))} -def main() -> None: +def main() -> int: sources: Final = workflow_sources() found: Final = collisions(sources) if found: - raise CheckRunNameCollision("Check-run names are not unique:\n - " + "\n - ".join(found)) + print("ERROR: Check-run names are not unique:\n - " + "\n - ".join(found), file=sys.stderr) + return 1 print(f"Check-run names are unique across {len(sources)} workflows") + return 0 if __name__ == "__main__": - try: - main() - except CheckRunNameCollision as exc: - print(f"ERROR: {exc}", file=sys.stderr) - sys.exit(1) + sys.exit(main()) From 5bbc83e3defa48db50546aa81ba60a9cdb45ed58 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 00:04:31 -0700 Subject: [PATCH 35/49] fix: model the check-run names GitHub really publishes The collision sweep read a job's name as its `name:` or bare job id, which is wrong for a matrix job that sets no name: GitHub publishes `build (3.12)`, one per combination. That missed real duplicates and invented ones that don't exist. It also crossed every matrix value while ignoring `exclude`, so it checked combinations no job ever runs. Four smaller gaps went with it. Boolean matrix values reached a name as `True` rather than `true`. A `format()` whose arguments cannot fill its placeholders raised straight out of the script instead of leaving the name unresolved. A job calling a reusable workflow only ever chained one level, and a call outside the repo fell back to the caller's own name, which GitHub never posts. A job whose `name:` was not a string failed validation and silently dropped every job in that file, so the sweep now renders any scalar and reports a file it cannot read instead of skipping it. --- .../check_workflow_job_name_collisions.py | 174 +++++++++++----- .../test_workflow_job_name_collisions.py | 196 +++++++++++++++++- 2 files changed, 319 insertions(+), 51 deletions(-) diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py index 9dc13c2ea0c..4eef3df1c11 100644 --- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -11,19 +11,27 @@ clash completely, so the context quietly stops meaning what the ruleset intended and commit ed5761daef4ae17152446d182c860630c38b7268 carried both check runs. This invariant has to be enforced here because CI cannot enforce it on itself. -A job publishes its `name:` when it sets one and its job id otherwise. Two shapes -expand that further. A name carrying `${{ ... }}` publishes one check run per -combination the matrix produces, with each `include` row's values held together -rather than crossed with the other rows', so two shard lists that overlap collide -even though their templates read differently. Each expression is evaluated per combination -over the pieces a job name can hold: string literals, `matrix.`, `format()`, -`==` and `!=`, and the ` && || ` idiom, which is how the shards -reach their real ` / Run tests` names rather than staying opaque. A job -calling a local reusable workflow publishes one check run per job of the callee, -named ` / `, which is why a caller's name never collides with a -plain job that happens to match it. An expression nothing resolves stays in the -string, so two jobs carrying the same unresolved template still compare equal and -their collision is still caught. +A job publishes its `name:` when it sets one, and otherwise its job id plus the +values of the combination it runs, the way GitHub writes `build (3.12)`. A name +carrying `${{ ... }}` publishes one check run per combination the matrix +produces: `exclude` rows drop combinations before `include` rows fold into the +survivors, and each `include` row's values stay together rather than crossing +with the other rows', so two shard lists that overlap collide even though their +templates read differently. Each expression is evaluated per combination over the +pieces a job name can hold: string literals, `matrix.`, `format()`, `==` and +`!=`, and the ` && || ` idiom, which is how the shards reach their +real ` / Run tests` names rather than staying opaque. An expression +nothing resolves stays in the string, so two jobs carrying the same unresolved +template still compare equal and their collision is still caught. + +A job calling a local reusable workflow publishes one check run per job of the +callee, named ` / ` and chained through however many levels of +local calls it takes, which is why a caller's name never collides with a plain +job that happens to match it. A call this sweep cannot follow, to a workflow +outside the repo or to one that is not there, publishes nothing rather than a +name GitHub never posts. A file under `.github/workflows/` that does not read as +a workflow at all is reported rather than skipped, since skipping it silently +would hide every job it holds. """ import itertools @@ -31,6 +39,7 @@ import operator import re import sys from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass from pathlib import Path from types import MappingProxyType from typing import Final @@ -46,13 +55,19 @@ LITERAL: Final = re.compile(r"^'(?P[^']*)'$") FORMAT_CALL: Final = re.compile(r"^format\((?P.*)\)$", re.DOTALL) COMPARISON: Final = re.compile(r"^(?P.+?)\s*(?P==|!=)\s*(?P.+)$", re.DOTALL) NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({}) +NO_CALLERS: Final[frozenset[str]] = frozenset() SCALAR: Final = (str, int, float) MATRIX_DIRECTIVES: Final = frozenset({"include", "exclude"}) LOCAL_CALL_PREFIX: Final = "./" +@dataclass(frozen=True, slots=True) +class Unreadable: + reason: str + + class Job(BaseModel): - name: str | None = None + name: object = None uses: str | None = None strategy: Mapping[str, object] = Field(default_factory=dict) @@ -61,15 +76,20 @@ class Workflow(BaseModel): jobs: Mapping[str, Job] = Field(default_factory=dict) -def parse(text: str) -> tuple[Workflow, object] | None: - """The workflow plus its raw `on:` value, or None when the file is not a workflow.""" - parsed: Final = yaml.safe_load(text) +def scalar_text(value: object) -> str: + """A YAML scalar the way GitHub renders it, so `true` never reaches a name as `True`.""" + return str(value).lower() if isinstance(value, bool) else str(value) + + +def parse(source: str) -> tuple[Workflow, object] | Unreadable: + """The workflow plus its raw `on:` value, or why the file does not read as one.""" + parsed: Final = yaml.safe_load(source) if not isinstance(parsed, dict): - return None + return Unreadable("its top level is not a mapping of workflow keys") try: return Workflow.model_validate(parsed), parsed.get(True, parsed.get("on")) - except ValidationError: - return None + except ValidationError as error: + return Unreadable(f"{error.error_count()} of its job definitions have a shape GitHub would reject") def events(raw_on: object) -> frozenset[str]: @@ -89,23 +109,28 @@ def publishes_check_runs(raw_on: object) -> bool: def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...]: return tuple( - (str(key), tuple(str(value) for value in values if isinstance(value, SCALAR))) + (str(key), tuple(scalar_text(value) for value in values if isinstance(value, SCALAR))) for key, values in matrix.items() if str(key) not in MATRIX_DIRECTIVES and isinstance(values, Sequence) and not isinstance(values, str) ) -def include_rows(matrix: Mapping[str, object]) -> tuple[Mapping[str, str], ...]: - rows: Final = matrix.get("include") +def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...]: + rows: Final = matrix.get(directive) if not isinstance(rows, Sequence) or isinstance(rows, str): return () return tuple( - MappingProxyType({str(key): str(value) for key, value in row.items() if isinstance(value, SCALAR)}) + MappingProxyType({str(key): scalar_text(value) for key, value in row.items() if isinstance(value, SCALAR)}) for row in rows if isinstance(row, Mapping) ) +def drops(row: Mapping[str, str], combination: Mapping[str, str]) -> bool: + """GitHub removes a combination that carries every value one `exclude` row names.""" + return all(combination.get(key) == value for key, value in row.items()) + + def extends(row: Mapping[str, str], combination: Mapping[str, str]) -> bool: """GitHub folds an `include` row into a combination only where it overwrites no listed value.""" return all(combination[key] == value for key, value in row.items() if key in combination) @@ -116,23 +141,29 @@ def extended(combination: Mapping[str, str], rows: Sequence[Mapping[str, str]]) return MappingProxyType({**combination, **additions}) +def crossed_values(listed: Sequence[tuple[str, tuple[str, ...]]]) -> tuple[Mapping[str, str], ...]: + if not listed: + return () + return tuple( + MappingProxyType(dict(zip((key for key, _ in listed), values))) + for values in itertools.product(*(values for _, values in listed)) + ) + + def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...]: - """One mapping per job the matrix produces, each `include` row's values staying together.""" + """One mapping per job the matrix produces, `exclude` applied before `include` as GitHub does.""" matrix: Final = job.strategy.get("matrix") if not isinstance(matrix, Mapping): return () - listed: Final = listed_values(matrix) - rows: Final = include_rows(matrix) - crossed: Final = ( - tuple( - MappingProxyType(dict(zip((key for key, _ in listed), values))) - for values in itertools.product(*(values for _, values in listed)) - ) - if listed - else () + rows: Final = directive_rows(matrix, "include") + dropped: Final = directive_rows(matrix, "exclude") + kept: Final = tuple( + combination + for combination in crossed_values(listed_values(matrix)) + if not any(drops(row, combination) for row in dropped) ) - standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in crossed)) - return (*(extended(combination, rows) for combination in crossed), *standalone) + standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in kept)) + return (*(extended(combination, rows) for combination in kept), *standalone) def scanned(state: tuple[int, bool], char: str) -> tuple[int, bool]: @@ -156,6 +187,14 @@ def split_outside(text: str, token: str) -> tuple[str, ...]: return tuple(text[start:end] for start, end in zip(starts, (*cuts, len(text)))) +def formatted(template: str, arguments: Sequence[str]) -> str | None: + """A `format()` whose placeholders the arguments cannot fill resolves to nothing, never a crash.""" + try: + return template.format(*arguments) + except (IndexError, KeyError, ValueError): + return None + + def value_of(text: str, values: Mapping[str, str]) -> str | None: expression: Final = text.strip() literal: Final = LITERAL.match(expression) @@ -171,7 +210,7 @@ def value_of(text: str, values: Mapping[str, str]) -> str | None: resolved: Final = tuple(argument for argument in arguments if argument is not None) if not resolved or len(resolved) != len(arguments): return None - return resolved[0].format(*resolved[1:]) + return formatted(resolved[0], resolved[1:]) def holds(condition: str, values: Mapping[str, str]) -> bool | None: @@ -211,27 +250,58 @@ def expand(template: str, job: Job) -> tuple[str, ...]: return tuple(dict.fromkeys(rendered(template, values) for values in combinations)) +def suffixed(job_id: str, combination: Mapping[str, str]) -> str: + """The name GitHub gives a job with no `name:`, its id plus the combination it runs.""" + return f"{job_id} ({', '.join(combination.values())})" if combination else job_id + + +def published_names(job_id: str, job: Job) -> tuple[str, ...]: + if job.name is not None: + return expand(scalar_text(job.name), job) + suffixes: Final = tuple(dict.fromkeys(suffixed(job_id, values) for values in matrix_combinations(job))) + return suffixes or (job_id,) + + def callee_path(job: Job) -> str | None: if job.uses is None or not job.uses.startswith(LOCAL_CALL_PREFIX): return None return job.uses[len(LOCAL_CALL_PREFIX) :].split("@")[0] -def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow]) -> tuple[str, ...]: - prefixes: Final = expand(job.name or job_id, job) - callee: Final = workflows.get(callee_path(job) or "") - if callee is None: +def job_names( + job_id: str, job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str] = NO_CALLERS +) -> tuple[str, ...]: + prefixes: Final = published_names(job_id, job) + if job.uses is None: return prefixes + path: Final = callee_path(job) + callee: Final = workflows.get(path or "") + if path is None or callee is None or path in callers: + return () suffixes: Final = tuple( name for callee_id, callee_job in callee.jobs.items() - for name in expand(callee_job.name or callee_id, callee_job) + for name in job_names(callee_id, callee_job, workflows, callers | {path}) ) return tuple(f"{prefix} / {suffix}" for prefix in prefixes for suffix in suffixes) +def readable(sources: Mapping[str, str]) -> Mapping[str, tuple[Workflow, object]]: + parsed: Final = {rel: parse(source) for rel, source in sources.items()} + return MappingProxyType({rel: entry for rel, entry in parsed.items() if not isinstance(entry, Unreadable)}) + + +def unreadable(sources: Mapping[str, str]) -> tuple[str, ...]: + parsed: Final = {rel: parse(source) for rel, source in sources.items()} + return tuple( + f"{rel} sits in the workflows directory but {entry.reason}, so none of its jobs were checked." + for rel, entry in sorted(parsed.items()) + if isinstance(entry, Unreadable) + ) + + def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]: - parsed: Final = {rel: entry for rel, text in sources.items() if (entry := parse(text)) is not None} + parsed: Final = readable(sources) workflows: Final = {rel: workflow for rel, (workflow, _) in parsed.items()} for rel, (workflow, raw_on) in parsed.items(): if not publishes_check_runs(raw_on): @@ -260,16 +330,26 @@ def workflow_sources() -> Mapping[str, str]: return {path.relative_to(REPO_ROOT).as_posix(): path.read_text() for path in sorted(WORKFLOWS_DIR.glob("*.y*ml"))} -def main() -> int: - sources: Final = workflow_sources() +def report(header: str, problems: Sequence[str]) -> None: + if problems: + print(f"ERROR: {header}:\n - " + "\n - ".join(problems), file=sys.stderr) + + +def exit_code(sources: Mapping[str, str]) -> int: + unread: Final = unreadable(sources) found: Final = collisions(sources) - if found: - print("ERROR: Check-run names are not unique:\n - " + "\n - ".join(found), file=sys.stderr) + report("Some workflows could not be read", unread) + report("Check-run names are not unique", found) + if unread or found: return 1 print(f"Check-run names are unique across {len(sources)} workflows") return 0 +def main() -> int: + return exit_code(workflow_sources()) + + if __name__ == "__main__": sys.exit(main()) diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py index 359072326bd..3dd8a3d4088 100644 --- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -1,6 +1,15 @@ from typing import Final -from check_workflow_job_name_collisions import callee_path, collisions, parse, published, workflow_sources +from check_workflow_job_name_collisions import ( + Unreadable, + callee_path, + collisions, + exit_code, + parse, + published, + unreadable, + workflow_sources, +) REUSABLE_BASE: Final = """on: workflow_call: @@ -53,6 +62,62 @@ jobs: - label: fast """ +NAMELESS_MATRIX: Final = """on: pull_request +jobs: + build: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12", "3.13"] +""" + +EXCLUDED_PAIR: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.os }}-${{ matrix.python-version }} + runs-on: ubuntu-latest + strategy: + matrix: + os: [ubuntu, macos] + python-version: ["3.12", "3.13"] + exclude: + - os: macos + python-version: "3.13" +""" + +EXCLUDED_KEY: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.os }}-${{ matrix.python-version }} + runs-on: ubuntu-latest + strategy: + matrix: + os: [ubuntu, macos] + python-version: ["3.12", "3.13"] + exclude: + - os: macos +""" + +BOOLEAN_MATRIX: Final = """on: pull_request +jobs: + unit: + name: cache ${{ matrix.cached }} + runs-on: ubuntu-latest + strategy: + matrix: + cached: [true, false] +""" + +UNFILLABLE_FORMAT: Final = """on: pull_request +jobs: + unit: + name: ${{ format('{0} {1}', matrix.shard) }} + runs-on: ubuntu-latest + strategy: + matrix: + shard: [core] +""" + def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None: assert collisions(workflow_sources()) == () @@ -60,7 +125,9 @@ def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None: def test_every_workflow_in_the_repo_parses_into_jobs() -> None: unparsed: Final = tuple( - rel for rel, source in workflow_sources().items() if (entry := parse(source)) is None or not entry[0].jobs + rel + for rel, source in workflow_sources().items() + if isinstance(entry := parse(source), Unreadable) or not entry[0].jobs ) assert unparsed == () @@ -68,7 +135,7 @@ def test_every_workflow_in_the_repo_parses_into_jobs() -> None: def test_every_local_reusable_call_in_the_repo_resolves_to_a_workflow() -> None: sources: Final = workflow_sources() - parsed: Final = tuple(entry for text in sources.values() if (entry := parse(text)) is not None) + parsed: Final = tuple(entry for text in sources.values() if not isinstance(entry := parse(text), Unreadable)) unresolved: Final = tuple( job.uses for workflow, _ in parsed @@ -249,15 +316,28 @@ def test_two_jobs_sharing_one_unresolvable_template_still_collide() -> None: assert "is published by 2 jobs" in found[0] -def test_a_file_that_is_not_a_workflow_is_ignored() -> None: +def test_a_file_that_is_not_a_workflow_is_reported_rather_than_skipped() -> None: sources: Final = { "notes.yml": "just a string\n", "a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", } + found: Final = unreadable(sources) + + assert len(found) == 1 + assert "notes.yml" in found[0] assert collisions(sources) == () +def test_a_workflow_holding_a_job_shape_github_would_reject_is_reported() -> None: + sources: Final = {"a.yml": "on: pull_request\njobs:\n test:\n uses: [not, a, string]\n"} + + found: Final = unreadable(sources) + + assert len(found) == 1 + assert "a.yml" in found[0] + + def test_a_conditional_name_expands_to_the_branch_each_matrix_value_takes() -> None: names: Final = frozenset( name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER}) @@ -310,3 +390,111 @@ def test_every_workflow_in_the_repo_resolves_every_expression_in_its_job_names() unresolved: Final = tuple(f"{owner}: {name}" for name, owner in published(workflow_sources()) if "${{" in name) assert unresolved == () + + +def test_a_matrix_job_with_no_name_publishes_the_id_and_values_github_appends() -> None: + names: Final = frozenset(name for name, _ in published({"a.yml": NAMELESS_MATRIX})) + + assert names == frozenset({"build (3.12)", "build (3.13)"}) + + +def test_a_matrix_job_with_no_name_does_not_collide_with_a_plain_job_carrying_its_id() -> None: + sources: Final = { + "a.yml": NAMELESS_MATRIX, + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + + +def test_a_matrix_job_with_no_name_collides_with_the_suffixed_name_github_writes() -> None: + sources: Final = { + "a.yml": NAMELESS_MATRIX, + "b.yml": "on: pull_request\njobs:\n legacy:\n name: build (3.13)\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`build (3.13)` is published by 2 jobs" in found[0] + + +def test_an_excluded_combination_publishes_no_check_run() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": EXCLUDED_PAIR})) + + assert names == frozenset({"ubuntu-3.12", "ubuntu-3.13", "macos-3.12"}) + + +def test_an_exclude_row_naming_one_key_drops_every_combination_carrying_it() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": EXCLUDED_KEY})) + + assert names == frozenset({"ubuntu-3.12", "ubuntu-3.13"}) + + +def test_a_boolean_matrix_value_renders_the_way_github_writes_it() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": BOOLEAN_MATRIX})) + + assert names == frozenset({"cache true", "cache false"}) + + +def test_a_format_call_its_arguments_cannot_fill_leaves_the_name_unresolved() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": UNFILLABLE_FORMAT})) + + assert names == frozenset({"${{ format('{0} {1}', matrix.shard) }}"}) + + +def test_a_call_to_a_workflow_outside_the_repo_publishes_no_name() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n unit:\n uses: BerriAI/other/.github/workflows/base.yml@main\n", + "b.yml": "on: pull_request\njobs:\n unit:\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"unit"}) + assert collisions(sources) == () + + +def test_a_chain_of_local_reusable_calls_publishes_every_level_of_the_chain() -> None: + sources: Final = { + ".github/workflows/leaf.yml": ( + "on:\n workflow_call:\njobs:\n run:\n name: Leaf\n runs-on: ubuntu-latest\n" + ), + ".github/workflows/mid.yml": ( + "on:\n workflow_call:\njobs:\n call:\n name: Mid\n uses: ./.github/workflows/leaf.yml\n" + ), + "top.yml": "on: pull_request\njobs:\n top:\n name: Top\n uses: ./.github/workflows/mid.yml\n", + } + + names: Final = frozenset(name for name, _ in published(sources)) + + assert names == frozenset({"Top / Mid / Leaf"}) + + +def test_a_job_name_that_is_not_a_string_still_publishes_the_value_github_renders() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n sweep:\n name: 2024\n runs-on: ubuntu-latest\n", + "b.yml": 'on: pull_request\njobs:\n other:\n name: "2024"\n runs-on: ubuntu-latest\n', + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`2024` is published by 2 jobs" in found[0] + + +def test_the_check_fails_when_a_file_in_the_workflows_directory_cannot_be_read() -> None: + assert exit_code({"notes.yml": "just a string\n"}) == 1 + + +def test_the_check_fails_when_two_jobs_publish_one_check_run_name() -> None: + plain: Final = "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n" + + assert exit_code({"a.yml": plain, "b.yml": plain}) == 1 + + +def test_the_check_passes_when_every_file_reads_and_every_name_is_unique() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n sweep:\n runs-on: ubuntu-latest\n", + } + + assert exit_code(sources) == 0 From a08fb489cdb2375b25b19fdca0a0f17a90942574 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sun, 6 Sep 2026 00:38:07 -0700 Subject: [PATCH 36/49] test(ui): accept trailing slashes in migration sidebar links --- tests/e2e/ui/tests/migration/migratedPages.spec.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts index 547330190bd..bc96d9a6c66 100644 --- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts +++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts @@ -29,12 +29,12 @@ async function expectRendered(page: Page) { } /** - * Click a migrated page's sidebar link. Migrated items render as ; + * Click a migrated page's sidebar link. Migrated items render as ; * nested ones live under collapsible groups whose children only render while the * group is open, so expand collapsed groups until the link is clickable. */ async function clickSidebar(page: Page, segment: string) { - const link = sidebar(page).locator(`a[href$="/ui/${segment}"]`).first(); + const link = sidebar(page).locator(`a[href$="/ui/${segment}"], a[href$="/ui/${segment}/"]`).first(); const collapsedGroups = sidebar(page).getByRole("button", { expanded: false }); for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) { const stillCollapsed = await collapsedGroups.count(); From b1cc1e8dae77a6ff1d0517bcedde40cddd55a9d4 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sun, 6 Sep 2026 00:59:18 -0700 Subject: [PATCH 37/49] test(ui): assert navigation through accessible links and page content --- tests/e2e/ui/fixtures/migratedPages.ts | 190 +++++++++++++----- tests/e2e/ui/helpers/navigation.ts | 24 +++ tests/e2e/ui/tests/migration/README.md | 31 ++- .../ui/tests/migration/migratedPages.spec.ts | 126 +++++------- tests/e2e/ui/tests/navigation/sidebar.spec.ts | 59 +++++- 5 files changed, 282 insertions(+), 148 deletions(-) diff --git a/tests/e2e/ui/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts index 58939ca2b9a..bce09b49e10 100644 --- a/tests/e2e/ui/fixtures/migratedPages.ts +++ b/tests/e2e/ui/fixtures/migratedPages.ts @@ -1,50 +1,142 @@ -/** - * Source of truth for the App Router migration E2E suites. - * - * Add an entry (legacy sidebar page id -> route segment) once a page's migration - * has MERGED to the branch under test. Consumers pick it up automatically: - * - migration smoke (tests/migration/migratedPages.spec.ts), via MIGRATED_E2E_SEGMENTS: - * default mount: npm run e2e:migration - * server-root-path mount: SERVER_ROOT_PATH=/ npm run e2e:migration:root - * - navigation specs that assert per-page URLs (tests/navigation/sidebar.spec.ts) - * - * Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts. - */ -export const MIGRATED_E2E_PAGES: Record = { - "api-keys": "api-keys", - models: "models-and-endpoints", - api_ref: "api-reference", - "llm-playground": "playground", - projects: "projects", - "access-groups": "access-groups", - budgets: "budgets", - workflows: "workflows", - "guardrails-monitor": "guardrails-monitor", - "mcp-servers": "mcp-servers", - "search-tools": "search-tools", - "tag-management": "tag-management", - "vector-stores": "vector-stores", - memory: "memory", - policies: "policies", - guardrails: "guardrails", - prompts: "prompts", - "tool-policies": "tool-policies", - skills: "skills", - caching: "caching", - "cost-tracking": "cost-tracking", - "transform-request": "transform-request", - "ui-theme": "ui-theme", - logs: "logs", - "admin-panel": "admin-panel", - "logging-and-alerts": "logging-and-alerts", - "model-hub-table": "model-hub-table", - new_usage: "usage", - usage: "old-usage", - agents: "agents", - "router-settings": "router-settings", - users: "users", - teams: "teams", - organizations: "organizations", -}; +export type MigratedPage = Readonly<{ + segment: string; + linkName: string | RegExp; + group?: string; + content: Readonly<{ role: "heading" | "tab" | "button"; name: string }> | Readonly<{ text: string }>; + unlicensedText?: string; +}>; -export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))]; +export const MIGRATED_E2E_PAGES: Readonly> = { + "api-keys": { segment: "api-keys", linkName: "Virtual Keys", content: { role: "heading", name: "Virtual Keys" } }, + models: { + segment: "models-and-endpoints", + linkName: "Models + Endpoints", + content: { role: "heading", name: "Model Management" }, + }, + api_ref: { + segment: "api-reference", + linkName: "API Reference", + content: { role: "heading", name: "OpenAI Compatible Proxy: API Reference" }, + }, + "llm-playground": { segment: "playground", linkName: "Playground", content: { role: "tab", name: "Chat" } }, + projects: { + segment: "projects", + linkName: /^Projects(?: Beta)?$/, + content: { role: "heading", name: "Projects" }, + }, + "access-groups": { + segment: "access-groups", + linkName: "Access Groups", + content: { role: "heading", name: "Access Groups" }, + }, + budgets: { segment: "budgets", linkName: "Budgets", content: { role: "heading", name: "Budgets" } }, + workflows: { + segment: "workflows", + linkName: "Workflow Runs", + group: "Agentic", + content: { text: "Workflow Runs" }, + }, + "guardrails-monitor": { + segment: "guardrails-monitor", + linkName: "Guardrails Monitor", + content: { role: "heading", name: "Guardrails Monitor" }, + }, + "mcp-servers": { + segment: "mcp-servers", + linkName: "MCP Servers", + content: { role: "heading", name: "MCP Servers" }, + }, + "search-tools": { + segment: "search-tools", + linkName: "Search Tools", + group: "Tools", + content: { role: "heading", name: "Search Tools" }, + }, + "tag-management": { + segment: "tag-management", + linkName: "Tag Management", + group: "Experimental", + content: { role: "heading", name: "Tag Management" }, + }, + "vector-stores": { + segment: "vector-stores", + linkName: "Vector Stores", + group: "Tools", + content: { role: "heading", name: "Vector Store Management" }, + }, + memory: { segment: "memory", linkName: "Memory", group: "Agentic", content: { role: "heading", name: "Memory" } }, + policies: { segment: "policies", linkName: "Policies", content: { role: "tab", name: "Policy Simulator" } }, + guardrails: { segment: "guardrails", linkName: "Guardrails", content: { role: "tab", name: "Guardrails" } }, + prompts: { + segment: "prompts", + linkName: "Prompts", + group: "Experimental", + content: { role: "button", name: "Add New Prompt" }, + }, + "tool-policies": { + segment: "tool-policies", + linkName: "Tool Policies", + group: "Tools", + content: { role: "heading", name: "Tool Policies" }, + }, + skills: { segment: "skills", linkName: "Skills", content: { role: "heading", name: "Skills" } }, + caching: { segment: "caching", linkName: "Response Cache", content: { role: "tab", name: "Cache Settings" } }, + "cost-tracking": { + segment: "cost-tracking", + linkName: "Cost Tracking", + group: "Settings", + content: { text: "Cost Tracking Settings" }, + }, + "transform-request": { + segment: "transform-request", + linkName: "API Playground", + group: "Experimental", + content: { role: "heading", name: "Playground" }, + }, + "ui-theme": { + segment: "ui-theme", + linkName: "UI Theme", + group: "Settings", + content: { role: "heading", name: "UI Theme Customization" }, + }, + logs: { segment: "logs", linkName: "Logs", content: { role: "heading", name: "Request Logs" } }, + "admin-panel": { + segment: "admin-panel", + linkName: "Admin Settings", + group: "Settings", + content: { role: "heading", name: "Admin Access" }, + }, + "logging-and-alerts": { + segment: "logging-and-alerts", + linkName: "Logging & Alerts", + group: "Settings", + content: { role: "tab", name: "Logging Callbacks" }, + }, + "model-hub-table": { + segment: "model-hub-table", + linkName: "AI Hub", + content: { role: "heading", name: "AI Hub" }, + }, + new_usage: { segment: "usage", linkName: "Usage", content: { role: "heading", name: "Usage View" } }, + usage: { + segment: "old-usage", + linkName: "Old Usage", + group: "Experimental", + content: { role: "tab", name: "All Up" }, + }, + agents: { segment: "agents", linkName: "Agents", group: "Agentic", content: { role: "heading", name: "Agents" } }, + "router-settings": { + segment: "router-settings", + linkName: "Router Settings", + group: "Settings", + content: { role: "heading", name: "Routing Settings" }, + }, + users: { segment: "users", linkName: "Internal Users", content: { role: "tab", name: "Users" } }, + teams: { segment: "teams", linkName: "Teams", content: { role: "heading", name: "Teams" } }, + organizations: { + segment: "organizations", + linkName: "Organizations", + content: { text: "Click on an organization ID to view its details." }, + unlicensedText: "This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key here.", + }, +}; diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index a58ece16f9c..4a7c4e7baa9 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -1,5 +1,29 @@ import { Page } from "../fixtures/pages"; import { Page as PlaywrightPage, expect } from "@playwright/test"; +import { UI_BASE_URL } from "../constants"; + +export const sidebarLink = (page: PlaywrightPage, name: string | RegExp) => + page.getByRole("complementary").getByRole("link", { name, exact: true }); + +export async function clickSidebarLink(page: PlaywrightPage, name: string | RegExp, groupName?: string): Promise { + const link = sidebarLink(page, name); + if (groupName && !(await link.isVisible())) { + const group = page.getByRole("complementary").getByRole("button", { name: groupName, exact: true }); + await expect(group).toBeVisible(); + if ((await group.getAttribute("aria-expanded")) === "false") { + await group.click(); + } + } + await link.click(); +} + +export async function expectUiRoute(page: PlaywrightPage, segment: string): Promise { + const root = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, ""); + const expected = new URL(`${root}/ui/${segment}`, UI_BASE_URL); + await expect(page, `navigate to ${expected.pathname}`).toHaveURL( + (url) => url.origin === expected.origin && url.pathname.replace(/\/+$/, "") === expected.pathname, + ); +} /** * Navigates to a specific page using the page query parameter. diff --git a/tests/e2e/ui/tests/migration/README.md b/tests/e2e/ui/tests/migration/README.md index d6b33598ec4..59933502463 100644 --- a/tests/e2e/ui/tests/migration/README.md +++ b/tests/e2e/ui/tests/migration/README.md @@ -1,17 +1,25 @@ # App Router migration smoke A growing E2E smoke for pages migrated from the legacy `?page=` switch to App -Router path routes. For each migrated page it clicks the page's sidebar link, checks -the URL is the path route and the page renders, reloads it, then clicks off to a -legacy page and back to confirm navigation still works. It runs in two situations: -the default mount and a non-root `SERVER_ROOT_PATH` mount. +Router path routes. For each page it clicks the sidebar link by its accessible +name, verifies the destination's content, reloads it, then visits Virtual Keys +and returns. It runs at the default mount and a non-root `SERVER_ROOT_PATH` mount + +Link selection does not depend on `href` formatting. URL assertions compare the +origin and pathname, allowing a trailing slash, query string, and fragment while +rejecting another route or mount. Reloads must return a successful document, +and each journey must finish without uncaught browser errors ## Adding a page -When a page's migration merges, add its route segment to -`tests/e2e/ui/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES` -in `ui/litellm-dashboard/src/utils/migratedPages.ts`). Both suites pick it up -automatically. +Add an entry to `tests/e2e/ui/fixtures/migratedPages.ts`, keyed by the legacy page +ID. Specify its route segment, accessible link name, sidebar group if collapsed, +and distinctive visible content such as a heading or tab. Keep expectations +independent of the application's route table so an incorrect destination fails +the test. Both navigation suites use this fixture + +For a licensed-only page, `unlicensedText` describes the expected upgrade notice. +The authenticated session's license claim determines which content must render ## Running @@ -31,4 +39,9 @@ SERVER_ROOT_PATH=/litellm npm run e2e:migration:root ``` `globalSetup` logs in once per role; the admin storage state is reused for these -tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login`. +tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login` + +`tests/navigation/sidebar.spec.ts` also checks the navigation helpers against +equivalent link formats on the live dashboard and a deep link containing a query +string and fragment. The link-format cases change only the rendered `href` +attribute to exercise the locator contract; destination pages and APIs remain live diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts index bc96d9a6c66..f2dbd2dbc77 100644 --- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts +++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts @@ -1,105 +1,73 @@ import { test, expect, type Page } from "@playwright/test"; -import { MIGRATED_E2E_SEGMENTS } from "../../fixtures/migratedPages"; +import { MIGRATED_E2E_PAGES, type MigratedPage } from "../../fixtures/migratedPages"; import { ADMIN_STORAGE_PATH } from "../../constants"; -import { dismissFeedbackPopup } from "../../helpers/navigation"; +import { clickSidebarLink, dismissFeedbackPopup, expectUiRoute, sidebarLink } from "../../helpers/navigation"; +import { proxyIsPremium } from "../../helpers/premium"; -/** - * App Router migration smoke as a user journey: start where the proxy lands you, - * click a migrated page in the sidebar, confirm it routed and rendered, reload it - * (the check a wrong server_root_path breaks), bounce to a legacy page and back, - * and, once two pages are migrated, navigate directly between two migrated pages. - * - * Driven by MIGRATED_E2E_SEGMENTS, so it grows as pages are migrated. Set - * SERVER_ROOT_PATH (e.g. "/litellm") to exercise the non-root mount; leave it - * unset for the default mount. Boot the proxy with the matching value first. - */ -const ROOT = process.env.SERVER_ROOT_PATH ?? ""; +const ROOT = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, ""); +const apiKeys = MIGRATED_E2E_PAGES["api-keys"]; -const esc = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); -const pathRe = (segment: string) => new RegExp(`${esc(ROOT)}/ui/${esc(segment)}/?($|\\?)`); -// Scope nav lookups to the sidebar (a `complementary` landmark). The top bar -// now renders a breadcrumb whose current-page item is also a "Virtual Keys" -// link, so an unscoped locator would match two elements. -const sidebar = (page: Page) => page.getByRole("complementary"); -const virtualKeysLink = (page: Page) => sidebar(page).getByRole("link", { name: "Virtual Keys", exact: true }); - -/** The dashboard shell is present (sidebar rendered); page didn't 404 / crash. */ -async function expectRendered(page: Page) { - await expect(virtualKeysLink(page)).toBeVisible({ timeout: 20_000 }); +async function expectContent(page: Page, destination: MigratedPage): Promise { + await expect(sidebarLink(page, apiKeys.linkName)).toBeVisible({ timeout: 20_000 }); + const main = page.getByRole("main"); + if (destination.unlicensedText && !proxyIsPremium()) { + await expect(main.getByText(destination.unlicensedText, { exact: true })).toBeVisible(); + return; + } + const content = destination.content; + const landmark = + "role" in content + ? main.getByRole(content.role, { name: content.name, exact: true }) + : main.getByText(content.text, { exact: true }); + await expect(landmark).toBeVisible(); } -/** - * Click a migrated page's sidebar link. Migrated items render as ; - * nested ones live under collapsible groups whose children only render while the - * group is open, so expand collapsed groups until the link is clickable. - */ -async function clickSidebar(page: Page, segment: string) { - const link = sidebar(page).locator(`a[href$="/ui/${segment}"], a[href$="/ui/${segment}/"]`).first(); - const collapsedGroups = sidebar(page).getByRole("button", { expanded: false }); - for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) { - const stillCollapsed = await collapsedGroups.count(); - if (stillCollapsed === 0) break; - await collapsedGroups.first().click(); - await expect(collapsedGroups).toHaveCount(stillCollapsed - 1); - } - await link.click(); +async function navigateToDestination(page: Page, destination: MigratedPage): Promise { + await clickSidebarLink(page, destination.linkName, destination.group); + await expectUiRoute(page, destination.segment); + await dismissFeedbackPopup(page); + await expectContent(page, destination); } test.use({ storageState: ADMIN_STORAGE_PATH }); test.describe("App Router migrated pages", () => { - for (const segment of MIGRATED_E2E_SEGMENTS) { - test(`${segment}: sidebar nav, reload, and round-trip via the api-keys landing`, async ({ page }) => { + for (const destination of Object.values(MIGRATED_E2E_PAGES)) { + test(`${destination.segment}: sidebar nav, reload, and round-trip via the api-keys landing`, async ({ page }) => { const pageErrors: string[] = []; - page.on("pageerror", (e) => pageErrors.push(String(e))); + page.on("pageerror", (error) => pageErrors.push(String(error))); - // 1. Start where the proxy lands us. - await page.goto(`${ROOT}/ui/`); + const landing = await page.goto(`${ROOT}/ui/`); + expect(landing?.ok(), "dashboard document loads successfully").toBe(true); await dismissFeedbackPopup(page); - await expectRendered(page); + await expectContent(page, apiKeys); - // 2. Click the migrated page in the sidebar -> path route + rendered. - await clickSidebar(page, segment); - await expect(page).toHaveURL(pathRe(segment)); - await expectRendered(page); - // 3. Reload the path route directly; a wrong server_root_path 404s here. - await page.reload(); + await navigateToDestination(page, destination); + + const reloaded = await page.reload(); + expect(reloaded?.ok(), `${destination.segment} document loads on reload`).toBe(true); + await expectUiRoute(page, destination.segment); await dismissFeedbackPopup(page); - await expect(page).toHaveURL(pathRe(segment)); - await expectRendered(page); - // 4. Click the Virtual Keys sidebar link to the api-keys landing (now a path route), then back. - await virtualKeysLink(page).click(); - await expect(page).toHaveURL(pathRe("api-keys")); - await dismissFeedbackPopup(page); - await expectRendered(page); - // 5. Click back to the migrated page. - await clickSidebar(page, segment); - await expect(page).toHaveURL(pathRe(segment)); - await expectRendered(page); - expect(pageErrors, `page errors during ${segment} journey`).toEqual([]); + await expectContent(page, destination); + + await navigateToDestination(page, apiKeys); + await navigateToDestination(page, destination); + expect(pageErrors, `page errors during ${destination.segment} journey`).toEqual([]); }); } test("navigates directly between two migrated pages", async ({ page }) => { - test.skip(MIGRATED_E2E_SEGMENTS.length < 2, "needs >= 2 migrated pages"); - const [first, second] = MIGRATED_E2E_SEGMENTS; const pageErrors: string[] = []; - page.on("pageerror", (e) => pageErrors.push(String(e))); + page.on("pageerror", (error) => pageErrors.push(String(error))); - await page.goto(`${ROOT}/ui/`); + const landing = await page.goto(`${ROOT}/ui/`); + expect(landing?.ok(), "dashboard document loads successfully").toBe(true); await dismissFeedbackPopup(page); + await expectContent(page, apiKeys); - await clickSidebar(page, first); - await expect(page).toHaveURL(pathRe(first)); - await expectRendered(page); - await clickSidebar(page, second); - await expect(page).toHaveURL(pathRe(second)); - await expectRendered(page); - // Back to the first migrated page. - await clickSidebar(page, first); - await expect(page).toHaveURL(pathRe(first)); - await expectRendered(page); - - expect(pageErrors, "page errors during migrated -> migrated nav").toEqual([]); + for (const destination of [apiKeys, MIGRATED_E2E_PAGES.models, apiKeys]) { + await navigateToDestination(page, destination); + } + expect(pageErrors, "page errors during migrated page navigation").toEqual([]); }); }); diff --git a/tests/e2e/ui/tests/navigation/sidebar.spec.ts b/tests/e2e/ui/tests/navigation/sidebar.spec.ts index b220dc09ae2..eaa4985c659 100644 --- a/tests/e2e/ui/tests/navigation/sidebar.spec.ts +++ b/tests/e2e/ui/tests/navigation/sidebar.spec.ts @@ -3,7 +3,13 @@ import { Role } from "../../fixtures/roles"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { Page } from "../../fixtures/pages"; import { menuLabelToPage } from "../../fixtures/menuMappings"; -import { navigateToPage } from "../../helpers/navigation"; +import { + clickSidebarLink, + dismissFeedbackPopup, + expectUiRoute, + navigateToPage, + sidebarLink, +} from "../../helpers/navigation"; import { MIGRATED_E2E_PAGES } from "../../fixtures/migratedPages"; import type { Page as PlaywrightPage } from "@playwright/test"; @@ -11,7 +17,7 @@ const sidebarButtons = { [Role.ProxyAdmin]: [ "Virtual Keys", "Playground", - "Models", + "Models + Endpoints", "Usage", "Teams", "Internal Users", @@ -22,9 +28,9 @@ const sidebarButtons = { /** Migrated pages live at a path route; legacy pages keep the ?page= query param. */ async function expectPageUrl(page: PlaywrightPage, pageKey: string): Promise { - const migratedSegment = MIGRATED_E2E_PAGES[pageKey]; - if (migratedSegment) { - await expect(page).toHaveURL(new RegExp(`/ui/${migratedSegment}/?($|\\?)`)); + const migratedPage = MIGRATED_E2E_PAGES[pageKey]; + if (migratedPage) { + await expectUiRoute(page, migratedPage.segment); } else { await expect(page).toHaveURL(new RegExp(`[?&]page=${pageKey}(&|$)`)); } @@ -51,12 +57,7 @@ for (const { role, storage } of roles) { throw new Error(`No page mapping found for menu label: ${buttonLabel}`); } - // Sidebar items are links inside the `complementary` landmark; scoping - // there avoids the top-bar breadcrumb, which also links the page name. - const tab = page.getByRole("complementary").getByRole("link", { name: buttonLabel }); - await expect(tab).toBeVisible(); - - await tab.click(); + await clickSidebarLink(page, buttonLabel); await expectPageUrl(page, expectedPage); } @@ -81,5 +82,41 @@ for (const { role, storage } of roles) { await navigateToPage(page, Page.LlmPlayground); await expectPageUrl(page, Page.LlmPlayground); }); + + for (const format of ["without trailing slash", "absolute with query and fragment", "relative"] as const) { + test(`sidebar locator tolerates hrefs ${format}`, async ({ page }) => { + await page.goto("/ui/"); + await dismissFeedbackPopup(page); + const link = sidebarLink(page, "Models + Endpoints"); + await expect(link).toBeVisible(); + const destination = new URL("/ui/models-and-endpoints/", page.url()); + const href = + format === "without trailing slash" + ? destination.pathname.replace(/\/$/, "") + : format === "relative" + ? "./models-and-endpoints/" + : `${destination.href}?source=navigation-smoke#overview`; + + await link.evaluate((element, value) => element.setAttribute("href", value), href); + await expect(link).toHaveAttribute("href", href); + await clickSidebarLink(page, "Models + Endpoints"); + + await expectUiRoute(page, "models-and-endpoints"); + await expect( + page.getByRole("main").getByRole("heading", { name: "Model Management", exact: true }), + ).toBeVisible(); + }); + } + + test("route assertion tolerates a query string and fragment on a deep link", async ({ page }) => { + const response = await page.goto("/ui/models-and-endpoints/?source=navigation-smoke#overview"); + expect(response?.ok()).toBe(true); + await expectUiRoute(page, "models-and-endpoints"); + await expect( + page.getByRole("main").getByRole("heading", { name: "Model Management", exact: true }), + ).toBeVisible(); + expect(new URL(page.url()).search).toBe("?source=navigation-smoke"); + expect(new URL(page.url()).hash).toBe("#overview"); + }); }); } From 5ca9e2605081d26503786cdd116e7fc8245c8903 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 01:09:36 -0700 Subject: [PATCH 38/49] fix(ci): leave check-run names the sweep cannot resolve out of the comparison A job name holding an expression the sweep could not resolve was compared as if it were the published name. Two jobs whose names differ per matrix value or per caller input were reported as a collision, and a matrix that was itself an expression collapsed onto the bare job id and did the same. Model what a job publishes as known names beside the reasons the rest stay unknown. Anything the sweep cannot work out contributes no name and is reported as a note instead of guessed at. An expression over contexts that are fixed for the whole run still compares, so two jobs sharing one of those are still caught. --- .../check_workflow_job_name_collisions.py | 173 ++++++++++++++---- .../test_workflow_job_name_collisions.py | 151 ++++++++++++++- 2 files changed, 280 insertions(+), 44 deletions(-) diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py index 4eef3df1c11..1e81acc4f90 100644 --- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -20,16 +20,24 @@ with the other rows', so two shard lists that overlap collide even though their templates read differently. Each expression is evaluated per combination over the pieces a job name can hold: string literals, `matrix.`, `format()`, `==` and `!=`, and the ` && || ` idiom, which is how the shards reach their -real ` / Run tests` names rather than staying opaque. An expression -nothing resolves stays in the string, so two jobs carrying the same unresolved -template still compare equal and their collision is still caught. +real ` / Run tests` names rather than staying opaque. + +Whatever the sweep cannot work out is left out of the comparison and reported +instead of guessed, because a guess that lands wrong fails a workflow GitHub +would have published perfectly well. A name left holding an expression over +`inputs`, `needs` or `matrix` reads differently from every job that runs it, so +it is one of those; an expression over `github` and the other contexts fixed for +the whole run reads the same everywhere, so two jobs carrying it still clash and +the check still says so. A matrix that is itself an expression or that lists +values which are not scalars, and a call this sweep cannot follow, go in the same +bucket. The cost is that a real clash hiding behind one of them goes unseen, +which leaves a merge no worse off than before this check existed, where the +opposite direction would block work that was fine. A job calling a local reusable workflow publishes one check run per job of the callee, named ` / ` and chained through however many levels of local calls it takes, which is why a caller's name never collides with a plain -job that happens to match it. A call this sweep cannot follow, to a workflow -outside the repo or to one that is not there, publishes nothing rather than a -name GitHub never posts. A file under `.github/workflows/` that does not read as +job that happens to match it. A file under `.github/workflows/` that does not read as a workflow at all is reported rather than skipped, since skipping it silently would hide every job it holds. """ @@ -54,6 +62,8 @@ MATRIX_REF: Final = re.compile(r"^matrix\.(?P[\w-]+)$") LITERAL: Final = re.compile(r"^'(?P[^']*)'$") FORMAT_CALL: Final = re.compile(r"^format\((?P.*)\)$", re.DOTALL) COMPARISON: Final = re.compile(r"^(?P.+?)\s*(?P==|!=)\s*(?P.+)$", re.DOTALL) +CONTEXT_REF: Final = re.compile(r"\b(?P[a-z][\w-]*)\s*\.") +RUN_FIXED_CONTEXTS: Final = frozenset({"github", "vars", "env", "runner", "secrets"}) NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({}) NO_CALLERS: Final[frozenset[str]] = frozenset() SCALAR: Final = (str, int, float) @@ -66,6 +76,19 @@ class Unreadable: reason: str +@dataclass(frozen=True, slots=True) +class Opaque: + reason: str + + +@dataclass(frozen=True, slots=True) +class Names: + """The check-run names a job publishes, beside the reasons the rest of them stay unknown.""" + + known: tuple[str, ...] = () + unknown: tuple[str, ...] = () + + class Job(BaseModel): name: object = None uses: str | None = None @@ -107,12 +130,23 @@ def publishes_check_runs(raw_on: object) -> bool: return events(raw_on) != frozenset({"workflow_call"}) -def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...]: - return tuple( - (str(key), tuple(scalar_text(value) for value in values if isinstance(value, SCALAR))) - for key, values in matrix.items() - if str(key) not in MATRIX_DIRECTIVES and isinstance(values, Sequence) and not isinstance(values, str) +def scalar_list(value: object) -> tuple[str, ...] | Opaque: + """One matrix key's values, or why the combinations it produces cannot be worked out.""" + if not isinstance(value, Sequence) or isinstance(value, str): + return Opaque("a matrix key holds something other than a list of values") + if any(not isinstance(item, SCALAR) for item in value): + return Opaque("a matrix key lists values that are not plain scalars") + return tuple(scalar_text(item) for item in value) + + +def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...] | Opaque: + listed: Final = tuple( + (str(key), scalar_list(values)) for key, values in matrix.items() if str(key) not in MATRIX_DIRECTIVES ) + opaque: Final = next((values for _, values in listed if isinstance(values, Opaque)), None) + if opaque is not None: + return opaque + return tuple((key, values) for key, values in listed if not isinstance(values, Opaque)) def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...]: @@ -150,17 +184,20 @@ def crossed_values(listed: Sequence[tuple[str, tuple[str, ...]]]) -> tuple[Mappi ) -def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...]: +def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...] | Opaque: """One mapping per job the matrix produces, `exclude` applied before `include` as GitHub does.""" matrix: Final = job.strategy.get("matrix") - if not isinstance(matrix, Mapping): + if matrix is None: return () + if not isinstance(matrix, Mapping): + return Opaque("the matrix itself comes from an expression") + listed: Final = listed_values(matrix) + if isinstance(listed, Opaque): + return listed rows: Final = directive_rows(matrix, "include") dropped: Final = directive_rows(matrix, "exclude") kept: Final = tuple( - combination - for combination in crossed_values(listed_values(matrix)) - if not any(drops(row, combination) for row in dropped) + combination for combination in crossed_values(listed) if not any(drops(row, combination) for row in dropped) ) standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in kept)) return (*(extended(combination, rows) for combination in kept), *standalone) @@ -243,11 +280,33 @@ def rendered(template: str, values: Mapping[str, str]) -> str: return EXPRESSION.sub(lambda span: resolved_span(span, values), template) -def expand(template: str, job: Job) -> tuple[str, ...]: +def stable(name: str) -> bool: + """An expression over the run's own contexts reads the same from every job, so two of them still clash.""" + return all( + reference.group("root") in RUN_FIXED_CONTEXTS + for span in EXPRESSION.finditer(name) + for reference in CONTEXT_REF.finditer(span.group("body")) + ) + + +def comparable(name: str) -> bool: + """A leftover expression over `inputs`, `needs` or `matrix` names a different check run per job.""" + return EXPRESSION.search(name) is None or stable(name) + + +def settled(names: Sequence[str]) -> Names: + return Names( + tuple(name for name in names if comparable(name)), + tuple(f"its name stays `{name}`" for name in names if not comparable(name)), + ) + + +def expand(template: str, job: Job) -> Names: combinations: Final = matrix_combinations(job) - if not combinations: - return (rendered(template, NO_MATRIX),) - return tuple(dict.fromkeys(rendered(template, values) for values in combinations)) + if isinstance(combinations, Opaque): + return Names((), (combinations.reason,)) + over: Final = combinations or (NO_MATRIX,) + return settled(tuple(dict.fromkeys(rendered(template, values) for values in over))) def suffixed(job_id: str, combination: Mapping[str, str]) -> str: @@ -255,11 +314,14 @@ def suffixed(job_id: str, combination: Mapping[str, str]) -> str: return f"{job_id} ({', '.join(combination.values())})" if combination else job_id -def published_names(job_id: str, job: Job) -> tuple[str, ...]: +def published_names(job_id: str, job: Job) -> Names: if job.name is not None: return expand(scalar_text(job.name), job) - suffixes: Final = tuple(dict.fromkeys(suffixed(job_id, values) for values in matrix_combinations(job))) - return suffixes or (job_id,) + combinations: Final = matrix_combinations(job) + if isinstance(combinations, Opaque): + return Names((), (combinations.reason,)) + suffixes: Final = tuple(dict.fromkeys(suffixed(job_id, values) for values in combinations)) + return Names(suffixes or (job_id,)) def callee_path(job: Job) -> str | None: @@ -268,22 +330,40 @@ def callee_path(job: Job) -> str | None: return job.uses[len(LOCAL_CALL_PREFIX) :].split("@")[0] -def job_names( - job_id: str, job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str] = NO_CALLERS -) -> tuple[str, ...]: +def joined(groups: Sequence[Names]) -> Names: + return Names( + tuple(name for group in groups for name in group.known), + tuple(reason for group in groups for reason in group.unknown), + ) + + +def call_blocker(job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str]) -> str | None: + path: Final = callee_path(job) + if path is None: + return "it calls a reusable workflow outside this repository" + if path in callers: + return f"its call to {path} loops back on itself" + return None if path in workflows else f"it calls {path}, which this checkout does not hold" + + +def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str] = NO_CALLERS) -> Names: prefixes: Final = published_names(job_id, job) if job.uses is None: return prefixes - path: Final = callee_path(job) - callee: Final = workflows.get(path or "") - if path is None or callee is None or path in callers: - return () - suffixes: Final = tuple( - name - for callee_id, callee_job in callee.jobs.items() - for name in job_names(callee_id, callee_job, workflows, callers | {path}) + blocker: Final = call_blocker(job, workflows, callers) + if blocker is not None: + return Names((), (*prefixes.unknown, blocker)) + path: Final = callee_path(job) or "" + suffixes: Final = joined( + tuple( + job_names(callee_id, callee_job, workflows, callers | {path}) + for callee_id, callee_job in workflows[path].jobs.items() + ) + ) + return Names( + tuple(f"{prefix} / {suffix}" for prefix in prefixes.known for suffix in suffixes.known), + (*prefixes.unknown, *suffixes.unknown), ) - return tuple(f"{prefix} / {suffix}" for prefix in prefixes for suffix in suffixes) def readable(sources: Mapping[str, str]) -> Mapping[str, tuple[Workflow, object]]: @@ -300,15 +380,29 @@ def unreadable(sources: Mapping[str, str]) -> tuple[str, ...]: ) -def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]: +def scanned_jobs(sources: Mapping[str, str]) -> Iterator[tuple[str, str, Names]]: parsed: Final = readable(sources) workflows: Final = {rel: workflow for rel, (workflow, _) in parsed.items()} for rel, (workflow, raw_on) in parsed.items(): if not publishes_check_runs(raw_on): continue for job_id, job in workflow.jobs.items(): - for name in job_names(job_id, job, workflows): - yield name, f"{rel} job `{job_id}`" + yield rel, job_id, job_names(job_id, job, workflows) + + +def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]: + for rel, job_id, names in scanned_jobs(sources): + for name in names.known: + yield name, f"{rel} job `{job_id}`" + + +def blind_spots(sources: Mapping[str, str]) -> tuple[str, ...]: + """Jobs whose published names GitHub decides at run time, which no offline sweep can compare.""" + return tuple( + f"{rel} job `{job_id}` publishes a name this check cannot work out because {reason}." + for rel, job_id, names in scanned_jobs(sources) + for reason in sorted(names.unknown) + ) def owners_by_name(sources: Mapping[str, str]) -> Iterator[tuple[str, tuple[str, ...]]]: @@ -338,6 +432,9 @@ def report(header: str, problems: Sequence[str]) -> None: def exit_code(sources: Mapping[str, str]) -> int: unread: Final = unreadable(sources) found: Final = collisions(sources) + blind: Final = blind_spots(sources) + if blind: + print("NOTE: names left out of the comparison:\n - " + "\n - ".join(blind)) report("Some workflows could not be read", unread) report("Check-run names are not unique", found) if unread or found: diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py index 3dd8a3d4088..39e61ff48d1 100644 --- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -2,6 +2,7 @@ from typing import Final from check_workflow_job_name_collisions import ( Unreadable, + blind_spots, callee_path, collisions, exit_code, @@ -301,9 +302,9 @@ def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None: assert "`Analyze (go)` is published by 2 jobs" in found[0] -def test_two_jobs_sharing_one_unresolvable_template_still_collide() -> None: +def test_two_jobs_sharing_a_template_over_the_run_itself_still_collide() -> None: template: Final = ( - "on: pull_request\njobs:\n {job}:\n name: ${{{{ matrix.shard }}}}\n runs-on: ubuntu-latest\n" + "on: pull_request\njobs:\n {job}:\n name: ${{{{ github.event_name }}}}-build\n runs-on: ubuntu-latest\n" ) sources: Final = { "a.yml": template.format(job="one"), @@ -316,6 +317,19 @@ def test_two_jobs_sharing_one_unresolvable_template_still_collide() -> None: assert "is published by 2 jobs" in found[0] +def test_two_jobs_sharing_a_template_that_reads_per_job_are_not_called_a_collision() -> None: + template: Final = ( + "on: pull_request\njobs:\n {job}:\n name: ${{{{ matrix.shard }}}}\n runs-on: ubuntu-latest\n" + ) + sources: Final = { + "a.yml": template.format(job="one"), + "b.yml": template.format(job="two"), + } + + assert collisions(sources) == () + assert len(blind_spots(sources)) == 2 + + def test_a_file_that_is_not_a_workflow_is_reported_rather_than_skipped() -> None: sources: Final = { "notes.yml": "just a string\n", @@ -437,13 +451,14 @@ def test_a_boolean_matrix_value_renders_the_way_github_writes_it() -> None: assert names == frozenset({"cache true", "cache false"}) -def test_a_format_call_its_arguments_cannot_fill_leaves_the_name_unresolved() -> None: - names: Final = frozenset(name for name, _ in published({"unit.yml": UNFILLABLE_FORMAT})) +def test_a_format_call_its_arguments_cannot_fill_publishes_nothing_to_compare() -> None: + sources: Final = {"unit.yml": UNFILLABLE_FORMAT} - assert names == frozenset({"${{ format('{0} {1}', matrix.shard) }}"}) + assert frozenset(name for name, _ in published(sources)) == frozenset() + assert "its name stays" in blind_spots(sources)[0] -def test_a_call_to_a_workflow_outside_the_repo_publishes_no_name() -> None: +def test_a_call_to_a_workflow_outside_the_repo_is_reported_rather_than_guessed() -> None: sources: Final = { "a.yml": "on: pull_request\njobs:\n unit:\n uses: BerriAI/other/.github/workflows/base.yml@main\n", "b.yml": "on: pull_request\njobs:\n unit:\n runs-on: ubuntu-latest\n", @@ -451,6 +466,7 @@ def test_a_call_to_a_workflow_outside_the_repo_publishes_no_name() -> None: assert frozenset(name for name, _ in published(sources)) == frozenset({"unit"}) assert collisions(sources) == () + assert "outside this repository" in blind_spots(sources)[0] def test_a_chain_of_local_reusable_calls_publishes_every_level_of_the_chain() -> None: @@ -498,3 +514,126 @@ def test_the_check_passes_when_every_file_reads_and_every_name_is_unique() -> No } assert exit_code(sources) == 0 + + +def test_two_callers_of_one_reusable_workflow_named_from_its_inputs_do_not_collide() -> None: + sources: Final = { + ".github/workflows/callee.yml": ( + "on:\n workflow_call:\njobs:\n run:\n name: ${{ inputs.suite }}\n runs-on: ubuntu-latest\n" + ), + "caller.yml": ( + "on: pull_request\njobs:\n" + " alpha:\n name: A\n uses: ./.github/workflows/callee.yml\n with:\n suite: alpha\n" + " beta:\n name: A\n uses: ./.github/workflows/callee.yml\n with:\n suite: beta\n" + ), + } + + assert collisions(sources) == () + assert len(blind_spots(sources)) == 2 + + +def test_a_matrix_that_is_itself_an_expression_never_collapses_onto_the_bare_job_id() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n strategy:\n" + " matrix: ${{ fromJson(needs.plan.outputs.matrix) }}\n runs-on: ubuntu-latest\n" + ), + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"build"}) + assert collisions(sources) == () + assert "the matrix itself comes from an expression" in blind_spots(sources)[0] + + +def test_a_matrix_listing_objects_never_collapses_onto_the_bare_job_id() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n strategy:\n matrix:\n target:\n" + " - os: ubuntu\n - os: windows\n runs-on: ubuntu-latest\n" + ), + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"build"}) + assert collisions(sources) == () + assert "not plain scalars" in blind_spots(sources)[0] + + +def test_a_call_to_a_workflow_file_the_checkout_does_not_hold_is_reported() -> None: + sources: Final = {"a.yml": "on: pull_request\njobs:\n unit:\n uses: ./.github/workflows/gone.yml\n"} + + assert collisions(sources) == () + assert "which this checkout does not hold" in blind_spots(sources)[0] + + +def test_reusable_workflows_calling_each_other_in_a_loop_are_reported_not_followed() -> None: + sources: Final = { + ".github/workflows/a.yml": ( + "on:\n workflow_call:\njobs:\n call:\n name: A\n uses: ./.github/workflows/b.yml\n" + ), + ".github/workflows/b.yml": ( + "on:\n workflow_call:\njobs:\n call:\n name: B\n uses: ./.github/workflows/a.yml\n" + ), + "top.yml": "on: pull_request\njobs:\n top:\n name: Top\n uses: ./.github/workflows/a.yml\n", + } + + assert collisions(sources) == () + assert any("loops back on itself" in spot for spot in blind_spots(sources)) + + +def test_a_caller_still_publishes_the_callee_jobs_it_can_read() -> None: + sources: Final = { + ".github/workflows/callee.yml": ( + "on:\n workflow_call:\njobs:\n" + " lint:\n name: Lint\n runs-on: ubuntu-latest\n" + " suite:\n name: ${{ inputs.suite }}\n runs-on: ubuntu-latest\n" + ), + "caller.yml": "on: pull_request\njobs:\n call:\n name: A\n uses: ./.github/workflows/callee.yml\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"A / Lint"}) + assert len(blind_spots(sources)) == 1 + + +def test_a_name_the_check_cannot_work_out_is_reported_without_failing_the_check() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n unit:\n uses: BerriAI/other/.github/workflows/base.yml@main\n", + } + + assert blind_spots(sources) != () + assert exit_code(sources) == 0 + + +def test_a_caller_whose_own_name_is_unreadable_publishes_none_of_its_callee_names() -> None: + sources: Final = { + ".github/workflows/callee.yml": ( + "on:\n workflow_call:\njobs:\n lint:\n name: Lint\n runs-on: ubuntu-latest\n" + ), + "caller.yml": ( + "on: pull_request\njobs:\n call:\n name: ${{ matrix.suite }}\n" + " uses: ./.github/workflows/callee.yml\n" + ), + "other.yml": "on: pull_request\njobs:\n plain:\n name: Lint\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"Lint"}) + assert collisions(sources) == () + assert "its name stays" in blind_spots(sources)[0] + + +def test_an_include_row_naming_a_listed_key_extends_only_the_combinations_it_matches() -> None: + sources: Final = { + "unit.yml": ( + "on: pull_request\njobs:\n unit:\n" + " name: ${{ matrix.python-version }} ${{ matrix.label }}\n" + " runs-on: ubuntu-latest\n strategy:\n matrix:\n" + ' python-version: ["3.12", "3.13"]\n' + " include:\n" + ' - python-version: "3.12"\n' + " label: fast\n" + ) + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"3.12 fast"}) + assert len(blind_spots(sources)) == 1 From 904542a559c1f297c63a5b299557d340c6f3f350 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 01:27:03 -0700 Subject: [PATCH 39/49] fix(ci): stop guessing at names built from contexts the sweep cannot read Three ways the sweep could fail a workflow GitHub would publish fine. `github.workflow` and `github.job` were counted as fixed for the whole run, so two jobs naming themselves after the workflow they sit in were reported as a collision. `runner` and `vars` were wrong the same way. Drop the exception entirely: a name still holding an expression is one GitHub resolves per job, so it is nothing to compare, which is what the rest of the module already does. `format()` was resolved with Python's semantics, so an attribute lookup crashed the script and a width specifier padded a name GitHub never pads. Fill `{0}` holes and escaped braces, and treat anything richer as unresolved. A matrix `include` or `exclude` row holding a value that is not a scalar lost that key and became an empty row, which excludes every combination. Report the row instead of quietly reshaping the matrix around it. --- .../check_workflow_job_name_collisions.py | 52 +++++++------- .../test_workflow_job_name_collisions.py | 70 +++++++++++++++++-- 2 files changed, 91 insertions(+), 31 deletions(-) diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py index 1e81acc4f90..9976bd9aa41 100644 --- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -24,13 +24,12 @@ real ` / Run tests` names rather than staying opaque. Whatever the sweep cannot work out is left out of the comparison and reported instead of guessed, because a guess that lands wrong fails a workflow GitHub -would have published perfectly well. A name left holding an expression over -`inputs`, `needs` or `matrix` reads differently from every job that runs it, so -it is one of those; an expression over `github` and the other contexts fixed for -the whole run reads the same everywhere, so two jobs carrying it still clash and -the check still says so. A matrix that is itself an expression or that lists -values which are not scalars, and a call this sweep cannot follow, go in the same -bucket. The cost is that a real clash hiding behind one of them goes unseen, +would have published perfectly well. A name still holding an expression once the +combination is filled in is one GitHub resolves per job, so it is one of those: +guessing that two jobs sharing such a template clash would fail workflows over a +context this sweep cannot read. A matrix that is itself an expression or that +lists values which are not scalars, an `include` or `exclude` row shaped the same +way, and a call this sweep cannot follow, go in the same bucket. The cost is that a real clash hiding behind one of them goes unseen, which leaves a merge no worse off than before this check existed, where the opposite direction would block work that was fine. @@ -62,8 +61,7 @@ MATRIX_REF: Final = re.compile(r"^matrix\.(?P[\w-]+)$") LITERAL: Final = re.compile(r"^'(?P[^']*)'$") FORMAT_CALL: Final = re.compile(r"^format\((?P.*)\)$", re.DOTALL) COMPARISON: Final = re.compile(r"^(?P.+?)\s*(?P==|!=)\s*(?P.+)$", re.DOTALL) -CONTEXT_REF: Final = re.compile(r"\b(?P[a-z][\w-]*)\s*\.") -RUN_FIXED_CONTEXTS: Final = frozenset({"github", "vars", "env", "runner", "secrets"}) +GITHUB_PLACEHOLDER: Final = re.compile(r"\{\{|\}\}|\{\d+\}") NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({}) NO_CALLERS: Final[frozenset[str]] = frozenset() SCALAR: Final = (str, int, float) @@ -149,15 +147,17 @@ def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, . return tuple((key, values) for key, values in listed if not isinstance(values, Opaque)) -def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...]: +def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...] | Opaque: + """One `include` or `exclude` row, or why the combinations they shape cannot be worked out.""" rows: Final = matrix.get(directive) if not isinstance(rows, Sequence) or isinstance(rows, str): return () - return tuple( - MappingProxyType({str(key): scalar_text(value) for key, value in row.items() if isinstance(value, SCALAR)}) - for row in rows - if isinstance(row, Mapping) - ) + mappings: Final = tuple(row for row in rows if isinstance(row, Mapping)) + if len(mappings) != len(rows): + return Opaque(f"a matrix `{directive}` row is not a mapping of values") + if any(not isinstance(value, SCALAR) for row in mappings for value in row.values()): + return Opaque(f"a matrix `{directive}` row holds a value that is not a plain scalar") + return tuple(MappingProxyType({str(key): scalar_text(value) for key, value in row.items()}) for row in mappings) def drops(row: Mapping[str, str], combination: Mapping[str, str]) -> bool: @@ -195,7 +195,11 @@ def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...] | Opaque: if isinstance(listed, Opaque): return listed rows: Final = directive_rows(matrix, "include") + if isinstance(rows, Opaque): + return rows dropped: Final = directive_rows(matrix, "exclude") + if isinstance(dropped, Opaque): + return dropped kept: Final = tuple( combination for combination in crossed_values(listed) if not any(drops(row, combination) for row in dropped) ) @@ -225,7 +229,10 @@ def split_outside(text: str, token: str) -> tuple[str, ...]: def formatted(template: str, arguments: Sequence[str]) -> str | None: - """A `format()` whose placeholders the arguments cannot fill resolves to nothing, never a crash.""" + """GitHub's `format()` fills `{0}`-style holes and escapes braces, so anything richer resolves to nothing.""" + residue: Final = GITHUB_PLACEHOLDER.sub("", template) + if "{" in residue or "}" in residue: + return None try: return template.format(*arguments) except (IndexError, KeyError, ValueError): @@ -280,18 +287,9 @@ def rendered(template: str, values: Mapping[str, str]) -> str: return EXPRESSION.sub(lambda span: resolved_span(span, values), template) -def stable(name: str) -> bool: - """An expression over the run's own contexts reads the same from every job, so two of them still clash.""" - return all( - reference.group("root") in RUN_FIXED_CONTEXTS - for span in EXPRESSION.finditer(name) - for reference in CONTEXT_REF.finditer(span.group("body")) - ) - - def comparable(name: str) -> bool: - """A leftover expression over `inputs`, `needs` or `matrix` names a different check run per job.""" - return EXPRESSION.search(name) is None or stable(name) + """A name still holding an expression is one GitHub resolves per job, so it is nothing to compare.""" + return EXPRESSION.search(name) is None def settled(names: Sequence[str]) -> Names: diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py index 39e61ff48d1..44a5adbbb33 100644 --- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -302,7 +302,7 @@ def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None: assert "`Analyze (go)` is published by 2 jobs" in found[0] -def test_two_jobs_sharing_a_template_over_the_run_itself_still_collide() -> None: +def test_two_jobs_sharing_a_template_over_the_run_itself_are_reported_rather_than_guessed() -> None: template: Final = ( "on: pull_request\njobs:\n {job}:\n name: ${{{{ github.event_name }}}}-build\n runs-on: ubuntu-latest\n" ) @@ -311,10 +311,72 @@ def test_two_jobs_sharing_a_template_over_the_run_itself_still_collide() -> None "b.yml": template.format(job="two"), } - found: Final = collisions(sources) + assert collisions(sources) == () + assert len(blind_spots(sources)) == 2 - assert len(found) == 1 - assert "is published by 2 jobs" in found[0] + +def test_a_name_reading_the_workflow_it_sits_in_is_not_called_a_collision() -> None: + template: Final = ( + "on: pull_request\njobs:\n {job}:\n name: ${{{{ github.workflow }}}} / build\n runs-on: ubuntu-latest\n" + ) + sources: Final = { + "a.yml": template.format(job="one"), + "b.yml": template.format(job="two"), + } + + assert collisions(sources) == () + assert exit_code(sources) == 0 + + +def test_a_format_call_python_accepts_but_github_does_not_publishes_nothing_to_compare() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: ${{ format('{0.real}', matrix.shard) }}\n" + " runs-on: ubuntu-latest\n strategy:\n matrix:\n shard: [core]\n" + ), + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_a_format_call_padding_its_argument_publishes_nothing_to_compare() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: ${{ format('{0:>8}', matrix.shard) }}\n" + " runs-on: ubuntu-latest\n strategy:\n matrix:\n shard: [core]\n" + ), + "b.yml": "on: pull_request\njobs:\n two:\n name: ' core'\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_an_exclude_row_that_is_not_a_mapping_is_reported_rather_than_skipped() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n v: [1, 2]\n exclude:\n - oops\n" + ), + "b.yml": "on: pull_request\njobs:\n other:\n name: build (1)\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_an_exclude_row_holding_a_non_scalar_never_drops_every_combination() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n v: [1, 2]\n exclude:\n - cfg: {k: 1}\n" + ), + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + assert blind_spots(sources) != () def test_two_jobs_sharing_a_template_that_reads_per_job_are_not_called_a_collision() -> None: From aa2c41f489eb3eba1de9983d641239b131057daa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 01:59:03 -0700 Subject: [PATCH 40/49] fix(ci): stop the name-collision check failing legal workflows An expression at `jobs..strategy` is legal on GitHub, but the model required a mapping there, so a workflow using one made the whole file unreadable and turned code-quality red. That job's names are now a blind spot like any other name the sweep cannot work out offline. A matrix whose `name:` holds no matrix value publishes that one name once per combination, which leaves a required context just as ambiguous as two jobs sharing a name, so it now reports instead of deduping. A file that does not parse as one YAML document is reported the way the module already promised, rather than escaping as a traceback. --- .../check_workflow_job_name_collisions.py | 52 ++++++++++++----- .../test_workflow_job_name_collisions.py | 58 +++++++++++++++++++ 2 files changed, 94 insertions(+), 16 deletions(-) diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py index 9976bd9aa41..941039e8340 100644 --- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -1,10 +1,12 @@ -"""Catch two workflow jobs that publish check runs under the same name. +"""Catch workflow jobs that publish check runs under the same name. A ruleset's required status check names a check run and GitHub matches it by that name alone. When two jobs publish the same name the required context stops mapping to the job that proves it: the commit carries two check runs under one name and nothing says which one the ruleset required. Both being green hides the clash completely, so the context quietly stops meaning what the ruleset intended. +One job lands in the same place when its `name:` holds no matrix value, since +every combination it runs then reports under that one name. `.github/workflows/auto-close-duplicates.yml` shipped a job id `test` while `.github/workflows/test-mcp.yml` already published the required `test` context, @@ -29,16 +31,17 @@ combination is filled in is one GitHub resolves per job, so it is one of those: guessing that two jobs sharing such a template clash would fail workflows over a context this sweep cannot read. A matrix that is itself an expression or that lists values which are not scalars, an `include` or `exclude` row shaped the same -way, and a call this sweep cannot follow, go in the same bucket. The cost is that a real clash hiding behind one of them goes unseen, -which leaves a merge no worse off than before this check existed, where the -opposite direction would block work that was fine. +way, a whole `strategy:` that comes from an expression, and a call this sweep +cannot follow, go in the same bucket. The cost is that a real clash hiding behind +one of them goes unseen, which leaves a merge no worse off than before this check +existed, where the opposite direction would block work that was fine. A job calling a local reusable workflow publishes one check run per job of the callee, named ` / ` and chained through however many levels of local calls it takes, which is why a caller's name never collides with a plain -job that happens to match it. A file under `.github/workflows/` that does not read as -a workflow at all is reported rather than skipped, since skipping it silently -would hide every job it holds. +job that happens to match it. A file under `.github/workflows/` that does not +read as one workflow at all is reported rather than skipped, since skipping it +silently would hide every job it holds. """ import itertools @@ -90,7 +93,7 @@ class Names: class Job(BaseModel): name: object = None uses: str | None = None - strategy: Mapping[str, object] = Field(default_factory=dict) + strategy: object = Field(default_factory=dict) class Workflow(BaseModel): @@ -104,7 +107,10 @@ def scalar_text(value: object) -> str: def parse(source: str) -> tuple[Workflow, object] | Unreadable: """The workflow plus its raw `on:` value, or why the file does not read as one.""" - parsed: Final = yaml.safe_load(source) + try: + parsed: Final = yaml.safe_load(source) + except yaml.YAMLError: + return Unreadable("it does not read as one YAML document") if not isinstance(parsed, dict): return Unreadable("its top level is not a mapping of workflow keys") try: @@ -186,6 +192,8 @@ def crossed_values(listed: Sequence[tuple[str, tuple[str, ...]]]) -> tuple[Mappi def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...] | Opaque: """One mapping per job the matrix produces, `exclude` applied before `include` as GitHub does.""" + if not isinstance(job.strategy, Mapping): + return Opaque("its whole `strategy` comes from an expression") matrix: Final = job.strategy.get("matrix") if matrix is None: return () @@ -304,7 +312,7 @@ def expand(template: str, job: Job) -> Names: if isinstance(combinations, Opaque): return Names((), (combinations.reason,)) over: Final = combinations or (NO_MATRIX,) - return settled(tuple(dict.fromkeys(rendered(template, values) for values in over))) + return settled(tuple(rendered(template, values) for values in over)) def suffixed(job_id: str, combination: Mapping[str, str]) -> str: @@ -408,13 +416,25 @@ def owners_by_name(sources: Mapping[str, str]) -> Iterator[tuple[str, tuple[str, yield name, tuple(owner for _, owner in pairs) +def clash(name: str, owners: Sequence[str]) -> str | None: + """Why one name is ambiguous, whether two jobs carry it or one job repeats it over its matrix.""" + jobs: Final = tuple(dict.fromkeys(owners)) + if len(jobs) > 1: + return ( + f"`{name}` is published by {len(jobs)} jobs: {', '.join(jobs)}. A required status check matching " + f"that name cannot say which job proves it; give one of them a distinct `name:` or job id." + ) + if len(owners) > 1: + return ( + f"`{name}` is published {len(owners)} times by {jobs[0]}, once per matrix combination. A required " + f"status check matching that name cannot say which run proves it; put a matrix value in its `name:`." + ) + return None + + def collisions(sources: Mapping[str, str]) -> tuple[str, ...]: - return tuple( - f"`{name}` is published by {len(owners)} jobs: {', '.join(owners)}. A required status check matching " - f"that name cannot say which job proves it; give one of them a distinct `name:` or job id." - for name, owners in owners_by_name(sources) - if len(owners) > 1 - ) + found: Final = tuple(clash(name, owners) for name, owners in owners_by_name(sources)) + return tuple(message for message in found if message is not None) def workflow_sources() -> Mapping[str, str]: diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py index 44a5adbbb33..ab6b1d2965e 100644 --- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -699,3 +699,61 @@ def test_an_include_row_naming_a_listed_key_extends_only_the_combinations_it_mat assert frozenset(name for name, _ in published(sources)) == frozenset({"3.12 fast"}) assert len(blind_spots(sources)) == 1 + + +def test_a_job_whose_whole_strategy_is_an_expression_is_reported_rather_than_rejecting_the_file() -> None: + sources: Final = { + "plan.yml": ( + "on: pull_request\njobs:\n plan:\n name: Plan\n runs-on: ubuntu-latest\n" + " fan:\n strategy: ${{ fromJSON(needs.plan.outputs.strategy) }}\n runs-on: ubuntu-latest\n" + ) + } + + assert unreadable(sources) == () + assert frozenset(name for name, _ in published(sources)) == frozenset({"Plan"}) + assert "`strategy` comes from an expression" in blind_spots(sources)[0] + assert exit_code(sources) == 0 + + +def test_one_job_publishing_one_name_for_every_matrix_combination_is_a_collision() -> None: + sources: Final = { + "unit.yml": ( + "on: pull_request\njobs:\n build:\n name: Run tests\n runs-on: ubuntu-latest\n" + ' strategy:\n matrix:\n python-version: ["3.12", "3.13"]\n' + ) + } + + found: Final = collisions(sources) + assert len(found) == 1 + assert "`Run tests` is published 2 times by unit.yml job `build`" in found[0] + assert exit_code(sources) == 1 + + +def test_a_name_carrying_a_matrix_value_publishes_one_name_per_combination_without_colliding() -> None: + sources: Final = { + "unit.yml": ( + "on: pull_request\njobs:\n build:\n name: Run tests ${{ matrix.python-version }}\n" + ' runs-on: ubuntu-latest\n strategy:\n matrix:\n python-version: ["3.12", "3.13"]\n' + ) + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"Run tests 3.12", "Run tests 3.13"}) + assert collisions(sources) == () + assert exit_code(sources) == 0 + + +def test_a_file_that_is_not_valid_yaml_is_reported_rather_than_raising() -> None: + sources: Final = {"broken.yml": "jobs:\n build: [\n"} + + assert unreadable(sources) == ( + "broken.yml sits in the workflows directory but it does not read as one YAML " + "document, so none of its jobs were checked.", + ) + assert exit_code(sources) == 1 + + +def test_a_file_holding_two_yaml_documents_is_reported_rather_than_raising() -> None: + sources: Final = {"two.yml": "on: pull_request\n---\non: push\n"} + + assert len(unreadable(sources)) == 1 + assert exit_code(sources) == 1 From 5a27e11263b2e267475ab39e287ead25920d43d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 02:32:50 -0700 Subject: [PATCH 41/49] fix(ci): compare names one workflow run settles the same way A `name:` whose only leftover expressions read a `github.` property other than `github.job` is filled in identically for every job of the run that publishes it, so two jobs of one workflow carrying it land on the same check run. Those names now compare against the other jobs of their own file instead of sitting in the blind-spot bucket. They stay out of the comparison across files, where two workflows can run on different events --- .../check_workflow_job_name_collisions.py | 67 ++++++++++++-- .../test_workflow_job_name_collisions.py | 91 ++++++++++++++++++- 2 files changed, 148 insertions(+), 10 deletions(-) diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py index 941039e8340..482f8386536 100644 --- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -27,12 +27,16 @@ real ` / Run tests` names rather than staying opaque. Whatever the sweep cannot work out is left out of the comparison and reported instead of guessed, because a guess that lands wrong fails a workflow GitHub would have published perfectly well. A name still holding an expression once the -combination is filled in is one GitHub resolves per job, so it is one of those: -guessing that two jobs sharing such a template clash would fail workflows over a -context this sweep cannot read. A matrix that is itself an expression or that -lists values which are not scalars, an `include` or `exclude` row shaped the same -way, a whole `strategy:` that comes from an expression, and a call this sweep -cannot follow, go in the same bucket. The cost is that a real clash hiding behind +combination is filled in is usually one GitHub resolves per job, so it is one of +those: guessing that two jobs sharing such a template clash would fail workflows +over a context this sweep cannot read. The exception is a name whose leftover +expressions all read a `github.` property other than `github.job`, which one run +fills in the same way for every job in it, so those are compared against the +other jobs of their own workflow and stay out of the comparison across files, +where two workflows can run on different events. A matrix that is itself an +expression or that lists values which are not scalars, an `include` or `exclude` +row shaped the same way, a whole `strategy:` that comes from an expression, and a +call this sweep cannot follow, go in the same bucket. The cost is that a real clash hiding behind one of them goes unseen, which leaves a merge no worse off than before this check existed, where the opposite direction would block work that was fine. @@ -64,6 +68,7 @@ MATRIX_REF: Final = re.compile(r"^matrix\.(?P[\w-]+)$") LITERAL: Final = re.compile(r"^'(?P[^']*)'$") FORMAT_CALL: Final = re.compile(r"^format\((?P.*)\)$", re.DOTALL) COMPARISON: Final = re.compile(r"^(?P.+?)\s*(?P==|!=)\s*(?P.+)$", re.DOTALL) +RUN_WIDE: Final = re.compile(r"^github\.(?!job\b)[\w.]+$") GITHUB_PLACEHOLDER: Final = re.compile(r"\{\{|\}\}|\{\d+\}") NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({}) NO_CALLERS: Final[frozenset[str]] = frozenset() @@ -88,6 +93,7 @@ class Names: known: tuple[str, ...] = () unknown: tuple[str, ...] = () + local: tuple[str, ...] = () class Job(BaseModel): @@ -300,10 +306,17 @@ def comparable(name: str) -> bool: return EXPRESSION.search(name) is None +def run_wide(name: str) -> bool: + """A name whose leftover expressions one workflow run fills in the same way for every job in it.""" + return all(RUN_WIDE.match(span.group("body").strip()) is not None for span in EXPRESSION.finditer(name)) + + def settled(names: Sequence[str]) -> Names: + unresolved: Final = tuple(name for name in names if not comparable(name)) return Names( tuple(name for name in names if comparable(name)), - tuple(f"its name stays `{name}`" for name in names if not comparable(name)), + tuple(f"its name stays `{name}`" for name in unresolved if not run_wide(name)), + tuple(name for name in unresolved if run_wide(name)), ) @@ -340,9 +353,15 @@ def joined(groups: Sequence[Names]) -> Names: return Names( tuple(name for group in groups for name in group.known), tuple(reason for group in groups for reason in group.unknown), + tuple(name for group in groups for name in group.local), ) +def tagged(names: Names) -> tuple[tuple[str, bool], ...]: + """Each name a job publishes beside whether only its own workflow's run settles it.""" + return (*((name, False) for name in names.known), *((name, True) for name in names.local)) + + def call_blocker(job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str]) -> str | None: path: Final = callee_path(job) if path is None: @@ -366,9 +385,15 @@ def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow], callers: for callee_id, callee_job in workflows[path].jobs.items() ) ) + composed: Final = tuple( + (f"{prefix} / {suffix}", prefix_local or suffix_local) + for prefix, prefix_local in tagged(prefixes) + for suffix, suffix_local in tagged(suffixes) + ) return Names( - tuple(f"{prefix} / {suffix}" for prefix in prefixes.known for suffix in suffixes.known), + tuple(name for name, is_local in composed if not is_local), (*prefixes.unknown, *suffixes.unknown), + tuple(name for name, is_local in composed if is_local), ) @@ -432,9 +457,33 @@ def clash(name: str, owners: Sequence[str]) -> str | None: return None +def local_published(sources: Mapping[str, str]) -> Iterator[tuple[tuple[str, str], str]]: + """Names their own workflow's run settles, keyed by the file whose run settles them.""" + for rel, job_id, names in scanned_jobs(sources): + for name in names.local: + yield (rel, name), f"job `{job_id}`" + + +def local_clash(rel: str, name: str, owners: Sequence[str]) -> str | None: + """Why one workflow's own run lands several of its jobs on one check run.""" + if len(owners) < 2: + return None + jobs: Final = tuple(dict.fromkeys(owners)) + return ( + f"`{name}` is published {len(owners)} times inside {rel}, by {', '.join(jobs)}. One run fills that " + f"expression in the same way throughout, so they all land on one check run; make the names differ." + ) + + +def local_clashes(sources: Mapping[str, str]) -> tuple[str, ...]: + grouped: Final = itertools.groupby(sorted(local_published(sources)), key=operator.itemgetter(0)) + found: Final = tuple(local_clash(rel, name, tuple(owner for _, owner in pairs)) for (rel, name), pairs in grouped) + return tuple(message for message in found if message is not None) + + def collisions(sources: Mapping[str, str]) -> tuple[str, ...]: found: Final = tuple(clash(name, owners) for name, owners in owners_by_name(sources)) - return tuple(message for message in found if message is not None) + return (*(message for message in found if message is not None), *local_clashes(sources)) def workflow_sources() -> Mapping[str, str]: diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py index ab6b1d2965e..aeb10985e1b 100644 --- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -302,7 +302,7 @@ def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None: assert "`Analyze (go)` is published by 2 jobs" in found[0] -def test_two_jobs_sharing_a_template_over_the_run_itself_are_reported_rather_than_guessed() -> None: +def test_two_workflows_sharing_a_run_wide_template_are_not_called_a_collision() -> None: template: Final = ( "on: pull_request\njobs:\n {job}:\n name: ${{{{ github.event_name }}}}-build\n runs-on: ubuntu-latest\n" ) @@ -311,10 +311,99 @@ def test_two_jobs_sharing_a_template_over_the_run_itself_are_reported_rather_tha "b.yml": template.format(job="two"), } + assert collisions(sources) == () + assert blind_spots(sources) == () + + +def test_two_jobs_of_one_workflow_sharing_a_run_wide_template_are_a_collision() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n" + " two:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n" + ), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "is published 2 times inside a.yml, by job `one`, job `two`" in found[0] + assert exit_code(sources) == 1 + + +def test_a_run_wide_template_carrying_a_matrix_value_does_not_collide_inside_one_workflow() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}-${{ matrix.shard }}\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n shard: [core, extras]\n" + ), + } + + assert collisions(sources) == () + assert blind_spots(sources) == () + + +def test_a_run_wide_name_repeated_over_a_matrix_by_one_job_is_a_collision() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n shard: [core, extras]\n" + ), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "is published 2 times inside a.yml, by job `one`" in found[0] + + +def test_a_name_reading_the_job_it_sits_in_stays_out_of_the_comparison() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.job }}-build\n runs-on: ubuntu-latest\n" + " two:\n name: ${{ github.job }}-build\n runs-on: ubuntu-latest\n" + ), + } + assert collisions(sources) == () assert len(blind_spots(sources)) == 2 +def test_a_run_wide_caller_name_collides_through_the_workflow_it_calls() -> None: + sources: Final = { + ".github/workflows/a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}\n uses: ./.github/workflows/c.yml\n" + " two:\n name: ${{ github.event_name }}\n uses: ./.github/workflows/c.yml\n" + ), + ".github/workflows/c.yml": "on:\n workflow_call:\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "github.event_name }} / build` is published 2 times inside .github/workflows/a.yml" in found[0] + + +def test_a_run_wide_name_inside_a_called_workflow_collides_under_the_caller() -> None: + sources: Final = { + ".github/workflows/a.yml": ("on: pull_request\njobs:\n one:\n uses: ./.github/workflows/c.yml\n"), + ".github/workflows/c.yml": ( + "on:\n workflow_call:\njobs:\n" + " build:\n name: ${{ github.event_name }}\n runs-on: ubuntu-latest\n" + " lint:\n name: ${{ github.event_name }}\n runs-on: ubuntu-latest\n" + ), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`one / ${{ github.event_name }}` is published 2 times inside .github/workflows/a.yml" in found[0] + + def test_a_name_reading_the_workflow_it_sits_in_is_not_called_a_collision() -> None: template: Final = ( "on: pull_request\njobs:\n {job}:\n name: ${{{{ github.workflow }}}} / build\n runs-on: ubuntu-latest\n" From 4edf6f4dc530f88afe09d2bb986a4bdc5a962cc0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 02:59:16 -0700 Subject: [PATCH 42/49] fix(ci): keep an expression matrix directive out of the comparison `include:` or `exclude:` written as `${{ ... }}` read back as a string, and the sweep treated that as the directive being absent, so it expanded every combination GitHub would have dropped. A job whose `name:` holds no matrix value then looked like it repeated one name across combinations that never run. An absent directive still means no rows; anything that is not a list of rows now joins the names left out of the comparison --- .../check_workflow_job_name_collisions.py | 4 ++- .../test_workflow_job_name_collisions.py | 32 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py index 482f8386536..ae2c1d80c8f 100644 --- a/tests/code_coverage_tests/check_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -162,8 +162,10 @@ def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, . def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...] | Opaque: """One `include` or `exclude` row, or why the combinations they shape cannot be worked out.""" rows: Final = matrix.get(directive) - if not isinstance(rows, Sequence) or isinstance(rows, str): + if rows is None: return () + if not isinstance(rows, Sequence) or isinstance(rows, str): + return Opaque(f"a matrix `{directive}` is itself an expression rather than a list of rows") mappings: Final = tuple(row for row in rows if isinstance(row, Mapping)) if len(mappings) != len(rows): return Opaque(f"a matrix `{directive}` row is not a mapping of values") diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py index aeb10985e1b..0f5ba43bd7a 100644 --- a/tests/code_coverage_tests/test_workflow_job_name_collisions.py +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -846,3 +846,35 @@ def test_a_file_holding_two_yaml_documents_is_reported_rather_than_raising() -> assert len(unreadable(sources)) == 1 assert exit_code(sources) == 1 + + +def test_an_exclude_that_is_itself_an_expression_is_reported_rather_than_ignored() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: build\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n python: ['3.11', '3.12']\n" + " exclude: ${{ fromJson(vars.SKIP) }}\n" + ), + } + + found: Final = blind_spots(sources) + + assert collisions(sources) == () + assert len(found) == 1 + assert "a matrix `exclude` is itself an expression" in found[0] + + +def test_an_include_that_is_itself_an_expression_is_reported_rather_than_ignored() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: build-${{ matrix.python }}\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n python: ['3.11']\n" + " include: ${{ fromJson(vars.EXTRA) }}\n" + ), + } + + found: Final = blind_spots(sources) + + assert collisions(sources) == () + assert len(found) == 1 + assert "a matrix `include` is itself an expression" in found[0] From f40c3842dc4cebbf1eb02379a492acb1d0c87605 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 10:08:19 -0700 Subject: [PATCH 43/49] fix(router): serialize heuristic tuning quota filters for Prisma --- litellm/repositories/model_repository.py | 3 +-- .../repositories/test_repositories.py | 22 ++++++++++++++++++- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index 540f197044f..d24eb8ffc62 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -4,7 +4,6 @@ Model repository for database operations on LiteLLM_ProxyModelTable. import json from collections.abc import Mapping, Sequence -from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.models.model import LiteLLM_ProxyModelTable @@ -109,7 +108,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): async def find_all_except(self, model_id: str) -> Sequence[LiteLLM_ProxyModelTable]: """Find every model except the row currently being updated.""" records: Final = await self.table.find_many( - where=MappingProxyType({"model_id": MappingProxyType({"not": model_id})}) + where={"model_id": {"not": model_id}} # mutable-ok: Prisma requires plain dicts for query serialization ) return tuple(self._to_model_list(records)) diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index 758d379f22c..a890d7ceed0 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -4,10 +4,13 @@ Tests for gateway repository layer. import json from datetime import datetime -from typing import Any, Dict, List, Optional +from types import SimpleNamespace +from typing import Any, Dict, Final, List, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from prisma import models as prisma_models +from prisma.builder import QueryBuilder from litellm.models.base import DomainModel from litellm.models.budget import LiteLLM_BudgetTable @@ -307,6 +310,23 @@ class TestModelRepository: client = MockPrismaClient() return ModelRepository(client) + @pytest.mark.asyncio + async def test_find_all_except_serializes_exclusion_for_prisma(self) -> None: + find_many: Final = AsyncMock(return_value=[]) + client: Final = SimpleNamespace( + db=SimpleNamespace(litellm_proxymodeltable=SimpleNamespace(find_many=find_many)) + ) + + await ModelRepository(client).find_all_except("current-model") + + find_many.assert_awaited_once() + query: Final = QueryBuilder( + method="find_many", + model=prisma_models.LiteLLM_ProxyModelTable, + arguments=find_many.call_args.kwargs, + ).build_query() + assert 'where: { model_id: { not: "current-model" } }' in " ".join(query.split()) + def test_table_is_wrapped_for_config_sync(self, repo): from litellm.proxy.common_utils.config_sync_pubsub import ( _PublishOnWriteActions, From 217cb126236cb2a73245f7c3ac5525f88dfc0f7c Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 7 Sep 2026 10:43:45 -0700 Subject: [PATCH 44/49] refactor(rust): remove per-request enablement arguments (#39928) * refactor(rust): remove per-request enablement arguments * fix(rust): remove ignored transcription enablement * refactor(rust): remove OCR-specific bridge controls --- .../litellm_core_utils/get_litellm_params.py | 9 +-- litellm/llms/custom_httpx/llm_http_handler.py | 11 +-- litellm/ocr/main.py | 11 +-- litellm/rust_bridge/chat_completions.py | 3 +- litellm/rust_bridge/configuration.py | 33 +------- litellm/rust_bridge/ocr.py | 46 +++-------- litellm/rust_bridge/transcription.py | 1 - litellm/types/router.py | 1 - .../test_rust_bridge_messages.py | 26 +----- .../test_get_litellm_params.py | 29 ------- .../chat/test_anthropic_chat_handler.py | 8 +- .../chat/test_bedrock_converse_handler.py | 12 ++- .../custom_httpx/test_llm_http_handler.py | 22 +++-- tests/test_litellm/ocr/test_rust_bridge.py | 80 ++++++++----------- .../responses/test_rust_bridge_websocket.py | 21 ++--- .../rust_bridge/test_chat_completions.py | 38 ++++----- .../rust_bridge/test_configuration.py | 61 +++----------- .../test_audio_transcription_rust_bridge.py | 4 +- tests/test_litellm/test_utils.py | 46 ----------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 - 20 files changed, 109 insertions(+), 357 deletions(-) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 389e6f7f501..1fd79db15a6 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -21,13 +21,10 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( } ) -# The per-deployment Rust opt-in. -RUST_KWARG_KEY: Final = "rust" - # Keys `completion()` forwards from its own kwargs into `get_litellm_params`, # which are otherwise invisible to it because that call site passes explicit # named arguments rather than `**kwargs`. -FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY}) +FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS # Pre-define optional kwargs keys as frozenset for O(1) lookups # These are extracted from kwargs only if present, avoiding unnecessary .get() calls @@ -58,10 +55,6 @@ OPTIONAL_KWARGS_KEYS: Final = ( "itpm", "otpm", "use_xai_oauth", - # The per-deployment Rust opt-in. `all_litellm_params` keeps it out - # of the provider body; this keeps it *in* litellm_params, which is - # where the chat completions handlers read it from. - RUST_KWARG_KEY, } ) | AWS_CREDENTIAL_KWARGS_KEYS diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f281c249c72..552b1549db8 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -163,13 +163,10 @@ from litellm.utils import ( def _rust_responses_websocket_enabled( custom_llm_provider: str | None, - litellm_params: GenericLiteLLMParams, ) -> bool: from litellm.rust_bridge.configuration import rust_enabled - raw_request_override: Final = litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - return custom_llm_provider == "openai" and rust_enabled(request_override=request_override) + return custom_llm_provider == "openai" and rust_enabled() from .http_handler import get_shared_realtime_ssl_context @@ -2403,9 +2400,7 @@ class BaseLLMHTTPHandler: return None from litellm.rust_bridge.configuration import rust_enabled - raw_request_override: Final = litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - if not rust_enabled(request_override=request_override): + if not rust_enabled(): return None if has_agentic_hook: return None @@ -6514,7 +6509,7 @@ class BaseLLMHTTPHandler: @asynccontextmanager async def _backend_connection(): - if _rust_responses_websocket_enabled(custom_llm_provider, litellm_params): + if _rust_responses_websocket_enabled(custom_llm_provider): from litellm.rust_bridge import responses_websocket as rust_responses_websocket rust_backend: Final = await rust_responses_websocket.connect( diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 6c68971f8d5..56c2292d00c 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -29,6 +29,7 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.rust_bridge.configuration import rust_enabled from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -196,12 +197,6 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS -def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool: - raw_request_override: Final = prepared_request.litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override) - - def _rust_bridge_optional_params( prepared_request: _PreparedOCRRequest, resolve_secret: Callable[[str], str | None], @@ -430,7 +425,7 @@ async def aocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): + if _rust_ocr_supported(prepared) and rust_enabled(): from litellm.secret_managers.main import get_secret_str rust_response: Final = await _run_rust_aocr( @@ -702,7 +697,7 @@ def ocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): + if _rust_ocr_supported(prepared) and rust_enabled(): from litellm.secret_managers.main import get_secret_str rust_response: Final = _run_rust_ocr( diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index c599667ab17..674bd8847f7 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -247,8 +247,7 @@ def rust_chat_completions_accepts( return False if stream: return False - request_override: Final = litellm_params.get("rust") if litellm_params is not None else None - if not rust_enabled(request_override=request_override if isinstance(request_override, bool) else None): + if not rust_enabled(): return False if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params): verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path") diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 515ab6edef1..5582027bb5d 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -1,13 +1,11 @@ from __future__ import annotations import os -import warnings from typing import Final DEFAULT_RUST_ENABLED: Final = False _TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"}) _GLOBAL_ENV_NAME: Final = "LITELLM_RUST" -_LEGACY_OCR_ENV_NAME: Final = "LITELLM_USE_RUST_OCR" class _RustConfiguration: @@ -26,49 +24,24 @@ def _parse_env_bool(value: str | None) -> bool | None: def resolve_rust_enabled( *, - request_override: bool | None, process_override: bool | None, environment_override: bool | None, - legacy_environment_override: bool | None = None, release_default: bool = DEFAULT_RUST_ENABLED, ) -> bool: - if request_override is not None: - return request_override if process_override is not None: return process_override if environment_override is not None: return environment_override - if legacy_environment_override is not None: - return legacy_environment_override return release_default -def rust_enabled(*, request_override: bool | None = None) -> bool: - if request_override is not None: - return request_override - process_override: Final = _CONFIGURATION.override - if process_override is not None: - return process_override - global_override: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)) - legacy_override: Final = None if global_override is not None else _parse_env_bool(os.getenv(_LEGACY_OCR_ENV_NAME)) - if legacy_override is not None: - warnings.warn( - f"{_LEGACY_OCR_ENV_NAME} is deprecated; use {_GLOBAL_ENV_NAME} instead", - DeprecationWarning, - stacklevel=2, - ) +def rust_enabled() -> bool: return resolve_rust_enabled( - request_override=None, - process_override=None, - environment_override=global_override, - legacy_environment_override=legacy_override, + process_override=_CONFIGURATION.override, + environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)), ) -def rust_ocr_enabled(*, request_override: bool | None = None) -> bool: - return rust_enabled(request_override=request_override) - - def reset_rust_configuration() -> None: _CONFIGURATION.override = None diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 86038438f57..b7fdb5a98ef 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -7,12 +7,9 @@ from typing import Final, Protocol, cast # noqa: TID251 # native extension exp import httpx -from litellm.rust_bridge import configuration as _configuration +from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds -rust_ocr_enabled = _configuration.rust_ocr_enabled -rust = _configuration.rust - class RustOcr(Protocol): def __call__( @@ -44,49 +41,24 @@ class RustAocr(Protocol): raise NotImplementedError -class _Unset: - pass +def _as_ocr(value: object) -> RustOcr | None: + return cast(RustOcr, value) if callable(value) else None -_UNSET: Final[_Unset] = _Unset() +def _as_aocr(value: object) -> RustAocr | None: + return cast(RustAocr, value) if callable(value) else None -_rust_ocr_impl: RustOcr | None = None -_rust_aocr_impl: RustAocr | None = None - - -def set_rust_ocr( - *, - ocr: RustOcr | None | _Unset = _UNSET, - aocr: RustAocr | None | _Unset = _UNSET, -) -> None: - global _rust_ocr_impl, _rust_aocr_impl - if not isinstance(ocr, _Unset): - _rust_ocr_impl = ocr - if not isinstance(aocr, _Unset): - _rust_aocr_impl = aocr +_OCR: Final = NativeBinding("ocr", validate=_as_ocr) +_AOCR: Final = NativeBinding("aocr", validate=_as_aocr) def load_rust_ocr() -> RustOcr | None: - if _rust_ocr_impl is not None: - return _rust_ocr_impl - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustOcr, native_bridge.ocr) + return _OCR.load() def load_rust_aocr() -> RustAocr | None: - if _rust_aocr_impl is not None: - return _rust_aocr_impl - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustAocr, getattr(native_bridge, "aocr", None)) + return _AOCR.load() def ocr( diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index 3d71f6f8a50..6c81786accd 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -56,7 +56,6 @@ _STATE: Final = _RustTranscriptionState() def configure_rust_transcription( - enabled: bool = True, *, transcription: RustTranscription | None | _Unset = _UNSET, atranscription: RustAtranscription | None | _Unset = _UNSET, diff --git a/litellm/types/router.py b/litellm/types/router.py index 2dad22751de..0db482d8a58 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -307,7 +307,6 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): """ custom_llm_provider: str | None = None - rust: bool | None = None tpm: int | None = None rpm: int | None = None itpm: int | None = None diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index b2cf253d164..a30474245c6 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -126,16 +126,6 @@ def test_load_rust_messages_returns_injected_impl(): assert rust_messages.load_rust_messages() is bridge -def test_bare_rust_still_toggles_ocr(): - from litellm.rust_bridge.ocr import rust_ocr_enabled - - litellm.rust(True) - assert rust_ocr_enabled() is True - - litellm.rust(False) - assert rust_ocr_enabled() is False - - def test_load_rust_amessages_returns_injected_impl(): bridge = RecordingAsyncMessages() litellm.rust(True) @@ -214,7 +204,7 @@ async def test_amessages_wrapper_forwards_args(): def _gate(**overrides): kwargs = { "custom_llm_provider": "azure_ai", - "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True), + "litellm_params": GenericLiteLLMParams(api_key="sk-azure"), "has_agentic_hook": False, "model": "claude-sonnet-4-5", "api_key": "sk-azure", @@ -282,18 +272,6 @@ async def test_gate_uses_process_enable_without_request_override(): assert bridge.calls[0]["custom_llm_provider"] == "azure_ai" -@pytest.mark.asyncio -async def test_gate_skips_rust_when_flag_false(): - bridge = ExplodingAsyncMessages() - litellm.rust(True) - rust_messages.set_rust_messages(amessages=bridge) - - response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False)) - - assert response is None - assert bridge.calls == 0 - - @pytest.mark.asyncio async def test_gate_invokes_rust_for_native_anthropic_provider(): bridge = RecordingAsyncMessages() @@ -302,7 +280,7 @@ async def test_gate_invokes_rust_for_native_anthropic_provider(): response = await _gate( custom_llm_provider="anthropic", - litellm_params=GenericLiteLLMParams(api_key="sk-ant", rust=True), + litellm_params=GenericLiteLLMParams(api_key="sk-ant"), api_key="sk-ant", api_base="https://api.anthropic.com", headers={"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"}, diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index 956da571d43..fb4cb494bee 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -215,32 +215,3 @@ class TestMetadataFallsBackToLitellmMetadata: assert result["metadata"] is not litellm_metadata result["metadata"].pop("trace_id") assert litellm_metadata == {"trace_id": "trace-1"} - - -class TestRustOptIn: - """`rust: true` is a litellm param, so it has to reach `litellm_params`. - - `all_litellm_params` keeps it out of the provider body; without it also - being carried into `litellm_params` the chat completions handlers cannot - see the opt-in and the Rust path is silently never taken. - """ - - def test_rust_is_an_optional_kwargs_key(self): - assert "rust" in _OPTIONAL_KWARGS_KEYS - - def test_rust_is_forwarded_from_completion_kwargs(self): - from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS - - assert "rust" in FORWARDED_KWARGS_KEYS - - def test_rust_survives_into_litellm_params(self): - params = get_litellm_params(rust=True) - assert params["rust"] is True - - def test_rust_is_absent_when_the_deployment_did_not_set_it(self): - assert "rust" not in get_litellm_params() - - def test_rust_stays_out_of_the_provider_body(self): - from litellm.types.utils import all_litellm_params - - assert "rust" in all_litellm_params diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index 043537f8c1f..b4b173b20c3 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2256,7 +2256,7 @@ class TestRustChatCompletionsHook: def _reset_bridge(self, monkeypatch): from litellm.rust_bridge import chat_completions as bridge - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") bridge.set_rust_chat_completions( chat_completions=None, achat_completions=None, decline=None ) @@ -2282,7 +2282,7 @@ class TestRustChatCompletionsHook: "logging_obj": MagicMock(), "optional_params": {"max_tokens": 16}, "timeout": 30.0, - "litellm_params": {"rust": True}, + "litellm_params": {}, "acompletion": False, "headers": {}, "client": None, @@ -2366,7 +2366,8 @@ class TestRustChatCompletionsHook: ) assert seen["call"][0]["optional_params"]["max_tokens"] == 7 - def test_without_the_opt_in_the_core_is_never_consulted(self): + def test_without_the_opt_in_the_core_is_never_consulted(self, monkeypatch): + monkeypatch.setenv("LITELLM_RUST", "0") from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -2587,6 +2588,7 @@ class TestRustChatCompletionsHook: def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch): """The suppression must not swallow the log on the ordinary path.""" + monkeypatch.setenv("LITELLM_RUST", "0") from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion from litellm.llms.anthropic.chat.transformation import AnthropicConfig diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index c4d6896b17b..f34b8eb1fb9 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -48,7 +48,7 @@ RESOLVED_CREDENTIALS = Credentials( @pytest.fixture(autouse=True) def reset_bridge(monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") bridge.set_rust_chat_completions( chat_completions=None, achat_completions=None, decline=None ) @@ -87,7 +87,7 @@ def _completion_kwargs(**overrides): "optional_params": {"maxTokens": 16}, "acompletion": False, "timeout": 30.0, - "litellm_params": {"rust": True}, + "litellm_params": {}, "extra_headers": None, "client": None, "api_key": None, @@ -157,7 +157,8 @@ def test_the_core_receives_the_untranslated_openai_messages(): ] -def test_without_the_opt_in_the_core_is_never_consulted(): +def test_without_the_opt_in_the_core_is_never_consulted(monkeypatch): + monkeypatch.setenv("LITELLM_RUST", "0") seen = _inject() try: _run(litellm_params={}) @@ -401,9 +402,10 @@ def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines(): assert logging_obj.pre_call.call_count == 1 -def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(): +def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(monkeypatch): """The suppression must not swallow the log on a request the gate declined, so a deployment with no `rust` flag keeps exactly the log it always had.""" + monkeypatch.setenv("LITELLM_RUST", "0") logging_obj = MagicMock() response = _run( logging_obj=logging_obj, @@ -491,6 +493,7 @@ def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monke """With only `AWS_BEARER_TOKEN_BEDROCK` configured boto3 resolves no credentials at all. Preparing the Rust handoff must not dereference that None: the bearer token signs the request on its own.""" + monkeypatch.setenv("LITELLM_RUST", "0") monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") client = _sync_client_returning_converse_response() @@ -520,6 +523,7 @@ def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, co """The deployment's AWS profile does not exist, so resolving SigV4 credentials raises; a bearer-token deployment must still serve the request, since the bearer token alone signs it.""" + monkeypatch.setenv("LITELLM_RUST", "0") if configured_through == "env_var": monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") else: diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 023e2d8843f..5c6cc4a52da 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2689,20 +2689,18 @@ async def test_generic_http_handler_async_streaming_forwards_provider_response_h @pytest.mark.parametrize( - "custom_llm_provider, litellm_params, expected", - [ - ("openai", GenericLiteLLMParams(rust=True), True), - ("openai", GenericLiteLLMParams(), False), - ("openai", GenericLiteLLMParams(rust=False), False), - ("azure", GenericLiteLLMParams(rust=True), False), - ("hosted_vllm", GenericLiteLLMParams(rust=True), False), - (None, GenericLiteLLMParams(rust=True), False), - ], + "custom_llm_provider, enabled, expected", + [("openai", True, True), ("openai", False, False), ("azure", True, False), + ("hosted_vllm", True, False), (None, True, False)], ) -def test_the_rust_responses_websocket_needs_both_openai_and_the_rust_flag( - custom_llm_provider, litellm_params, expected +def test_the_rust_responses_websocket_needs_openai_and_process_enablement( + custom_llm_provider, enabled, expected, monkeypatch ): - assert _rust_responses_websocket_enabled(custom_llm_provider, litellm_params) is expected + from litellm.rust_bridge import configuration + + configuration.reset_rust_configuration() + monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0") + assert _rust_responses_websocket_enabled(custom_llm_provider) is expected def test_a_plain_callback_does_not_advertise_a_pre_call_deployment_hook(monkeypatch): diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 1c2e07e0d24..d28184b96a6 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -17,6 +17,7 @@ from litellm.rust_bridge import configuration # explicitly via importlib rather than attribute traversal. ocr_main = importlib.import_module("litellm.ocr.main") rust_bridge = importlib.import_module("litellm.rust_bridge.ocr") +rust_bridge_bindings = importlib.import_module("litellm.rust_bridge.bindings") rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") MODEL = "mistral/mistral-ocr-latest" @@ -215,11 +216,13 @@ def build_prepared_request( @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge.set_rust_ocr(ocr=None, aocr=None) + rust_bridge._OCR.reset() + rust_bridge._AOCR.reset() configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - rust_bridge.set_rust_ocr(ocr=None, aocr=None) + rust_bridge._OCR.reset() + rust_bridge._AOCR.reset() configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -229,7 +232,7 @@ def fake_bridge(): """Enable the Rust path with an injected recording bridge (no native wheel).""" bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) return bridge @@ -238,34 +241,14 @@ def fake_async_bridge(): """Enable the async Rust path with an injected recording bridge.""" bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._AOCR.override(bridge) return bridge -def test_rust_toggles_flag(): - assert rust_bridge.rust_ocr_enabled() is False - litellm.rust(True) - assert rust_bridge.rust_ocr_enabled() is True - litellm.rust(False) - assert rust_bridge.rust_ocr_enabled() is False - - -def test_env_var_enables_rust_ocr(monkeypatch): - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert rust_bridge.rust_ocr_enabled() is True - - -def test_explicit_false_overrides_process_enable(): - litellm.rust(True) - - assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False - - def test_load_rust_ocr_returns_injected_impl(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) assert rust_bridge.load_rust_ocr() is bridge @@ -329,7 +312,7 @@ def test_native_bridge_available_reflects_loader(monkeypatch): def test_load_rust_aocr_returns_injected_impl(): bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._AOCR.override(bridge) assert rust_bridge.load_rust_aocr() is bridge @@ -338,7 +321,8 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) + rust_bridge._OCR.override(bridge) + rust_bridge._AOCR.override(async_bridge) litellm.rust(False) assert rust_bridge.load_rust_ocr() is bridge @@ -350,16 +334,18 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): def test_explicit_ocr_none_clears_injected_impl(monkeypatch): monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + rust_bridge_bindings, "get_native_bridge", lambda: None, ) bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) + rust_bridge._OCR.override(bridge) + rust_bridge._AOCR.override(async_bridge) - rust_bridge.set_rust_ocr(ocr=None, aocr=None) + rust_bridge._OCR.override(None) + rust_bridge._AOCR.override(None) assert rust_bridge.load_rust_ocr() is None assert rust_bridge.load_rust_aocr() is None @@ -368,7 +354,7 @@ def test_load_rust_ocr_none_when_extension_absent(monkeypatch): """With no injected impl and no compiled wheel, the loader returns None so the caller degrades to the Python path instead of raising ImportError.""" monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + rust_bridge_bindings, "get_native_bridge", lambda: None, ) @@ -385,7 +371,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + rust_bridge_bindings, "get_native_bridge", lambda: fake_module, ) @@ -406,7 +392,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) response = rust_bridge.ocr( model="mistral-ocr-latest", document=DOCUMENT, @@ -441,7 +427,7 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._AOCR.override(bridge) response = await rust_bridge.aocr( model="mistral-ocr-maas", document=DOCUMENT, @@ -470,7 +456,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): bridge = RecordingBridge() logging_obj = RecordingLogging() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) response = ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -503,7 +489,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request(api_key=None, timeout=None), @@ -516,7 +502,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): def test_run_rust_ocr_prefers_explicit_key_over_resolver(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) def _resolver(name: str) -> str | None: raise AssertionError(f"resolver should not be called for {name}") @@ -536,7 +522,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): bridge = RecordingBridge() resolver_calls = [] litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) def _resolver(name): resolver_calls.append(name) @@ -559,7 +545,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -586,7 +572,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) def _resolver(name: str) -> str | None: return { @@ -610,7 +596,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -628,7 +614,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -649,7 +635,7 @@ def test_run_rust_ocr_runs_pre_call_logging(): logging_obj = RecordingLogging() bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -737,7 +723,7 @@ def test_ocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=RaisingBridge()) + rust_bridge._OCR.override(RaisingBridge()) with pytest.raises(CapturedException): litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -783,7 +769,7 @@ async def test_aocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=RaisingAsyncBridge()) + rust_bridge._AOCR.override(RaisingAsyncBridge()) with pytest.raises(CapturedException): await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -812,9 +798,7 @@ def test_ocr_does_not_route_to_rust_when_disabled(): """With the flag off, the bridge must not be consulted even if an impl exists.""" bridge = RecordingBridge() litellm.rust(False) - rust_bridge.set_rust_ocr(ocr=bridge) - - assert rust_bridge.rust_ocr_enabled() is False + rust_bridge._OCR.override(bridge) # The impl stays available for injection, but the disabled flag gates usage, # so ocr() never reaches the Rust path (asserted via the enabled-path test). assert bridge.calls == [] diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index 4b446368dbe..74d96bda336 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -4,7 +4,6 @@ import pytest from litellm.llms.custom_httpx.llm_http_handler import _rust_responses_websocket_enabled from litellm.rust_bridge import configuration, responses_websocket -from litellm.types.router import GenericLiteLLMParams class _FakeNativeConnection: @@ -48,22 +47,12 @@ def reset_responses_websocket(): configuration.reset_rust_configuration() -def test_rust_websocket_bridge_is_disabled_without_flag() -> None: - assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) - assert not _rust_responses_websocket_enabled("anthropic", GenericLiteLLMParams(rust=True)) - assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=True)) - - -def test_explicit_false_overrides_process_enable() -> None: +def test_rust_websocket_bridge_uses_process_enablement() -> None: + configuration.rust(False) + assert not _rust_responses_websocket_enabled("openai") configuration.rust(True) - - assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=False)) - - -def test_process_enable_applies_without_request_override() -> None: - configuration.rust(True) - - assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) + assert _rust_responses_websocket_enabled("openai") + assert not _rust_responses_websocket_enabled("anthropic") @pytest.mark.asyncio diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 0489f4ff017..b2fd2e6dcc0 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -10,8 +10,8 @@ from __future__ import annotations import pytest import litellm -from litellm.rust_bridge import chat_completions as bridge from litellm.rust_bridge import configuration +from litellm.rust_bridge import chat_completions as bridge from litellm.types.utils import ModelResponse RUST_RESPONSE = { @@ -67,10 +67,11 @@ def _hide_native_bridge(monkeypatch): @pytest.fixture(autouse=True) -def reset_bridge(): +def reset_bridge(monkeypatch): """Every test starts with no injected callables, and leaves none behind.""" bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) configuration.reset_rust_configuration() + monkeypatch.setenv("LITELLM_RUST", "1") yield bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) configuration.reset_rust_configuration() @@ -112,7 +113,7 @@ def _accepts(**overrides) -> bool: "messages": MESSAGES, "optional_params": {"max_tokens": 16}, "custom_llm_provider": "anthropic", - "litellm_params": {"rust": True}, + "litellm_params": {}, "stream": None, } kwargs.update(overrides) @@ -126,23 +127,16 @@ class TestGate: bridge.set_rust_chat_completions(decline=gate) assert _accepts(litellm_params={}) is False assert _accepts(litellm_params=None) is False - assert _accepts(litellm_params={"rust": False}) is False assert gate.calls == [], "the gate must not be consulted before opt-in" def test_accepts_when_the_deployment_opted_in_and_the_core_agrees(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) assert _accepts() is True assert gate.calls[0]["model"] == "claude-sonnet-4-5" assert gate.calls[0]["custom_llm_provider"] == "anthropic" - def test_explicit_false_overrides_process_enable(self): - bridge.set_rust_chat_completions(decline=_RecordingDecline()) - configuration.rust(True) - - assert _accepts(litellm_params={"rust": False}) is False - def test_process_enable_applies_without_request_override(self): bridge.set_rust_chat_completions(decline=_RecordingDecline()) configuration.rust(True) @@ -155,7 +149,7 @@ class TestGate: assert _accepts(litellm_params={}) is True def test_declines_streaming_and_providers_off_the_path(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) assert _accepts(stream=True) is False @@ -170,10 +164,10 @@ class TestGate: handed `optional_params` only, so accepting here would send the request to Anthropic with the abuse-detection attribution silently missing. """ - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) - assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": "u-123"}}) is False + assert _accepts(litellm_params={"metadata": {"user_id": "u-123"}}) is False assert gate.calls == [], "the core must not be consulted for a request it cannot see the key of" # Bedrock's Converse transform reads no `user_id`, and an Anthropic request @@ -182,13 +176,13 @@ class TestGate: _accepts( custom_llm_provider="bedrock", model="bedrock/us-east-1/anthropic.claude-v2", - litellm_params={"rust": True, "metadata": {"user_id": "u-123"}}, + litellm_params={"metadata": {"user_id": "u-123"}}, ) is True ) - assert _accepts(litellm_params={"rust": True, "metadata": {"trace_id": "t-1"}}) is True - assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": None}}) is True - assert _accepts(litellm_params={"rust": True, "metadata": None}) is True + assert _accepts(litellm_params={"metadata": {"trace_id": "t-1"}}) is True + assert _accepts(litellm_params={"metadata": {"user_id": None}}) is True + assert _accepts(litellm_params={"metadata": None}) is True def test_declines_a_bedrock_request_while_the_proxy_owns_request_metadata(self, monkeypatch): """`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the @@ -196,7 +190,7 @@ class TestGate: evicting a caller-supplied one. The core can do neither, so an operator who armed `bedrock_request_metadata_fields` keeps the Python path. """ - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) bedrock = { @@ -213,17 +207,17 @@ class TestGate: assert _accepts(**bedrock) is True, "the decline follows the operator's opt-in alone" def test_declines_when_the_core_declines(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") bridge.set_rust_chat_completions(decline=_RecordingDecline("streaming")) assert _accepts() is False def test_declines_when_the_bridge_is_unavailable(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") _hide_native_bridge(monkeypatch) assert _accepts() is False def test_declines_when_the_gate_itself_raises(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") def exploding(**_kwargs): raise RuntimeError("boom") diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/test_litellm/rust_bridge/test_configuration.py index 15f69f95335..aff9d5acac1 100644 --- a/tests/test_litellm/rust_bridge/test_configuration.py +++ b/tests/test_litellm/rust_bridge/test_configuration.py @@ -10,7 +10,6 @@ from typing import Final import pytest from litellm.rust_bridge import configuration -from litellm.rust_bridge import ocr as rust_ocr @pytest.fixture(autouse=True) @@ -19,42 +18,31 @@ def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest ) -> Generator[None]: configuration.reset_rust_configuration() monkeypatch.delenv("LITELLM_RUST", raising=False) - monkeypatch.delenv("LITELLM_USE_RUST_OCR", raising=False) - rust_ocr.set_rust_ocr(ocr=None, aocr=None) yield configuration.reset_rust_configuration() - rust_ocr.set_rust_ocr(ocr=None, aocr=None) @pytest.mark.parametrize( - ("request_override", "process", "environment", "legacy_environment", "release_default", "expected"), + ("process", "environment", "release_default", "expected"), ( - (False, True, True, True, True, False), - (True, False, False, False, False, True), - (None, False, True, True, True, False), - (None, True, False, False, False, True), - (None, None, False, True, True, False), - (None, None, True, False, False, True), - (None, None, None, False, True, False), - (None, None, None, True, False, True), - (None, None, None, None, False, False), - (None, None, None, None, True, True), + (False, True, True, False), + (True, False, False, True), + (None, False, True, False), + (None, True, False, True), + (None, None, False, False), + (None, None, True, True), ), ) def test_resolution_precedence( - request_override: bool | None, process: bool | None, environment: bool | None, - legacy_environment: bool | None, release_default: bool, expected: bool, ) -> None: assert ( configuration.resolve_rust_enabled( - request_override=request_override, process_override=process, environment_override=environment, - legacy_environment_override=legacy_environment, release_default=release_default, ) is expected @@ -71,7 +59,6 @@ def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch) configuration.rust(True) assert configuration.rust_enabled() is True - assert configuration.rust_enabled(request_override=False) is False def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPatch) -> None: @@ -83,18 +70,8 @@ def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPat @pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) def test_invalid_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None: monkeypatch.setenv("LITELLM_RUST", value) - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") assert configuration.rust_enabled() is False - assert configuration.rust_ocr_enabled() is False - - -@pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) -def test_invalid_legacy_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None: - monkeypatch.setenv("LITELLM_USE_RUST_OCR", value) - - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert configuration.rust_enabled() is False def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytest.MonkeyPatch) -> None: @@ -104,40 +81,20 @@ def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytes assert executor.submit(configuration.rust_enabled).result() is True configuration.rust(False) assert executor.submit(configuration.rust_enabled).result() is False - assert executor.submit(configuration.rust_ocr_enabled).result() is False configuration.reset_rust_configuration() assert executor.submit(configuration.rust_enabled).result() is True - assert executor.submit(configuration.rust_ocr_enabled).result() is True def test_explicit_override_precedes_invalid_environment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_RUST", "sometimes") - assert configuration.rust_enabled(request_override=False) is False configuration.rust(True) assert configuration.rust_enabled() is True -def test_legacy_ocr_environment_is_deprecated_and_global(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert configuration.rust_enabled() is True - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert configuration.rust_ocr_enabled() is True - - -def test_global_environment_precedes_legacy_ocr_environment(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("LITELLM_RUST", "0") - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - - assert configuration.rust_enabled() is False - - -@pytest.mark.parametrize("environment_name", ("LITELLM_RUST", "LITELLM_USE_RUST_OCR")) @pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False"))) -def test_environment_controls_startup(environment_name: str, value: str, expected: str) -> None: - environment: Final = {**os.environ, environment_name: value} +def test_environment_controls_startup(value: str, expected: str) -> None: + environment: Final = {**os.environ, "LITELLM_RUST": value} result: Final = subprocess.run( ( sys.executable, diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index bbeb6c38f78..112464bda22 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -44,7 +44,7 @@ class AsyncBridge: def test_enabled_sync_bridge_receives_audio() -> None: bridge = SyncBridge() - rust_bridge.configure_rust_transcription(True, transcription=bridge) + rust_bridge.configure_rust_transcription(transcription=bridge) result = rust_bridge.transcription( model="mistral.voxtral-mini-3b-2507", audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, @@ -61,7 +61,7 @@ def test_enabled_sync_bridge_receives_audio() -> None: @pytest.mark.asyncio async def test_enabled_async_bridge() -> None: - rust_bridge.configure_rust_transcription(True, atranscription=AsyncBridge()) + rust_bridge.configure_rust_transcription(atranscription=AsyncBridge()) result = await rust_bridge.atranscription( model="mistral.voxtral-mini-3b-2507", audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d686b032ee0..271c84384b4 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -5239,52 +5239,6 @@ def test_client_side_timeout_marker_never_reaches_the_provider(): ) -def test_rust_flag_not_forwarded_as_provider_param(): - forwarded = get_non_default_completion_params({"rust": True, "temperature": 0.5}) - assert "rust" not in forwarded - - -def test_completion_does_not_leak_rust_flag_into_provider_request_body(): - mock_response = MagicMock() - mock_response.model_dump.return_value = { - "id": "chatcmpl-1", - "object": "chat.completion", - "created": 1234567890, - "model": "gpt-4o-mini", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "hi"}, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 1, - "completion_tokens": 1, - "total_tokens": 2, - }, - } - - mock_raw_response = MagicMock() - mock_raw_response.headers = {} - mock_raw_response.parse.return_value = mock_response - - mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = mock_raw_response - - litellm.completion( - model="openai/gpt-4o-mini", - messages=[{"role": "user", "content": "hi"}], - rust=True, - api_key="sk-test", - client=mock_client, - ) - - create_kwargs = mock_client.chat.completions.with_raw_response.create.call_args.kwargs - assert "rust" not in create_kwargs - assert "rust" not in (create_kwargs.get("extra_body") or {}) - - class _RecordingDeploymentFailureLogger(CustomLogger): def __init__(self) -> None: super().__init__() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b1534c19670..108e50b52a1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29562,8 +29562,6 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; - /** Rust */ - rust?: boolean | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Encryption Key Id */ @@ -39729,8 +39727,6 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; - /** Rust */ - rust?: boolean | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Encryption Key Id */ From 8da735410c7903b0e8c92033e444a76d32a03c6f Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 7 Sep 2026 10:44:22 -0700 Subject: [PATCH 45/49] test(ocr): trace callback lifecycle parity (#40063) --- .../crates/ai-gateway/src/ocr/hooks.rs | 12 +++ .../crates/ai-gateway/tests/ocr_lifecycle.rs | 52 +++++++++-- .../python-bridge/src/function_trace.rs | 28 ++++-- .../python-bridge/src/routes/definition.rs | 9 +- .../shared/tracing/native.py | 3 +- .../strategies/trace_parity/models.py | 1 + .../strategies/trace_parity/sdk/execution.py | 74 +++++++++++---- .../strategies/trace_parity/sdk/ocr/case.py | 91 +++++++++++++++++-- 8 files changed, 220 insertions(+), 50 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 446b323db3a..ed41f1ff9e7 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -265,6 +265,12 @@ impl CallLifecycleHooks for OcrLi Box::pin(async move { Ok(request) }) } + #[tracing::instrument( + name = "success_callback", + target = "litellm::function_trace", + level = "trace", + skip_all + )] fn async_log_success_event<'a>( &'a self, context: &'a CallLifecycleContext, @@ -288,6 +294,12 @@ impl CallLifecycleHooks for OcrLi }) } + #[tracing::instrument( + name = "failure_callback", + target = "litellm::function_trace", + level = "trace", + skip_all + )] fn async_log_failure_event<'a>( &'a self, context: &'a CallLifecycleContext, diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs index c3a89f4394d..60e90ed2a7c 100644 --- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs +++ b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs @@ -11,9 +11,13 @@ use litellm_ai_gateway::integrations::custom_logger::{ use litellm_ai_gateway::integrations::types::RequestMetadata; use litellm_ai_gateway::ocr::{OcrRequest, ocr}; use litellm_core::error::Error; +#[cfg(feature = "trace-parity")] +use litellm_core::observability::FunctionTrace; use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; +#[cfg(feature = "trace-parity")] +use tracing::instrument::WithSubscriber; async fn read_http_headers(socket: &mut TcpStream) -> String { let mut request = Vec::new(); @@ -320,14 +324,17 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { GuardrailEventHook::PreCall, GuardrailEventHook::DuringCall, ])); - let response = ocr(OcrRequest { + #[cfg(feature = "trace-parity")] + let trace = FunctionTrace::default(); + let api_base = format!("http://{addr}"); + let call = ocr(OcrRequest { model: "mistral-ocr-latest", document: json!({ "type": "document_url", "document_url": "https://example.com/doc.pdf" }), api_key: Some("sk-test"), - api_base: Some(&format!("http://{addr}")), + api_base: Some(&api_base), custom_llm_provider: Some("mistral"), extra_headers: None, optional_params: Map::new(), @@ -339,9 +346,10 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { ..Default::default() }, litellm_call_id: Some("ocr-call-1"), - }) - .await - .expect("ocr request succeeds"); + }); + #[cfg(feature = "trace-parity")] + let call = call.with_subscriber(trace.dispatcher()); + let response = call.await.expect("ocr request succeeds"); assert_eq!(response["pages"][0]["markdown"], "ok"); assert_eq!( @@ -359,6 +367,16 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { error_kind: None, }] ); + #[cfg(feature = "trace-parity")] + assert_eq!( + trace + .events() + .iter() + .filter(|event| event.function.ends_with("_callback")) + .map(|event| event.function) + .collect::>(), + vec!["success_callback"] + ); let request = server.await.expect("server task completes"); assert!(request.contains(r#""guarded_pre":true"#), "{request}"); @@ -388,14 +406,17 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { }); let logger = Arc::new(RecordingOcrLogger::default()); - let err = ocr(OcrRequest { + #[cfg(feature = "trace-parity")] + let trace = FunctionTrace::default(); + let api_base = format!("http://{addr}"); + let call = ocr(OcrRequest { model: "mistral-ocr-latest", document: json!({ "type": "document_url", "document_url": "https://example.com/doc.pdf" }), api_key: Some("sk-test"), - api_base: Some(&format!("http://{addr}")), + api_base: Some(&api_base), custom_llm_provider: Some("mistral"), extra_headers: None, optional_params: Map::new(), @@ -404,9 +425,10 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { guardrails: Vec::new(), request_metadata: RequestMetadata::default(), litellm_call_id: Some("ocr-call-2"), - }) - .await - .expect_err("provider error propagates"); + }); + #[cfg(feature = "trace-parity")] + let call = call.with_subscriber(trace.dispatcher()); + let err = call.await.expect_err("provider error propagates"); assert!(matches!(err, Error::Http { status: 500, .. })); server.await.expect("server task completes"); @@ -421,6 +443,16 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { error_kind: Some("HttpError".to_string()), }] ); + #[cfg(feature = "trace-parity")] + assert_eq!( + trace + .events() + .iter() + .filter(|event| event.function.ends_with("_callback")) + .map(|event| event.function) + .collect::>(), + vec!["failure_callback"] + ); } #[tokio::test] diff --git a/litellm-rust/crates/python-bridge/src/function_trace.rs b/litellm-rust/crates/python-bridge/src/function_trace.rs index ea7d9f4993e..bc3c962f7a3 100644 --- a/litellm-rust/crates/python-bridge/src/function_trace.rs +++ b/litellm-rust/crates/python-bridge/src/function_trace.rs @@ -1,3 +1,4 @@ +use std::fmt::Display; use std::future::Future; use litellm_core::observability::{FunctionTrace, FunctionTraceEvent}; @@ -6,17 +7,32 @@ use tracing::instrument::WithSubscriber; #[derive(Serialize)] pub(crate) struct TracedResponse { - response: T, + #[serde(skip_serializing_if = "Option::is_none")] + response: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, trace: Vec, } pub(crate) async fn capture( future: impl Future>, -) -> Result, E> { +) -> Result, E> +where + E: Display, +{ let trace = FunctionTrace::default(); - let response = future.with_subscriber(trace.dispatcher()).await?; - Ok(TracedResponse { - response, - trace: trace.events(), + let result = future.with_subscriber(trace.dispatcher()).await; + let events = trace.events(); + Ok(match result { + Ok(response) => TracedResponse { + response: Some(response), + error: None, + trace: events, + }, + Err(error) => TracedResponse { + response: None, + error: Some(error.to_string()), + trace: events, + }, }) } diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 3285da14d5f..bc51647cbad 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -486,10 +486,11 @@ asyncio.run(exercise()) let code = CString::new( r#" result = routes.echo("traced") -assert result == { - "response": "traced", - "trace": [{"function": "execute_echo", "depth": 0}], -} +assert result["response"] == "traced", result +assert [event["function"] for event in result["trace"]] == ["execute_echo"], result +failure = routes.echo("error") +assert failure["error"] == "invalid request: synthetic error", failure +assert [event["function"] for event in failure["trace"]] == ["execute_echo"], failure "#, ) .expect("Python source should not contain null bytes"); diff --git a/tests/rust-python-harness/shared/tracing/native.py b/tests/rust-python-harness/shared/tracing/native.py index 4f988f65294..688995cbc4b 100644 --- a/tests/rust-python-harness/shared/tracing/native.py +++ b/tests/rust-python-harness/shared/tracing/native.py @@ -19,7 +19,8 @@ class _TraceEventPayload(BaseModel): class TraceResponsePayload(BaseModel): model_config = ConfigDict(strict=True, extra="forbid") - response: object + response: object = None + error: str | None = None trace: tuple[_TraceEventPayload, ...] | list[_TraceEventPayload] diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index da560f99730..04659b25382 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -16,6 +16,7 @@ TraceFailureSource = Literal["python", "rust", "harness"] class RouteFixture: kwargs: dict[str, object] provider_responses: tuple[RecordedHttpResponse, ...] + expected_failure: bool = False @dataclass(frozen=True, slots=True) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index f8d7c55d4e2..eb6c9233565 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -2,15 +2,16 @@ from __future__ import annotations import asyncio from collections.abc import Awaitable +from dataclasses import dataclass from pathlib import Path from typing import Final, Protocol, cast from ....shared.parity.replay import replay_server from ....shared.reporting.models import Surface -from ....shared.tracing.native import native_trace_events +from ....shared.tracing.native import TraceResponsePayload, native_trace_events from ....shared.tracing.profiler import FunctionTraceEvent, profile_python from ....shared.tracing.steps import Engine, pipeline_projection -from ..models import RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario +from ..models import RouteFixture, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario from ..reporting import TraceComparisonArtifact @@ -18,9 +19,22 @@ class SdkCall(Protocol): def __call__(self, **kwargs: object) -> object: ... +@dataclass(frozen=True, slots=True) +class _CollectedTrace: + events: tuple[FunctionTraceEvent, ...] + error: str | None = None + + def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> object: async def invoke_async() -> object: - return await cast(Awaitable[object], function(**kwargs)) + try: + return await cast(Awaitable[object], function(**kwargs)) + finally: + await asyncio.sleep(0) + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + await GLOBAL_LOGGING_WORKER.stop() if asynchronous: return asyncio.run(invoke_async()) @@ -48,16 +62,30 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa return cast(SdkCall, getattr(owner, spec.python_entrypoints[int(asynchronous)])) +def _python_invocation_error(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> str | None: + try: + _invoke(function, kwargs, asynchronous=asynchronous) + except Exception as error: + return f"{type(error).__name__}: {error}" + return None + + def _collect( - function: SdkCall, kwargs: dict[str, object], engine: Engine, *, asynchronous: bool -) -> tuple[FunctionTraceEvent, ...]: + function: SdkCall, + fixture: RouteFixture, + engine: Engine, + *, + asynchronous: bool, +) -> _CollectedTrace: + kwargs: Final = fixture.kwargs if engine == "rust": - return native_trace_events(_invoke(function, kwargs, asynchronous=asynchronous)) + payload: Final = TraceResponsePayload.model_validate(_invoke(function, kwargs, asynchronous=asynchronous)) + return _CollectedTrace(native_trace_events(payload), payload.error) import litellm with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: - _invoke(function, kwargs, asynchronous=asynchronous) - return tuple(profiler.events) + error: Final = _python_invocation_error(function, kwargs, asynchronous=asynchronous) + return _CollectedTrace(tuple(profiler.events), error) def collect_trace( @@ -68,22 +96,30 @@ def collect_trace( return function try: with replay_server() as provider: - fixture: Final = spec.fixture(engine, provider.url) - for response in fixture.provider_responses: + base_fixture: Final = spec.fixture(engine, provider.url) + for response in base_fixture.provider_responses: provider.enqueue_response(response) - kwargs: Final = { - **fixture.kwargs, - "api_key": "test-key", - "api_base": provider.url, - **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}), - } - events: Final = _collect(function, kwargs, engine, asynchronous=asynchronous) + fixture: Final = RouteFixture( + kwargs={ + **base_fixture.kwargs, + "api_key": "test-key", + "api_base": provider.url, + **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}), + }, + provider_responses=base_fixture.provider_responses, + expected_failure=base_fixture.expected_failure, + ) + collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) except Exception as error: return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}") - if not events: + if fixture.expected_failure and collected.error is None: + return TraceExecutionFailure(engine, "call succeeded but the scenario expects failure") + if not fixture.expected_failure and collected.error is not None: + return TraceExecutionFailure(engine, collected.error) + if not collected.events: return TraceExecutionFailure(engine, "trace is empty") - return events + return collected.events def _failure_message(result: tuple[FunctionTraceEvent, ...] | TraceExecutionFailure) -> str | None: diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index fe214f45339..2a4a1b3a152 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -19,6 +19,20 @@ COMMON_MAPPINGS: Final = ( mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), ) +SUCCESS_CALLBACK_SYNC_MAPPING: Final = mapping( + rust_span="success_callback", + python_frame=r"BoundedLoggingThreadPoolExecutor\.submit$", +) +SUCCESS_CALLBACK_ASYNC_MAPPING: Final = mapping( + rust_span="success_callback", + python_frame=r"Logging\.async_success_handler$", +) +FAILURE_CALLBACK_MAPPING: Final = mapping( + rust_span="failure_callback", + python_frame=r"Logging\.(?:async_)?failure_handler$", +) +IGNORED_SUCCESS_CALLBACK_MAPPING: Final = mapping(rust_span="success_callback") + SYNC_MAPPINGS: Final = ( *COMMON_MAPPINGS, mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"), @@ -39,6 +53,21 @@ ASYNC_MAPPINGS: Final = ( ), ) +CALLBACK_SUCCESS_SYNC_MAPPINGS: Final = (*SYNC_MAPPINGS, SUCCESS_CALLBACK_SYNC_MAPPING) +CALLBACK_SUCCESS_ASYNC_MAPPINGS: Final = (*ASYNC_MAPPINGS, SUCCESS_CALLBACK_ASYNC_MAPPING) +CALLBACK_FAILURE_SYNC_MAPPINGS: Final = ( + *COMMON_MAPPINGS, + mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"), + FAILURE_CALLBACK_MAPPING, +) +CALLBACK_FAILURE_ASYNC_MAPPINGS: Final = ( + *COMMON_MAPPINGS, + mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"), + mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"), + FAILURE_CALLBACK_MAPPING, +) + + AZURE_COMMON_MAPPINGS: Final = ( *COMMON_MAPPINGS[:7], mapping( @@ -99,6 +128,34 @@ def _mistral_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture(engine, "mistral/mistral-ocr-latest") +def _callback_fixture(engine: Engine, *, failure: bool) -> RouteFixture: + fixture: Final = _fixture(engine, "mistral/mistral-ocr-latest") + provider_responses: Final = ( + ( + RecordedHttpResponse.from_bytes( + 400, + (HttpHeader(name="content-type", value="application/json"),), + b'{"message":"trace callback provider failure"}', + ), + ) + if failure + else fixture.provider_responses + ) + return RouteFixture( + kwargs=fixture.kwargs, + provider_responses=provider_responses, + expected_failure=failure, + ) + + +def _mistral_callback_success_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _callback_fixture(engine, failure=False) + + +def _mistral_callback_failure_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _callback_fixture(engine, failure=True) + + def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture( engine, @@ -304,36 +361,50 @@ TRACE_SUITE: Final = TraceSuite( name="mistral", fixture=_mistral_fixture, mappings=COMMON_MAPPINGS, - sync_mappings=SYNC_MAPPINGS, - async_mappings=ASYNC_MAPPINGS, + sync_mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + ), + TraceScenario( + name="mistral-callback-success", + fixture=_mistral_callback_success_fixture, + mappings=COMMON_MAPPINGS, + sync_mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS, + async_mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS, + ), + TraceScenario( + name="mistral-callback-failure", + fixture=_mistral_callback_failure_fixture, + mappings=(*COMMON_MAPPINGS, FAILURE_CALLBACK_MAPPING), + sync_mappings=CALLBACK_FAILURE_SYNC_MAPPINGS, + async_mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS, ), TraceScenario( name="azure-ai", fixture=_azure_fixture, mappings=AZURE_COMMON_MAPPINGS, - sync_mappings=AZURE_SYNC_MAPPINGS, - async_mappings=AZURE_ASYNC_MAPPINGS, + sync_mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), TraceScenario( name="azure-document-intelligence", fixture=_azure_document_intelligence_fixture, mappings=DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS, - sync_mappings=DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, - async_mappings=DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, + sync_mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), TraceScenario( name="vertex-ai", fixture=_vertex_fixture, mappings=VERTEX_COMMON_MAPPINGS, - sync_mappings=VERTEX_SYNC_MAPPINGS, - async_mappings=VERTEX_ASYNC_MAPPINGS, + sync_mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), TraceScenario( name="vertex-deepseek", fixture=_vertex_deepseek_fixture, mappings=DEEPSEEK_COMMON_MAPPINGS, - sync_mappings=DEEPSEEK_SYNC_MAPPINGS, - async_mappings=DEEPSEEK_ASYNC_MAPPINGS, + sync_mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), ), ) From 2700ffe8c978229628767a45e2c83b7c1bb0297a Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 7 Sep 2026 11:12:34 -0700 Subject: [PATCH 46/49] fix: capture provider request id in failure logging payloads (#40045) * fix: capture provider request id in failure logging payloads * fix: include common provider request id headers * test: remove accidental formatter churn * test: type provider request id cases * test: cover azure and google request id headers * refactor: share provider request id headers from constants * chore: remove redundant request header comment * test: cover provider header lookup failures --- litellm/constants.py | 10 ++ litellm/litellm_core_utils/litellm_logging.py | 27 +++++ litellm/types/utils.py | 1 + .../test_litellm_logging.py | 111 +++++++++++++++++- 4 files changed, 148 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index ce744e9c58a..d53686e5e5b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1901,6 +1901,16 @@ HTTP_FRAMING_HEADERS: Final[frozenset[str]] = frozenset( } ) +PROVIDER_REQUEST_ID_HEADERS: Final[tuple[str, ...]] = ( + "x-amzn-requestid", + "x-request-id", + "request-id", + "x-ms-request-id", + "apim-request-id", + "x-goog-request-id", + "cf-ray", +) + # Browser-facing security headers that a malicious or misconfigured upstream # provider must not be able to set on the proxy's own response. BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 60abad17c03..ca2cca5360f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -42,6 +42,7 @@ from litellm.caching.caching_handler import LLMCachingHandler from litellm.constants import ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, + PROVIDER_REQUEST_ID_HEADERS, SENTRY_DENYLIST, SENTRY_PII_DENYLIST, ) @@ -255,6 +256,30 @@ _in_memory_loggers: Final[list[CustomLogger]] = [] _STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggingMetadata.__annotations__.keys()) + +def _get_provider_request_id(original_exception: Exception) -> str | None: + try: + error_response: Final = getattr(original_exception, "response", None) + header_sources: Final = ( + _get_response_headers(original_exception), + getattr(error_response, "headers", None), + getattr(original_exception, "litellm_response_headers", None), + ) + return next( + ( + str(value) + for expected_header_name in PROVIDER_REQUEST_ID_HEADERS + for headers in header_sources + if isinstance(headers, Mapping) + for header_name, value in headers.items() + if isinstance(header_name, str) and header_name.lower() == expected_header_name and value + ), + None, + ) + except Exception: + return None + + ### GLOBAL VARIABLES ### # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys @@ -5664,6 +5689,7 @@ class StandardLoggingPayloadSetup: rate_limit_category: Final = validate_rate_limit_category(getattr(original_exception, "category", None)) rate_limit_type: Final = validate_rate_limit_type(getattr(original_exception, "rate_limit_type", None)) budget_error: Final = original_exception if isinstance(original_exception, BudgetExceededError) else None + provider_request_id: Final = _get_provider_request_id(original_exception) if original_exception else None return StandardLoggingPayloadErrorInformation( error_code=error_status, @@ -5671,6 +5697,7 @@ class StandardLoggingPayloadSetup: llm_provider=_llm_provider_in_exception, traceback=_redact_string(traceback_info), error_message=_redact_string(error_message), + error_provider_request_id=provider_request_id, error_rate_limit_category=rate_limit_category, error_rate_limit_type=rate_limit_type, error_budget_entity_type=budget_error.entity_type if budget_error else None, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3a1422bb190..010ff18d166 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3064,6 +3064,7 @@ class StandardLoggingPayloadErrorInformation(TypedDict, total=False): llm_provider: str | None traceback: str | None error_message: str | None + error_provider_request_id: ReadOnly[str | None] # error_rate_limit_category: # For 429 / rate-limit errors, the source of the rate limit. One of the # string values defined by `litellm.exceptions.RateLimitErrorCategory` diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 3748b5a9d64..0fdca755685 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1,7 +1,9 @@ +import asyncio import contextlib +import datetime import os import sys -import asyncio +from typing import Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -6439,6 +6441,113 @@ def test_get_standard_logging_object_payload_survives_logging_obj_without_timing assert payload["hidden_params"]["litellm_overhead_time_ms"] is None +@pytest.mark.parametrize( + ("header_name", "header_source"), + ( + ("x-amzn-RequestId", "response"), + ("x-request-id", "response"), + ("request-id", "response"), + ("x-ms-request-id", "response"), + ("apim-request-id", "response"), + ("x-goog-request-id", "response"), + ("cf-ray", "response"), + ("X-Request-Id", "litellm_response_headers"), + ("X-MS-Request-ID", "headers"), + ), +) +def test_failure_standard_logging_payload_captures_provider_request_id( + logging_obj: LitellmLogging, + header_name: str, + header_source: Literal["response", "litellm_response_headers", "headers"], +): + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + request_id = "provider-request-123" + request = httpx.Request("POST", "https://provider.example/v1/chat/completions") + response = httpx.Response(429, headers={header_name: request_id}, request=request) + provider_error = httpx.HTTPStatusError("provider error", request=request, response=response) + if header_source == "litellm_response_headers": + response.headers.clear() + provider_error.litellm_response_headers = {header_name: request_id} + elif header_source == "headers": + response.headers.clear() + provider_error.headers = {header_name: request_id} + now = datetime.datetime.now() + + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "call-1", "model": "test-model", "messages": []}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="failure", + original_exception=provider_error, + ) + + assert payload is not None + assert payload["error_information"] is not None + assert payload["error_information"]["error_provider_request_id"] == request_id + + +def test_get_error_information_ignores_unsupported_headers() -> None: + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + request = httpx.Request("POST", "https://provider.example/v1/chat/completions") + response = httpx.Response(429, headers={"retry-after": "3"}, request=request) + provider_error = httpx.HTTPStatusError("provider error", request=request, response=response) + + error_information = StandardLoggingPayloadSetup.get_error_information(provider_error) + + assert error_information["error_provider_request_id"] is None + + +def test_get_error_information_uses_header_precedence_and_fallback() -> None: + from litellm.exceptions import RateLimitError + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + request = httpx.Request("POST", "https://provider.example/v1/chat/completions") + response = httpx.Response( + 429, + headers={"x-request-id": "response-id", "x-amzn-requestid": "amazon-id"}, + request=request, + ) + provider_error = RateLimitError( + message="provider error", + llm_provider="test-provider", + model="test-model", + response=response, + headers={"retry-after": "3"}, + ) + + error_information = StandardLoggingPayloadSetup.get_error_information(provider_error) + + assert error_information["error_provider_request_id"] == "amazon-id" + + +def test_get_error_information_ignores_malformed_headers() -> None: + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + provider_error = Exception("provider error") + provider_error.headers = [("x-request-id", "provider-request-123")] + + error_information = StandardLoggingPayloadSetup.get_error_information(provider_error) + + assert error_information["error_provider_request_id"] is None + + +def test_get_provider_request_id_ignores_header_lookup_errors() -> None: + from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id + + class HeaderLookupError(Exception): + @property + def response(self) -> object: + raise RuntimeError("headers unavailable") + + assert _get_provider_request_id(HeaderLookupError("provider error")) is None + + def test_get_standard_logging_object_payload_failure_status_keeps_overhead_none(logging_obj): """A post_call guardrail can fail the request after the upstream call succeeded; the failure payload keeps litellm_overhead_time_ms None, matching responses that carry their own _hidden_params.""" From 9acc01efce25bc83769a5f45f0c1de4991edcb3e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 7 Sep 2026 11:18:16 -0700 Subject: [PATCH 47/49] test(e2e/ui): cover member role and budget edits, member permission delegation, and team guardrail removal (#40042) * test(e2e/ui): cover member role and budget edits, member permission delegation, and team guardrail removal Three Playwright specs for the Teams flows enterprise customers hit most, each owning its fixtures and proving the mutation through a read-back rather than a toast. - teamMemberEdit: an admin edits a member's team role and per-member budget, and both survive a reload of the Members table - memberPermissions: a plain member is refused /key/generate for their team, a team admin grants it on the Member Permissions tab, and the member then creates a team key that serves a real completion - teamGuardrailRemoval: clearing a team's only guardrail on the Settings tab really clears it, and traffic the guardrail refused starts serving again * test(e2e/ui): make the new team specs safe to run in parallel Fixture ids came from Date.now(), so two repeats starting in the same millisecond minted the same user id: one got a 409 and the loser's teardown deleted the user the other was still signed in as. Ids now carry a random suffix. Also move the member-permissions setup inside the cleanup-protected block so a half-finished setup cannot leak a team, and close both browser contexts the test opens. --- .../proxy-admin/teamGuardrailRemoval.spec.ts | 145 ++++++++++++++ .../tests/proxy-admin/teamMemberEdit.spec.ts | 129 +++++++++++++ .../team-admin/memberPermissions.spec.ts | 177 ++++++++++++++++++ 3 files changed, 451 insertions(+) create mode 100644 tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts create mode 100644 tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts create mode 100644 tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts diff --git a/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts b/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts new file mode 100644 index 00000000000..f7930378b1e --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts @@ -0,0 +1,145 @@ +import { test, expect, type APIRequestContext, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { proxyIsPremium } from "../../helpers/premium"; +import { readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; + +const auth = () => ({ Authorization: `Bearer ${masterKey()}` }); + +async function guardrailId(request: APIRequestContext, name: string): Promise { + const res = await request.get("/v2/guardrails/list", { headers: auth() }); + expect(res.ok(), `GET /v2/guardrails/list (${res.status()})`).toBe(true); + const rows = (await res.json()).guardrails as { guardrail_id: string; guardrail_name: string | null }[]; + return rows.find((row) => row.guardrail_name === name)?.guardrail_id; +} + +async function teamGuardrails(page: PlaywrightPage, teamId: string): Promise { + const body = await readBack<{ team_info: { metadata: { guardrails?: string[] } | null } }>( + page, + `/team/info?team_id=${encodeURIComponent(teamId)}`, + ); + return body.team_info.metadata?.guardrails ?? []; +} + +async function keywordPromptStatus(request: APIRequestContext, apiKey: string, keyword: string): Promise { + const res = await request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" }, + data: { model: CHAT_MODEL_A, messages: [{ role: "user", content: `please tell me about ${keyword}` }] }, + }); + return res.status(); +} + +test.describe("Proxy Admin - Team guardrail removal", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Clearing a team's only guardrail on the Settings tab lets blocked traffic through again", async ({ + page, + request, + }) => { + test.skip(!proxyIsPremium(), "proxy under test is unlicensed, so team guardrails are premium-gated"); + + const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`; + const guardrailName = `e2e-team-guardrail-${stamp}`; + const bannedKeyword = `e2eteamban${stamp}`; + const teamAlias = `e2e-guardrail-team-${stamp}`; + + let teamId = ""; + let teamKey = ""; + try { + const guardrailRes = await request.post("/guardrails", { + headers: auth(), + data: { + guardrail: { + guardrail_name: guardrailName, + litellm_params: { + guardrail: "litellm_content_filter", + mode: "pre_call", + default_on: false, + blocked_words: [{ keyword: bannedKeyword, action: "BLOCK" }], + }, + }, + }, + }); + expect( + guardrailRes.ok(), + `POST /guardrails failed (${guardrailRes.status()}): ${await guardrailRes.text()}`, + ).toBe(true); + + const teamRes = await request.post("/team/new", { + headers: auth(), + data: { team_alias: teamAlias, models: [CHAT_MODEL_A], metadata: { guardrails: [guardrailName] } }, + }); + expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + teamId = (await teamRes.json()).team_id as string; + + const keyRes = await request.post("/key/generate", { headers: auth(), data: { team_id: teamId } }); + expect(keyRes.ok(), `POST /key/generate failed (${keyRes.status()}): ${await keyRes.text()}`).toBe(true); + teamKey = (await keyRes.json()).key as string; + + await expect + .poll(async () => keywordPromptStatus(request, teamKey, bannedKeyword), { + message: "the team's guardrail never started refusing the banned keyword", + timeout: 60_000, + }) + .toBe(400); + + await navigateToPage(page, Page.Teams); + await dismissFeedbackPopup(page); + await clickTeamId(page, teamId); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + const chip = page.locator('[data-slot="combobox-chip"]').filter({ hasText: guardrailName }); + await expect(chip).toBeVisible({ timeout: 10_000 }); + await chip.locator('[data-slot="combobox-chip-remove"]').click(); + await expect(chip).toHaveCount(0, { timeout: 10_000 }); + + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect + .poll(async () => teamGuardrails(page, teamId), { + message: "the team still carries a guardrail in /team/info after the save", + timeout: 20_000, + }) + .toEqual([]); + + await page.reload(); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + await expect(page.getByRole("combobox", { name: "Select guardrails" })).toBeVisible({ timeout: 15_000 }); + await expect( + page.locator('[data-slot="combobox-chip"]').filter({ hasText: guardrailName }), + "the removed guardrail is gone from the Settings tab after a reload", + ).toHaveCount(0); + + await expect + .poll(async () => keywordPromptStatus(request, teamKey, bannedKeyword), { + message: "the team key is still refused for a keyword whose guardrail was removed", + timeout: 60_000, + }) + .toBe(200); + + const served = await request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${teamKey}`, "Content-Type": "application/json" }, + data: { + model: CHAT_MODEL_A, + messages: [{ role: "user", content: `please tell me about ${bannedKeyword}` }], + }, + }); + expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); + } finally { + if (teamKey) { + await request.post("/key/delete", { headers: auth(), data: { keys: [teamKey] } }); + } + if (teamId) { + await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } }); + } + const id = await guardrailId(request, guardrailName); + if (id) { + await request.delete(`/guardrails/${id}`, { headers: auth() }); + } + } + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts b/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts new file mode 100644 index 00000000000..c7325e78f75 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts @@ -0,0 +1,129 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, masterKey } from "../../helpers/traffic"; + +interface TeamInfoResponse { + team_info: { + models: string[]; + members_with_roles: { user_id?: string; role?: string }[]; + }; + team_memberships: { + user_id: string; + litellm_budget_table: { max_budget: number | null } | null; + }[]; +} + +const auth = () => ({ Authorization: `Bearer ${masterKey()}` }); + +async function teamInfo(page: PlaywrightPage, teamId: string): Promise { + return readBack(page, `/team/info?team_id=${encodeURIComponent(teamId)}`); +} + +function roleOf(info: TeamInfoResponse, userId: string): string | undefined { + return info.team_info.members_with_roles.find((member) => member.user_id === userId)?.role; +} + +function budgetOf(info: TeamInfoResponse, userId: string): number | null | undefined { + return info.team_memberships.find((membership) => membership.user_id === userId)?.litellm_budget_table?.max_budget; +} + +function otherMembers(info: TeamInfoResponse, userId: string): string[] { + return info.team_info.members_with_roles + .filter((member) => member.user_id !== userId) + .map((member) => `${member.user_id}:${member.role}`) + .sort(); +} + +test.describe("Proxy Admin - Team member edit", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + const createdTeams: string[] = []; + const createdUsers: string[] = []; + + test.afterEach(async ({ request }) => { + for (const teamId of createdTeams.splice(0)) { + await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } }); + } + for (const userId of createdUsers.splice(0)) { + await request.post("/user/delete", { headers: auth(), data: { user_ids: [userId] } }); + } + }); + + test("Editing a member's role and per-member budget persists and survives a reload", async ({ page, request }) => { + const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`; + const memberId = `e2e-member-edit-${stamp}`; + const teamAlias = `e2e-member-edit-team-${stamp}`; + + createdUsers.push(memberId); + const created = await request.post("/user/new", { + headers: auth(), + data: { user_id: memberId, user_role: "internal_user", auto_create_key: false }, + }); + expect(created.ok(), `POST /user/new failed (${created.status()}): ${await created.text()}`).toBe(true); + + const teamRes = await request.post("/team/new", { + headers: auth(), + data: { + team_alias: teamAlias, + models: [CHAT_MODEL_A], + members_with_roles: [{ user_id: memberId, role: "admin" }], + }, + }); + expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + const teamId = (await teamRes.json()).team_id as string; + createdTeams.push(teamId); + + const before = await teamInfo(page, teamId); + expect(roleOf(before, memberId), "the member starts out as a team admin").toBe("admin"); + expect(budgetOf(before, memberId) ?? null, "the member starts out with no per-member budget").toBeNull(); + expect( + otherMembers(before, memberId).length, + "the team has another member for the edit to leave alone", + ).toBeGreaterThan(0); + + await navigateToPage(page, Page.Teams); + await dismissFeedbackPopup(page); + await clickTeamId(page, teamId); + await page.getByRole("tab", { name: "Members" }).click(); + + const memberRow = page.locator("tr", { hasText: memberId }).first(); + await expect(memberRow).toBeVisible({ timeout: 10_000 }); + await memberRow.getByTestId("edit-member").click(); + + const modal = page.getByRole("dialog", { name: "Edit Member" }); + await expect(modal).toBeVisible({ timeout: 10_000 }); + + await modal.getByLabel(/^Role/).click(); + await page.getByRole("option", { name: "User", exact: true }).click(); + await modal.getByLabel(/Team Member Budget \(USD\)/).fill("5"); + await modal.getByRole("button", { name: "Save Changes" }).click(); + + await expect(page.getByText("Team member updated successfully").first()).toBeVisible({ timeout: 10_000 }); + + await expect + .poll( + async () => { + const info = await teamInfo(page, teamId); + return [roleOf(info, memberId), budgetOf(info, memberId)]; + }, + { message: "the member's role and budget never landed in /team/info", timeout: 20_000 }, + ) + .toEqual(["user", 5]); + + await page.reload(); + await page.getByRole("tab", { name: "Members" }).click(); + const reloadedRow = page.locator("tr", { hasText: memberId }).first(); + await expect(reloadedRow).toBeVisible({ timeout: 15_000 }); + await expect(reloadedRow.getByText("user", { exact: true }), "role shown after a reload").toBeVisible(); + await expect(reloadedRow.getByText("$5.00"), "per-member budget shown after a reload").toBeVisible(); + + const after = await teamInfo(page, teamId); + expect(after.team_info.models, "model access untouched by a member edit").toEqual(before.team_info.models); + expect(otherMembers(after, memberId), "the rest of the roster untouched by a member edit").toEqual( + otherMembers(before, memberId), + ); + }); +}); diff --git a/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts b/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts new file mode 100644 index 00000000000..75fb3be9b64 --- /dev/null +++ b/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts @@ -0,0 +1,177 @@ +import { test, expect, type Browser, type BrowserContext, type Page as PlaywrightPage } from "@playwright/test"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; + +const PASSWORD = "E2e-Member-Perms-Pass-1!"; + +const auth = () => ({ Authorization: `Bearer ${masterKey()}` }); + +async function sessionKey(page: PlaywrightPage): Promise { + const cookie = (await page.context().cookies()).find((candidate) => candidate.name === "token"); + expect(cookie?.value, "logged-in session carries a token cookie").toBeTruthy(); + const payload = JSON.parse(Buffer.from(cookie!.value.split(".")[1], "base64url").toString("utf-8")) as { + key?: string; + }; + expect(payload.key, "session JWT carries the virtual key the dashboard calls with").toMatch(/^sk-/); + return payload.key!; +} + +async function signIn(browser: Browser, email: string): Promise { + const context = await browser.newContext({ storageState: { cookies: [], origins: [] } }); + const page = await context.newPage(); + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill(email); + await page.getByPlaceholder("Enter your password").fill(PASSWORD); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 }); + await dismissFeedbackPopup(page); + return context; +} + +test.describe("Team Admin - Member permissions", () => { + test.use({ storageState: { cookies: [], origins: [] } }); + + test("Granting /key/generate lets a plain member create a team key that serves traffic", async ({ + browser, + request, + }) => { + const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`; + const adminId = `e2e-perm-admin-${stamp}`; + const memberId = `e2e-perm-member-${stamp}`; + const adminEmail = `${adminId}@test.local`; + const memberEmail = `${memberId}@test.local`; + const teamAlias = `e2e-perm-team-${stamp}`; + + const createUser = async (userId: string, email: string): Promise => { + const created = await request.post("/user/new", { + headers: auth(), + data: { user_id: userId, user_email: email, user_role: "internal_user", auto_create_key: false }, + }); + expect(created.ok(), `POST /user/new for ${userId} (${created.status()}): ${await created.text()}`).toBe(true); + const password = await request.post("/user/update", { + headers: auth(), + data: { user_id: userId, password: PASSWORD }, + }); + expect(password.ok(), `POST /user/update for ${userId} (${password.status()})`).toBe(true); + }; + + let teamId = ""; + const createdKeys: string[] = []; + const contexts: BrowserContext[] = []; + try { + await createUser(adminId, adminEmail); + await createUser(memberId, memberEmail); + + const teamRes = await request.post("/team/new", { + headers: auth(), + data: { + team_alias: teamAlias, + models: [CHAT_MODEL_A], + members_with_roles: [ + { user_id: adminId, role: "admin" }, + { user_id: memberId, role: "user" }, + ], + }, + }); + expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + teamId = (await teamRes.json()).team_id as string; + + const memberContext = await signIn(browser, memberEmail); + contexts.push(memberContext); + const memberPage = memberContext.pages()[0]; + const memberSessionKey = await sessionKey(memberPage); + + const refused = await memberPage.request.post("/key/generate", { + headers: { Authorization: `Bearer ${memberSessionKey}`, "Content-Type": "application/json" }, + data: { team_id: teamId, key_alias: `e2e-perm-denied-${stamp}` }, + }); + expect(refused.status(), "a plain member cannot mint a team key before the grant").toBe(401); + expect(await refused.text()).toContain("/key/generate"); + + const adminContext = await signIn(browser, adminEmail); + contexts.push(adminContext); + const adminPage = adminContext.pages()[0]; + await navigateToPage(adminPage, Page.Teams); + await clickTeamId(adminPage, teamId); + await adminPage.getByRole("tab", { name: "Member Permissions" }).click(); + + for (const route of ["/key/generate", "/key/update"]) { + await adminPage.getByRole("row").filter({ hasText: route }).getByRole("checkbox").check(); + } + await adminPage.getByRole("button", { name: "Save Changes" }).click(); + await expect(adminPage.getByText("Permissions updated successfully").first()).toBeVisible({ timeout: 10_000 }); + + await expect + .poll( + async () => { + const res = await request.get(`/team/permissions_list?team_id=${encodeURIComponent(teamId)}`, { + headers: auth(), + }); + if (!res.ok()) return []; + return ((await res.json()).team_member_permissions ?? []) as string[]; + }, + { message: "the granted permissions never landed in /team/permissions_list", timeout: 20_000 }, + ) + .toEqual(expect.arrayContaining(["/key/generate", "/key/update"])); + + const keyAlias = `e2e-perm-key-${stamp}`; + await navigateToPage(memberPage, Page.ApiKeys); + await memberPage.getByRole("button", { name: /Create New Key/i }).click(); + await expect(memberPage.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); + await memberPage.getByLabel(/Key Name/).fill(keyAlias); + + const teamSelect = memberPage.getByTestId("team-dropdown").getByRole("combobox"); + await teamSelect.click(); + await memberPage.keyboard.type(teamAlias); + await memberPage.getByRole("option", { name: teamAlias }).first().click(); + + await memberPage.getByRole("combobox", { name: "Select models" }).click(); + await memberPage.getByRole("option", { name: "All Team Models", exact: true }).click(); + await memberPage.keyboard.press("Escape"); + + await memberPage.getByRole("button", { name: "Create Key", exact: true }).click(); + const saveDialog = memberPage.getByRole("dialog", { name: "Save your Key" }); + await expect(saveDialog).toBeVisible({ timeout: 15_000 }); + const apiKey = (await saveDialog.locator("pre").innerText()).trim(); + expect(apiKey).toMatch(/^sk-/); + createdKeys.push(apiKey); + await memberPage.keyboard.press("Escape"); + + await expect + .poll( + async () => { + const res = await request.get( + `/key/list?team_id=${encodeURIComponent(teamId)}&return_full_object=true&size=100`, + { headers: auth() }, + ); + if (!res.ok()) return null; + const row = ((await res.json()).keys as Record[]).find( + (candidate) => candidate.key_alias === keyAlias, + ); + return row ? [row.user_id, row.team_id] : null; + }, + { message: `key ${keyAlias} never appeared on the team with the member as its owner`, timeout: 20_000 }, + ) + .toEqual([memberId, teamId]); + + const served = await request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" }, + data: { model: CHAT_MODEL_A, messages: [{ role: "user", content: `member key ping ${stamp}` }] }, + }); + expect(served.status(), "the delegated key is a real key the gateway serves").toBe(200); + expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); + } finally { + for (const context of contexts) { + await context.close(); + } + for (const key of createdKeys) { + await request.post("/key/delete", { headers: auth(), data: { keys: [key] } }); + } + if (teamId) { + await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } }); + } + await request.post("/user/delete", { headers: auth(), data: { user_ids: [adminId, memberId] } }); + } + }); +}); From 5f2b4d27d712847c0f4fef5333f24848983686f5 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 7 Sep 2026 11:23:12 -0700 Subject: [PATCH 48/49] test(ocr): add SDK callback E2E parity (#40061) * test(ocr): cover SDK callback parity * test(ocr): compare callback kwargs --- litellm/ocr/main.py | 74 +++-- .../shared/parity/compare.py | 1 + .../shared/parity/models.py | 21 ++ .../shared/parity/runner.py | 10 +- .../shared/parity/test_parity.py | 52 ++- .../strategies/e2e_parity/__init__.py | 4 +- .../e2e_parity/sdk/ocr/test_sdk_parity.py | 295 +++++++++++++++++- .../unit_tests_mapping/cases/ocr.py | 4 + tests/test_litellm/ocr/test_rust_bridge.py | 22 ++ 9 files changed, 453 insertions(+), 30 deletions(-) diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 56c2292d00c..df3f9d2096b 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -29,6 +29,7 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.rust_bridge.bindings import native_exception_types from litellm.rust_bridge.configuration import rust_enabled from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -281,6 +282,33 @@ def _prepare_rust_ocr_call( ) +def _map_rust_ocr_error( + error: Exception, + prepared_request: _PreparedOCRRequest, + exception_types: tuple[type[BaseException], type[BaseException]] | None, +) -> Exception: + if exception_types is None: + return error + _, upstream_error = exception_types + if not isinstance(error, upstream_error): + return error + error_args: Final = cast( # cast-ok: BaseException.args is typed with Any in the standard library stubs + tuple[object, ...], error.args + ) + status_value: Final = error_args[0] if error_args else 0 + message_value: Final = error_args[1] if len(error_args) > 1 else str(error) + status: Final = status_value if isinstance(status_value, int) else 0 + message: Final = message_value if isinstance(message_value, str) else str(message_value) + error_factory: Final = cast( # cast-ok: the legacy provider interface leaves callable parameters untyped + Callable[..., Exception], prepared_request.provider_config.get_error_class + ) + return error_factory( + error_message=message, + status_code=status or 500, + headers={}, # mutable-ok: provider error factories require a concrete header dict + ) + + def _run_rust_ocr( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], @@ -291,16 +319,19 @@ def _run_rust_ocr( prepared_request=prepared_request, resolve_api_key=resolve_api_key, ) - rust_response: Final = rust_ocr_bridge.ocr( - model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, - ) + try: + rust_response: Final = rust_ocr_bridge.ocr( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout=prepared_request.effective_timeout, + ) + except Exception as error: + raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error if rust_response is None: return None return OCRResponse.model_validate(rust_response) @@ -316,16 +347,19 @@ async def _run_rust_aocr( prepared_request=prepared_request, resolve_api_key=resolve_api_key, ) - rust_response: Final = await rust_ocr_bridge.aocr( - model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, - ) + try: + rust_response: Final = await rust_ocr_bridge.aocr( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout=prepared_request.effective_timeout, + ) + except Exception as error: + raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error if rust_response is None: return None return OCRResponse.model_validate(rust_response) diff --git a/tests/rust-python-harness/shared/parity/compare.py b/tests/rust-python-harness/shared/parity/compare.py index adf85e5c8d7..88239ed042a 100644 --- a/tests/rust-python-harness/shared/parity/compare.py +++ b/tests/rust-python-harness/shared/parity/compare.py @@ -78,3 +78,4 @@ def assert_parity(baseline: Execution, candidate: Execution, baseline_user_agent validate_harness(baseline, candidate, baseline_user_agent) assert_request_parity(baseline.requests, candidate.requests) assert_value_parity(baseline.report, candidate.report) + assert_value_parity(baseline.callbacks, candidate.callbacks, path="$.callbacks") diff --git a/tests/rust-python-harness/shared/parity/models.py b/tests/rust-python-harness/shared/parity/models.py index 898b58d23ee..5612e218824 100644 --- a/tests/rust-python-harness/shared/parity/models.py +++ b/tests/rust-python-harness/shared/parity/models.py @@ -37,6 +37,25 @@ class SDKError(BaseModel): llm_provider: str | None +class CallbackObservation(BaseModel): + model_config = ConfigDict(frozen=True) + + hook: Literal[ + "log_success_event", + "async_log_success_event", + "log_failure_event", + "async_log_failure_event", + ] + phase: Literal["success", "failure"] + model: str | None + call_type: str | None + litellm_call_id: str | None + metadata: JsonValue + kwargs: JsonValue + payload: JsonValue + error: SDKError | None + + class SDKJsonChunk(BaseModel): model_config = ConfigDict(frozen=True) @@ -119,6 +138,7 @@ class Execution(BaseModel): requests: tuple[CapturedRequest, ...] report: SDKReport + callbacks: tuple[CallbackObservation, ...] | None = None class SDKCommand(BaseModel): @@ -133,6 +153,7 @@ class WorkerSuccess(BaseModel): status: Literal["ok"] = "ok" report: SDKReport + callbacks: tuple[CallbackObservation, ...] | None = None class WorkerFailure(BaseModel): diff --git a/tests/rust-python-harness/shared/parity/runner.py b/tests/rust-python-harness/shared/parity/runner.py index 43a583382cb..feae0201aee 100644 --- a/tests/rust-python-harness/shared/parity/runner.py +++ b/tests/rust-python-harness/shared/parity/runner.py @@ -39,9 +39,7 @@ class SubprocessRunner: return ( sys.executable, "-m", - ".".join( - self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts - ), + ".".join(self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts), "--parity-worker", provider_url, ) @@ -113,7 +111,11 @@ class SubprocessWorker: ) assert isinstance(result, WorkerSuccess) try: - return Execution(requests=self.provider.take_requests(len(responses)), report=result.report) + return Execution( + requests=self.provider.take_requests(len(responses)), + report=result.report, + callbacks=result.callbacks, + ) except AssertionError: self.provider.reset() raise diff --git a/tests/rust-python-harness/shared/parity/test_parity.py b/tests/rust-python-harness/shared/parity/test_parity.py index 83daccdf8ba..7c8355fbfd8 100644 --- a/tests/rust-python-harness/shared/parity/test_parity.py +++ b/tests/rust-python-harness/shared/parity/test_parity.py @@ -7,7 +7,7 @@ import pytest from pydantic import BaseModel, ConfigDict, JsonValue, PrivateAttr from .compare import assert_model_parity, assert_parity -from .models import CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report +from .models import CallbackObservation, CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report SENTINEL: Final = "python-parity-fallback" @@ -69,6 +69,56 @@ def test_parity_rejects_response_difference() -> None: assert_parity(python, rust, SENTINEL) +def test_parity_distinguishes_unobserved_callbacks_from_zero_events() -> None: + python: Final = _execution(user_agent=SENTINEL) + rust: Final = _execution(user_agent="litellm-rust").model_copy(update={"callbacks": ()}) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + +def test_parity_rejects_callback_payload_difference() -> None: + observation: Final = CallbackObservation( + hook="log_success_event", + phase="success", + model="test-model", + call_type="ocr", + litellm_call_id="test-call", + metadata={"profile": "success"}, + kwargs={"model": "test-model"}, + payload={"model": "test-model", "pages": []}, + error=None, + ) + python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)}) + rust: Final = _execution(user_agent="litellm-rust").model_copy( + update={"callbacks": (observation.model_copy(update={"payload": {"model": "changed", "pages": []}}),)} + ) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + +def test_parity_rejects_callback_kwargs_difference() -> None: + observation: Final = CallbackObservation( + hook="log_success_event", + phase="success", + model="test-model", + call_type="ocr", + litellm_call_id="test-call", + metadata={"profile": "success"}, + kwargs={"model": "test-model"}, + payload={"model": "test-model", "pages": []}, + error=None, + ) + python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)}) + rust: Final = _execution(user_agent="litellm-rust").model_copy( + update={"callbacks": (observation.model_copy(update={"kwargs": {"model": "changed"}}),)} + ) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + def test_parity_rejects_error_difference() -> None: python: Final = Execution( requests=(), diff --git a/tests/rust-python-harness/strategies/e2e_parity/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/__init__.py index f668e178eef..95346bbd496 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/__init__.py +++ b/tests/rust-python-harness/strategies/e2e_parity/__init__.py @@ -18,8 +18,8 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( coverage=Coverage.PARTIAL, module="tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.test_sdk_parity", note=( - "Recorded sync/async SDK parity; invalid-model provider errors differ, " - "and Reducto lacks a Rust contract." + "Recorded sync/async SDK parity with focused success/error callback profiles; " + "Reducto lacks a Rust contract, and known provider parity gaps remain." ), ), surface="sdk", diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py index e72980752f2..5c4abc78081 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py @@ -1,23 +1,31 @@ from __future__ import annotations import asyncio +import datetime +import queue import sys import tempfile +import time import traceback -from collections.abc import Callable, Coroutine, Generator +from collections.abc import Callable, Coroutine, Generator, Mapping from contextlib import contextmanager from enum import Enum from functools import partial from pathlib import Path from typing import Annotated, Final, Literal, cast +from urllib.parse import urlsplit, urlunsplit from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse from .....shared.parity.compare import assert_parity from .....shared.parity.fixtures.store import fixture_id, recorded_fixtures from .....shared.parity.models import ( + JSON_VALUE_ADAPTER, + CallbackObservation, + Execution, SDKCommand, SDKError, SDKReport, @@ -39,6 +47,9 @@ from .fixtures.config import configured_fixture_directory from .fixtures.models import OcrParityCase, OcrSdkInput API_KEY: Final = "test-key" +CALLBACK_DELAY_SECONDS: Final = 0.05 +CALLBACK_DRAIN_TIMEOUT_SECONDS: Final = 10.0 +CALLBACK_TERMINALS: Final[tuple[Literal["success", "failure"], ...]] = ("success", "failure") PYTHON_HTTP_SENTINEL: Final = "python-ocr-parity-fallback" PYTHON_VARIANT: Final = ExecutionVariant(name="Python", environment=(("LITELLM_RUST", "0"),)) RUST_VARIANT: Final = ExecutionVariant(name="Rust", environment=(("LITELLM_RUST", "1"),)) @@ -75,10 +86,148 @@ class InvalidOcrWorkerCase(BaseModel): case: InvalidOcrCase -OcrWorkerCase = Annotated[RecordedOcrWorkerCase | InvalidOcrWorkerCase, Field(discriminator="kind")] +class CallbackOcrWorkerCase(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["callback"] = "callback" + case: OcrParityCase + terminal: Literal["success", "failure"] + + +OcrWorkerCase = Annotated[ + RecordedOcrWorkerCase | InvalidOcrWorkerCase | CallbackOcrWorkerCase, + Field(discriminator="kind"), +] OCR_WORKER_CASE_ADAPTER: Final[TypeAdapter[OcrWorkerCase]] = TypeAdapter(OcrWorkerCase) +class RecordingCallback(CustomLogger): + def __init__(self) -> None: + self.message_logging: Final = True + self.turn_off_message_logging: Final = False + self._observations: Final[queue.SimpleQueue[CallbackObservation]] = queue.SimpleQueue() + + def _normalized_kwargs(self, value: object, key: str | None = None) -> JsonValue: + if value is None or isinstance(value, (bool, int, float)): + return value + if isinstance(value, str): + if key != "api_base": + return value + parsed: Final = urlsplit(value) + return urlunsplit(("", "", parsed.path, parsed.query, parsed.fragment)) + if isinstance(value, datetime.datetime): + return "datetime" + if isinstance(value, Exception): + return sdk_error_report(value).model_dump(mode="json") + if isinstance(value, BaseModel): + return self._normalized_kwargs(value.model_dump(mode="json"), key) + if isinstance(value, Mapping): + if any(not isinstance(map_key, str) for map_key in value): + raise TypeError("callback kwarg mappings must use string keys") + return { + map_key: self._normalized_kwargs(map_value, map_key) + for map_key, map_value in value.items() + } + if isinstance(value, (list, tuple)): + return [self._normalized_kwargs(item) for item in value] + raise TypeError(f"unsupported callback kwarg type: {type(value)}") + + def _record( + self, + hook: Literal[ + "log_success_event", + "async_log_success_event", + "log_failure_event", + "async_log_failure_event", + ], + phase: Literal["success", "failure"], + kwargs: dict[str, object], + response_obj: object, + ) -> None: + raw_litellm_params: Final = kwargs.get("litellm_params") + litellm_params: Final[Mapping[str, object]] = ( + cast(Mapping[str, object], raw_litellm_params) if isinstance(raw_litellm_params, Mapping) else {} + ) + raw_metadata: Final = litellm_params.get("metadata") + metadata_mapping: Final[Mapping[str, object]] = ( + cast(Mapping[str, object], raw_metadata) if isinstance(raw_metadata, Mapping) else {} + ) + metadata: Final = JSON_VALUE_ADAPTER.validate_python( + {key: metadata_mapping[key] for key in ("callback_profile", "sdk_route") if key in metadata_mapping} + ) + raw_error: Final = kwargs.get("exception") + error: Final = sdk_error_report(raw_error) if isinstance(raw_error, Exception) else None + normalized_kwargs: Final = self._normalized_kwargs(kwargs) + payload_source: Final = ( + response_obj.model_dump(mode="json") if isinstance(response_obj, BaseModel) else response_obj + ) + payload: Final = JSON_VALUE_ADAPTER.validate_python(payload_source) + raw_model: Final = kwargs.get("model") + raw_call_type: Final = kwargs.get("call_type") + raw_call_id: Final = kwargs.get("litellm_call_id") + self._observations.put( + CallbackObservation( + hook=hook, + phase=phase, + model=raw_model if isinstance(raw_model, str) else None, + call_type=str(raw_call_type) if raw_call_type is not None else None, + litellm_call_id=raw_call_id if isinstance(raw_call_id, str) else None, + metadata=metadata, + kwargs=normalized_kwargs, + payload=payload, + error=error, + ) + ) + + def log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + time.sleep(CALLBACK_DELAY_SECONDS) + self._record("log_success_event", "success", kwargs, response_obj) + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + await asyncio.sleep(CALLBACK_DELAY_SECONDS) + self._record("async_log_success_event", "success", kwargs, response_obj) + + def log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + time.sleep(CALLBACK_DELAY_SECONDS) + self._record("log_failure_event", "failure", kwargs, response_obj) + + async def async_log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + await asyncio.sleep(CALLBACK_DELAY_SECONDS) + self._record("async_log_failure_event", "failure", kwargs, response_obj) + + def observations(self) -> tuple[CallbackObservation, ...]: + observations: Final = tuple(self._observations.get_nowait() for _ in range(self._observations.qsize())) + return tuple(sorted(observations, key=lambda observation: observation.hook)) + + INVALID_OCR_CASES: Final = ( InvalidOcrCase( name="unsupported_provider", @@ -229,6 +378,106 @@ def _execute_invalid_sdk_case( return _execute_sdk_call(call_kwargs, route, event_loop) +def _callback_call_id(route: SDKRoute, terminal: Literal["success", "failure"]) -> str: + return f"ocr-callback-{route.value}-{terminal}" + + +def _callback_metadata(route: SDKRoute, terminal: Literal["success", "failure"]) -> dict[str, str]: + return {"callback_profile": terminal, "sdk_route": route.value} + + +def _drain_callback_delivery(route: SDKRoute, event_loop: asyncio.AbstractEventLoop) -> None: + if route is SDKRoute.AOCR: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + async def drain_async_callbacks() -> None: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=CALLBACK_DRAIN_TIMEOUT_SECONDS) + await GLOBAL_LOGGING_WORKER.stop() + + event_loop.run_until_complete(drain_async_callbacks()) + + from litellm.litellm_core_utils.thread_pool_executor import executor + + executor.shutdown(wait=True, cancel_futures=False) + + +def _execute_callback_sdk_case( + case: OcrParityCase, + route: SDKRoute, + terminal: Literal["success", "failure"], + mock_url: str, + event_loop: asyncio.AbstractEventLoop, +) -> WorkerSuccess: + callback: Final = RecordingCallback() + call_kwargs: Final = { + **_call_kwargs(case.litellm_input, mock_url, route), + "callbacks": [callback], + "litellm_call_id": _callback_call_id(route, terminal), + "litellm_trace_id": _callback_call_id(route, terminal), + "metadata": _callback_metadata(route, terminal), + } + report: Final = _execute_sdk_call(call_kwargs, route, event_loop) + _drain_callback_delivery(route, event_loop) + return WorkerSuccess(report=report, callbacks=callback.observations()) + + +def _assert_callback_lifecycle( + execution: Execution, + route: SDKRoute, + terminal: Literal["success", "failure"], +) -> None: + observations: Final = execution.callbacks + assert observations is not None, f"{route.value} {terminal} callbacks were not observed" + expected_hooks: Final = ( + (f"log_{terminal}_event",) + if route is SDKRoute.OCR + else ("async_log_success_event",) + if terminal == "success" + else ("async_log_failure_event", "log_failure_event") + ) + actual_hooks: Final = tuple(observation.hook for observation in observations) + assert actual_hooks == expected_hooks, ( + f"{route.value} {terminal} expected callback hooks {expected_hooks}, received {actual_hooks}" + ) + expected_call_id: Final = _callback_call_id(route, terminal) + expected_metadata: Final = _callback_metadata(route, terminal) + for observation in observations: + assert observation.phase == terminal + assert observation.call_type == route.value + assert observation.litellm_call_id == expected_call_id + assert observation.metadata == expected_metadata + assert observation.model + if terminal == "success": + assert isinstance(execution.report, SDKSuccess) + assert observation.payload == execution.report.response + assert observation.error is None + else: + assert isinstance(execution.report, SDKError) + assert observation.payload is None + assert observation.error is not None + assert observation.error.exception_type + assert observation.error.message + assert observation.error.status_code is not None + assert observation.error.status_code >= 400 + + +def _check_callback_ocr_sdk_parity( + case: OcrParityCase, + route: SDKRoute, + terminal: Literal["success", "failure"], + case_file: Path, + runner: SubprocessRunner, +) -> None: + with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers: + python_worker, rust_worker = workers + python: Final = python_worker.execute(case_file, route.value, case.provider_responses) + rust: Final = rust_worker.execute(case_file, route.value, case.provider_responses) + + _assert_callback_lifecycle(python, route, terminal) + _assert_callback_lifecycle(rust, route, terminal) + assert_parity(python, rust, PYTHON_HTTP_SENTINEL) + + def _check_recorded_ocr_sdk_parity( ocr_fixture: OcrParityCase, route: SDKRoute, @@ -276,6 +525,25 @@ def _write_worker_case(directory: Path, index: int, case: OcrWorkerCase) -> Path return case_file +def _callback_fixture( + fixtures: tuple[OcrParityCase, ...], + terminal: Literal["success", "failure"], +) -> OcrParityCase: + matching: Final = tuple( + fixture + for fixture in fixtures + if fixture.litellm_input.contract == "mistral" + and ( + all(response.status_code < 400 for response in fixture.provider_responses) + if terminal == "success" + else any(response.status_code >= 400 for response in fixture.provider_responses) + ) + ) + if not matching: + raise AssertionError(f"no recorded Mistral OCR {terminal} fixture is available for callback parity") + return min(matching, key=lambda fixture: fixture_id(fixture.litellm_input, fixture.litellm_input.model)) + + @contextmanager def parity_checks() -> Generator[tuple[E2ECheck, ...]]: fixtures: Final = tuple( @@ -298,6 +566,17 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]: _write_worker_case(directory, len(recorded_files) + index, InvalidOcrWorkerCase(case=case)) for index, case in enumerate(INVALID_OCR_CASES) ) + callback_cases: Final[tuple[tuple[Literal["success", "failure"], OcrParityCase], ...]] = tuple( + (terminal, _callback_fixture(fixtures, terminal)) for terminal in CALLBACK_TERMINALS + ) + callback_files: Final = tuple( + _write_worker_case( + directory, + len(recorded_files) + len(invalid_files) + index, + CallbackOcrWorkerCase(case=case, terminal=terminal), + ) + for index, (terminal, case) in enumerate(callback_cases) + ) with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers: recorded: Final = tuple( E2ECheck( @@ -315,7 +594,15 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]: for case, case_file in zip(INVALID_OCR_CASES, invalid_files, strict=True) for route in SDKRoute ) - yield (*recorded, *invalid) + callbacks: Final = tuple( + E2ECheck( + f"callback:{route.value}:{terminal}", + partial(_check_callback_ocr_sdk_parity, case, route, terminal, case_file, runner), + ) + for (terminal, case), case_file in zip(callback_cases, callback_files, strict=True) + for route in SDKRoute + ) + yield (*recorded, *invalid, *callbacks) def _execute_worker_command( @@ -333,6 +620,8 @@ def _execute_worker_command( return WorkerSuccess(report=_execute_sdk_case(recorded.litellm_input, route, mock_url, event_loop)) case InvalidOcrWorkerCase(case=invalid): return WorkerSuccess(report=_execute_invalid_sdk_case(invalid, route, mock_url, event_loop)) + case CallbackOcrWorkerCase(case=callback_case, terminal=terminal): + return _execute_callback_sdk_case(callback_case, route, terminal, mock_url, event_loop) except Exception: return WorkerFailure(error=traceback.format_exc()) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py index 3e6c4060134..0e771f0dc17 100644 --- a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py @@ -230,6 +230,10 @@ _HOST_ONLY_BRIDGE_EXCLUSIONS: Final = tuple( "test_ocr_exception_type_uses_resolved_provider_context", "Python wraps bridge exceptions into public errors.", ), + ( + "test_rust_upstream_error_uses_ocr_provider_error_mapping", + "Python maps native upstream errors through the selected OCR provider config.", + ), ("test_aocr_routes_to_async_rust_when_enabled", "Python selects and invokes the async native bridge."), ("test_aocr_exception_type_uses_resolved_provider_context", "Python wraps async bridge exceptions."), ("test_ocr_forwards_timeout_to_rust", "Python converts and forwards explicit timeouts."), diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index d28184b96a6..c34833221cc 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -9,6 +9,7 @@ import httpx import pytest import litellm +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge import configuration @@ -39,6 +40,10 @@ class CapturedException(Exception): pass +class RustUpstreamError(Exception): + pass + + class RecordingBridge: """A fake ``RustOcr`` callable that records the args it was handed.""" @@ -183,6 +188,9 @@ class FakeOCRConfig: ) -> str: return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" + def get_error_class(self, error_message: str, status_code: int, headers: dict[str, str]) -> BaseLLMException: + return BaseLLMException(status_code=status_code, message=error_message, headers=headers) + def build_prepared_request( *, @@ -486,6 +494,20 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): } +def test_rust_upstream_error_uses_ocr_provider_error_mapping(): + error = RustUpstreamError(400, '{"message":"invalid model"}') + + mapped = ocr_main._map_rust_ocr_error( + error, + build_prepared_request(), + (RuntimeError, RustUpstreamError), + ) + + assert isinstance(mapped, BaseLLMException) + assert mapped.status_code == 400 + assert mapped.message == '{"message":"invalid model"}' + + def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): bridge = RecordingBridge() litellm.rust(True) From 4b3355bdc6b9807d5d1c817094449df03b6af881 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 7 Sep 2026 11:30:46 -0700 Subject: [PATCH 49/49] test(e2e): prove the virtual key lifecycle on every gateway replica (#40023) * test(e2e): prove the virtual key lifecycle on every replica Walks one virtual key through create, read, partial update, clear, enforce and delete against a live proxy and database, reading every write back on every gateway replica. The management suite already had single write-then-read tests for keys, but none of them proved that a partial /key/update leaves the untouched fields alone, that an explicit null clears a field, or that a write is visible on more than the one gateway that took it. Adds read_back_everywhere to the shared ProxyClient: it polls a GET path on every URL in PROXY_REPLICA_URLS until each replica's parsed body satisfies the caller's predicate, and fails naming the replica that never converged. The CLEAR sentinel in the e2e models makes an explicit JSON null expressible in a body the transport otherwise strips of None fields. Documents /key/update's merge patch semantics on the endpoint docstring. * test(e2e): prove key revocation and field preservation on every replica Applies the findings from an adversarial review of the first commit. The delete step only checked that chat was refused on the gateway that took the write, so it would have passed while a sibling gateway kept serving the deleted key. It now serves one call from every replica first, so each has the key cached and the delete has something to revoke everywhere, then polls every replica for the refusal. The file also carried its own poll loop that tested the deadline before attempting, so it gave up one attempt early and skipped the attempt landing exactly on the deadline. It now shares the harness helper, which is generic over the polled value rather than over a parsed body, so the same loop covers both the info read-back and the chat refusal. The model the enforcement step registers now carries a unique marker in its alias, matching every other deployment this suite creates, so concurrent runs never share one model group. The docstring sentence claimed an explicit null clears any field. It does not: the metadata-backed fields merge into stored metadata, where a null is a silent no-op, and only the key's own columns clear. Regenerating the dashboard types picks up the corrected text. * fix(e2e): delete a deployment that never becomes servable Registering a model posts /model/new and then waits for every replica to list it. When that wait timed out the deployment already existed in the database but its id had never been returned, so no caller could delete it and the row outlived the run. It is now deleted before the failure propagates. Found by review on the key lifecycle suite, whose module fixture registers a deployment this way, but every caller of the shared helper had the same exposure. * docs(e2e): drop the duplicated notes from the lifecycle docstrings The delete method restated what the warm-up helper already explains, and the module restated the merge patch rule that the endpoint and the request model both document. --- .../key_management_endpoints.py | 5 + tests/e2e/coverage_registry/mgmt.yaml | 2 + .../e2e/management/test_key_lifecycle_e2e.py | 299 ++++++++++++++++++ tests/e2e/models.py | 36 ++- tests/e2e/proxy_client.py | 152 ++++++++- tests/e2e/test_proxy_client.py | 114 ++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 7 files changed, 592 insertions(+), 21 deletions(-) create mode 100644 tests/e2e/management/test_key_lifecycle_e2e.py diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ebcfab090b5..8e51e250319 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2840,6 +2840,11 @@ async def update_key_fn( """ Update an existing API key's parameters. + The body is a merge patch: a field left out keeps its stored value, and on the key's own columns + an explicit null clears it. The metadata-backed fields below are the exception, merging into the + stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself + replaces the stored metadata wholesale. + Parameters: - key: Optional[str] - The key to update. Either key or key_alias must be provided. - key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index d571fb36546..860d96a50b4 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -5,6 +5,8 @@ - {id: mgmt.key.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:2462", rationale: "Budget/model changes persist"} - {id: mgmt.key.update.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:2462", rationale: "Non-admin cannot escalate perms"} - {id: mgmt.key.update.happy_path, module: mgmt, tier: P1, surface: ui, assertions: [happy_path], source: "key_management_endpoints.py:2462", rationale: "Key edit through the dashboard"} +- {id: mgmt.key.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "key_management_endpoints.py:2829", rationale: "A partial /key/update changes only the field it names; alias, models, limits, budget window, team and metadata read back unchanged on every gateway replica"} +- {id: mgmt.key.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "key_management_endpoints.py:2829", rationale: "An explicit null on /key/update clears max_budget and budget_duration, and the derived budget_reset_at with it, on every gateway replica"} - {id: mgmt.key.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3122", rationale: "Deletion revokes future calls"} - {id: mgmt.key.delete.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:3122", rationale: "Non-owner cannot delete"} - {id: mgmt.key.info.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3380", rationale: "Info reflects all writes"} diff --git a/tests/e2e/management/test_key_lifecycle_e2e.py b/tests/e2e/management/test_key_lifecycle_e2e.py new file mode 100644 index 00000000000..4c8effc4d24 --- /dev/null +++ b/tests/e2e/management/test_key_lifecycle_e2e.py @@ -0,0 +1,299 @@ +"""Live e2e: one virtual key walked through its whole lifecycle, read back on every +gateway replica. + +Create, read, partial update, clear, enforce, delete: one method per step, and every +step creates its own team and key (both deleted on teardown) so a step reruns or skips +on its own. Writes go through the control plane; read-backs poll every URL in +PROXY_REPLICA_URLS until each replica converges, because a write that is visible on the +gateway that took it and stale on its neighbour is exactly the failure this file exists +to catch. Revocation is the slowest of those: a deleted key stays usable on the other +replicas until their auth cache entry expires, so the delete step polls each of them +rather than asserting once. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +import pytest + +from e2e_config import unique_marker +from e2e_http import Result, StreamingResponse, Success, UnknownApiError, unwrap +from lifecycle import ResourceManager +from management_client import MODEL_ACCESS_DENIED_MARKER, ManagementClient +from models import ( + CLEAR, + ChatBody, + ChatMessage, + KeyGenerateBody, + KeyGenerateResponse, + KeyInfo, + KeyInfoParams, + KeyInfoResponse, + KeyMetadata, + KeyUpdateBody, + LiteLLMParamsBody, + TeamNewBody, +) +from proxy_client import Converged, NotConverged, Poller, await_converged, await_converged_everywhere +from transport import Transport + +pytestmark = pytest.mark.e2e + +BACKING_MODEL: Final = "gpt-4o-mini" +DENIED_MODEL: Final = "gpt-5.5" +MAX_BUDGET: Final = 25.0 +TPM_LIMIT: Final = 313131 +RPM_LIMIT: Final = 323232 +UPDATED_RPM_LIMIT: Final = 424242 +BUDGET_DURATION: Final = "30d" + + +@dataclass(frozen=True, slots=True) +class CreatedKey: + written: KeyGenerateBody + response: KeyGenerateResponse + + @property + def key(self) -> str: + return self.response.key + + +@pytest.fixture(scope="module") +def mock_deployment(client: ManagementClient) -> Iterator[str]: + """A deployment that answers from a canned response, so the enforcement step needs no + provider key. The alias carries a unique marker, like every other model this suite + registers, so concurrent runs never share one model group.""" + model_name: Final = f"e2e-key-lifecycle-{unique_marker()}" + model_id: Final = client.proxy.create_model(model_name, LiteLLMParamsBody(model=BACKING_MODEL, mock_response="ok")) + try: + yield model_name + finally: + client.proxy.delete_model(model_id) + + +def _await[T](client: ManagementClient, poller: Poller[T], converged: Callable[[T], bool], failure: str) -> T: + outcome: Final = await_converged( + poller, + converged=converged, + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case Converged(result=result): + return result + case NotConverged(last_result=last): + pytest.fail(f"{failure}; last outcome: {last}") + + +def _chat_poller(transport: Transport, key: str, model: str) -> Poller[StreamingResponse]: + return lambda: transport.send( + "/chat/completions", + headers=transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"say hi {unique_marker()}")], + max_tokens=16, + ), + ) + + +def _create_key(client: ManagementClient, resources: ResourceManager, model: str) -> CreatedKey: + marker: Final = unique_marker() + team_id: Final = client.create_team(TeamNewBody(team_alias=f"e2e-key-lifecycle-team-{marker}")) + resources.defer(lambda: client.delete_team(team_id)) + written: Final = KeyGenerateBody( + key_alias=f"e2e-key-lifecycle-{marker}", + models=[model], + max_budget=MAX_BUDGET, + tpm_limit=TPM_LIMIT, + rpm_limit=RPM_LIMIT, + budget_duration=BUDGET_DURATION, + metadata=KeyMetadata(tag=marker), + team_id=team_id, + ) + response: Final = unwrap(client.generate_key(written)) + resources.defer(lambda: client.proxy.delete_key(response.key)) + return CreatedKey(written=written, response=response) + + +def _key_info_everywhere( + client: ManagementClient, key: str, settled: Callable[[KeyInfo], bool] +) -> Mapping[str, KeyInfo]: + def converged(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, Success) and settled(result.data.info) + + reads: Final = client.proxy.read_back_everywhere( + "/key/info", params=KeyInfoParams(key=key), response_type=KeyInfoResponse, converged=converged + ) + return MappingProxyType({replica: unwrap(read).info for replica, read in reads.items()}) + + +def _is_key_not_found(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, UnknownApiError) and result.status_code == 404 + + +def _assert_reads_back(info: KeyInfo, expected: KeyGenerateBody, replica: str) -> None: + for field, observed, wanted in ( + ("key_alias", info.key_alias, expected.key_alias), + ("models", info.models, expected.models), + ("max_budget", info.max_budget, expected.max_budget), + ("tpm_limit", info.tpm_limit, expected.tpm_limit), + ("rpm_limit", info.rpm_limit, expected.rpm_limit), + ("budget_duration", info.budget_duration, expected.budget_duration), + ("team_id", info.team_id, expected.team_id), + ("metadata", info.metadata, expected.metadata), + ): + assert observed == wanted, f"{replica}: /key/info reports {field}={observed!r}, expected {wanted!r}" + + +def _poll_chat_ok(client: ManagementClient, key: str, model: str) -> None: + _ = _await( + client, + _chat_poller(client.proxy.transport, key, model), + lambda outcome: outcome.ok, + f"chat on {model} never succeeded for the key before the deadline", + ) + + +def _warm_every_replica(client: ManagementClient, key: str, model: str) -> None: + """Serve one call from every replica, so each has the key in its auth cache. Without + this the revocation check below would only prove a replica rejects a key it never + knew, which is true of any random string.""" + for replica, transport in client.proxy.replicas.items(): + _ = _await( + client, + _chat_poller(transport, key, model), + lambda outcome: outcome.ok, + f"{replica}: chat on {model} never succeeded for the key before the deadline", + ) + + +def _assert_chat_rejected_everywhere(client: ManagementClient, key: str, model: str) -> None: + outcomes: Final = await_converged_everywhere( + {replica: _chat_poller(transport, key, model) for replica, transport in client.proxy.replicas.items()}, + converged=lambda outcome: outcome.status_code == 401, + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + for replica, outcome in outcomes.items(): + assert isinstance(outcome, Converged), ( + f"{replica}: the deleted key was still accepted on chat after " + f"{client.proxy.poll_timeout}s, last status {outcome.last_result.status_code}" + ) + + +class TestKeyLifecycle: + def test_create_echoes_every_field_written( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + response: Final = created.response + for field, observed, wanted in ( + ("key_alias", response.key_alias, created.written.key_alias), + ("models", response.models, created.written.models), + ("max_budget", response.max_budget, created.written.max_budget), + ("tpm_limit", response.tpm_limit, created.written.tpm_limit), + ("rpm_limit", response.rpm_limit, created.written.rpm_limit), + ("budget_duration", response.budget_duration, created.written.budget_duration), + ("team_id", response.team_id, created.written.team_id), + ("metadata", response.metadata, created.written.metadata), + ): + assert observed == wanted, f"/key/generate echoed {field}={observed!r}, sent {wanted!r}" + + def test_read_reflects_the_create_on_every_replica( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + infos: Final = _key_info_everywhere( + client, created.key, lambda info: info.key_alias == created.written.key_alias + ) + for replica, info in infos.items(): + _assert_reads_back(info, created.written, replica) + assert info.budget_reset_at is not None, ( + f"{replica}: /key/info reports no budget_reset_at for budget_duration={BUDGET_DURATION!r}" + ) + + @pytest.mark.covers("mgmt.key.update.preserves_unrelated_fields") + def test_partial_update_changes_only_the_named_field( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + before: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == RPM_LIMIT) + + _ = unwrap(client.update_key(KeyUpdateBody(key=created.key, rpm_limit=UPDATED_RPM_LIMIT))) + + after: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == UPDATED_RPM_LIMIT) + for replica, info in after.items(): + _assert_reads_back(info, created.written.model_copy(update={"rpm_limit": UPDATED_RPM_LIMIT}), replica) + assert info.budget_reset_at == before[replica].budget_reset_at, ( + f"{replica}: budget_reset_at moved from {before[replica].budget_reset_at!r} to " + f"{info.budget_reset_at!r} on a /key/update that did not name budget_duration" + ) + + @pytest.mark.covers("mgmt.key.update.clear_persists") + def test_explicit_null_clears_the_budget_and_its_reset_time( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + _ = _key_info_everywhere(client, created.key, lambda info: info.max_budget == MAX_BUDGET) + + _ = unwrap(client.update_key(KeyUpdateBody(key=created.key, max_budget=CLEAR, budget_duration=CLEAR))) + + cleared: Final = _key_info_everywhere(client, created.key, lambda info: info.max_budget is None) + for replica, info in cleared.items(): + assert info.budget_duration is None, ( + f"{replica}: budget_duration={info.budget_duration!r} survived an explicit null" + ) + assert info.budget_reset_at is None, ( + f"{replica}: clearing budget_duration left budget_reset_at={info.budget_reset_at!r}" + ) + _assert_reads_back( + info, created.written.model_copy(update={"max_budget": None, "budget_duration": None}), replica + ) + + def test_key_serves_its_model_and_is_denied_others( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + _poll_chat_ok(client, created.key, mock_deployment) + + denied: Final = client.chat_status(created.key, DENIED_MODEL, f"say hi {unique_marker()}") + assert denied.status_code == 403, ( + f"chat on {DENIED_MODEL!r} outside the key's model list must be denied 403, got " + f"{denied.status_code}: {denied.body[:300]}" + ) + assert MODEL_ACCESS_DENIED_MARKER in denied.body, ( + f"403 body must be a model-access denial, got: {denied.body[:300]}" + ) + + def test_delete_revokes_info_and_chat_on_every_replica( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + """The teardown's deferred delete fires again on the already-deleted key by + design: the deferred cleanup must survive this test failing before the + in-body delete, and a repeat /key/delete is a cheap no-op the warn-only + teardown absorbs.""" + created: Final = _create_key(client, resources, mock_deployment) + _warm_every_replica(client, created.key, mock_deployment) + + client.delete_key_strict(created.key) + + _ = client.proxy.read_back_everywhere( + "/key/info", + params=KeyInfoParams(key=created.key), + response_type=KeyInfoResponse, + converged=_is_key_not_found, + ) + _assert_chat_rejected_everywhere(client, created.key, mock_deployment) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index c01687d0b31..fa124bdf6bd 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -8,9 +8,9 @@ from __future__ import annotations from collections.abc import Sequence from datetime import datetime -from typing import Literal +from typing import Final, Literal -from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_validator +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator # ---------- keys ---------- @@ -49,6 +49,7 @@ class KeyMetadata(BaseModel): logging: list[KeyLoggingCallback] | None = None priority: str | None = None batch_enqueued_token_limit: int | None = None + tag: str | None = None class ObjectPermission(BaseModel): @@ -81,6 +82,14 @@ class KeyGenerateBody(BaseModel): class KeyGenerateResponse(BaseModel): key: str + key_alias: str | None = None + models: list[str] = [] + max_budget: float | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + budget_duration: str | None = None + team_id: str | None = None + metadata: KeyMetadata | None = None class KeyRegenerateBody(BaseModel): @@ -122,6 +131,7 @@ class KeyInfo(BaseModel): blocked: bool | None = None spend: float | None = None max_budget: float | None = None + budget_duration: str | None = None budget_reset_at: str | None = None budget_id: str | None = None litellm_budget_table: LiteLLMBudgetTable | None = None @@ -919,12 +929,34 @@ class CredentialCreateResponse(BaseModel): # ---------- key / team / user / organization management ---------- +class Cleared(BaseModel): + """An explicit JSON null in a merge-patch body. The transport drops `None` fields + before sending (`exclude_none`), so `None` means "leave the stored value alone"; a + field set to `CLEAR` reaches the wire as `null`, which tells the proxy to clear it.""" + + model_config = ConfigDict(frozen=True) + + @model_serializer + def _as_null(self) -> None: + return None + + +CLEAR: Final = Cleared() + + class KeyUpdateBody(BaseModel): + """POST /key/update is a merge patch: a field left `None` is dropped from the body and + keeps its stored value, `CLEAR` sends an explicit null that clears it (`budget_duration` + clears `budget_reset_at` with it), and `metadata` replaces the stored metadata wholesale.""" + key: str models: list[str] | None = None key_alias: str | None = None tpm_limit: int | None = None rpm_limit: int | None = None + max_budget: float | Cleared | None = None + budget_duration: str | Cleared | None = None + metadata: KeyMetadata | None = None class KeyBlockBody(BaseModel): diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 3b15e57dfcf..520cbfde5a9 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -16,6 +16,8 @@ from datetime import datetime from types import MappingProxyType from typing import Final +from pydantic import BaseModel + from e2e_http import ( AnthropicHeaders, AuthHeaders, @@ -149,9 +151,7 @@ def await_servable( last_result: Result[ModelsListResponse] | None = None while True: t = now() - phase_deadline = ( - started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds - ) + phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds remaining = phase_deadline - t if remaining <= 0: if ( @@ -164,9 +164,7 @@ def await_servable( poll_timeout = min(request_timeout, remaining) last_result = list_models(poll_timeout) - listed = isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ) + listed = isinstance(last_result, Success) and any(entry.id == model_name for entry in last_result.data.data) t = now() if not listed: first_seen_at = None @@ -179,9 +177,7 @@ def await_servable( elif t - first_seen_at >= db_sync_seconds: return Servable() - phase_deadline = ( - started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds - ) + phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds wait = min(interval, phase_deadline - now()) if wait > 0: sleep(wait) @@ -239,6 +235,88 @@ def servable_timeout_message( ) +type Poller[T] = Callable[[], T] + + +@dataclass(frozen=True, slots=True) +class Converged[T]: + result: T + + +@dataclass(frozen=True, slots=True) +class NotConverged[T]: + """The deadline passed without a read satisfying the predicate; `last_result` is + the final read, so the caller can tell a stale body from a failed request.""" + + last_result: T + + +type ConvergeOutcome[T] = Converged[T] | NotConverged[T] + + +def await_converged[T]( + poll: Poller[T], + *, + converged: Callable[[T], bool], + timeout: float, + interval: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> ConvergeOutcome[T]: + """Poll until a read satisfies `converged` or `timeout` elapses. + + Polls before testing the deadline, so a zero or already-spent budget still gets one + attempt, and sleeps only min(interval, time left), so the attempt that lands exactly + on the deadline is taken rather than skipped. Clock and sleep are injected.""" + deadline: Final = now() + timeout + while True: + result = poll() + if converged(result): + return Converged(result=result) + remaining = deadline - now() + if remaining <= 0: + return NotConverged(last_result=result) + sleep(min(interval, remaining)) + + +def await_converged_everywhere[T]( + pollers: Mapping[str, Poller[T]], + *, + converged: Callable[[T], bool], + timeout: float, + interval: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> Mapping[str, ConvergeOutcome[T]]: + """`await_converged` against every replica in turn, each with the full budget, so a + replica that lags behind the one a write landed on is polled until it catches up + rather than failing on its first stale read.""" + return MappingProxyType( + { + replica: await_converged( + poll, converged=converged, timeout=timeout, interval=interval, now=now, sleep=sleep + ) + for replica, poll in pollers.items() + } + ) + + +def first_lagging_replica[T]( + outcomes: Mapping[str, ConvergeOutcome[T]], +) -> tuple[str, NotConverged[T]] | None: + return next( + ((replica, outcome) for replica, outcome in outcomes.items() if isinstance(outcome, NotConverged)), + None, + ) + + +def converge_timeout_message(*, what: str, replica: str, timeout: float, last_result: object) -> str: + return ( + f"{what} on {replica} never converged within {timeout}s of the write " + f"(control/data-plane propagation issue); last read: {last_result}" + ) + + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport @@ -290,6 +368,52 @@ class ProxyClient: ) ).info + def read_back_everywhere[R: BaseModel]( + self, + path: str, + *, + params: BaseModel, + response_type: type[R], + converged: Callable[[Result[R]], bool], + ) -> Mapping[str, Result[R]]: + """GET `path` under the master key on every replica in PROXY_REPLICA_URLS (the + data-plane URL alone when the stack exports no per-gateway addresses), polling + each to poll_timeout until its read satisfies `converged`. Returns that read per + replica, or fails naming the first replica that never converged and its last + read. Behind a load balancer the single address proves one replica converged, + not all of them; only per-gateway addresses make this a fleet-wide proof.""" + outcomes: Final = await_converged_everywhere( + { + url: self._body_poller(transport, path, params, response_type) + for url, transport in self.replicas.items() + }, + converged=converged, + timeout=self.poll_timeout, + interval=self.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + lagging: Final = first_lagging_replica(outcomes) + if lagging is not None: + replica, outcome = lagging + raise AssertionError( + converge_timeout_message( + what=f"GET {path}", + replica=replica, + timeout=self.poll_timeout, + last_result=outcome.last_result, + ) + ) + return MappingProxyType( + {replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)} + ) + + @staticmethod + def _body_poller[R: BaseModel]( + transport: Transport, path: str, params: BaseModel, response_type: type[R] + ) -> Poller[Result[R]]: + return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type) + def model_info(self) -> list[ModelInfoEntry]: """Every configured deployment with the price the proxy resolved for it (config override merged over cost-map defaults).""" @@ -320,9 +444,7 @@ class ProxyClient: response_type=FileListResponse, ) - def list_fine_tuning_jobs( - self, key: str, params: FineTuningJobsParams - ) -> Result[FineTuningJobsResponse]: + def list_fine_tuning_jobs(self, key: str, params: FineTuningJobsParams) -> Result[FineTuningJobsResponse]: return self.transport.get( "/v1/fine_tuning/jobs", headers=self.transport.bearer(key), @@ -375,7 +497,11 @@ class ProxyClient: ) ).model_id written_at = time.monotonic() - self._await_model_servable(body.model_name, listed_for) + try: + self._await_model_servable(body.model_name, listed_for) + except BaseException: + self.delete_model(model_id) + raise settle_propagation(written_at) return model_id diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index a508c97b9fb..2caac58333f 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -1,9 +1,11 @@ -"""Harness coverage for the model barrier that gates on every replica. +"""Harness coverage for the barriers that gate on every replica. No proxy needed and no ``e2e`` marker: this pins that a model registered through the control plane only counts as servable once every configured replica lists it -on /v1/models, which is what keeps a two-gateway stack from handing a test a -model that one gateway has not reloaded yet. The fakes are plain pollers and an +on /v1/models, and that a management write only counts as read back once every +replica's read satisfies the caller's predicate, which is what keeps a two-gateway +stack from handing a test a model or a key that one gateway has not caught up on +yet. The fakes are plain pollers standing in for each replica's transport plus an injected clock, so nothing here monkeypatches anything. """ @@ -12,18 +14,33 @@ from __future__ import annotations from collections.abc import Iterable, Mapping from dataclasses import dataclass from itertools import chain, repeat +from types import MappingProxyType from typing import Final import pytest from e2e_config import parse_replica_urls -from e2e_http import Success -from models import ModelListEntry, ModelsListResponse -from proxy_client import ModelsPoller, NotServableOn, Servable, await_servable_everywhere +from e2e_http import Result, Success +from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse +from proxy_client import ( + Poller, + ConvergeOutcome, + Converged, + ModelsPoller, + NotConverged, + NotServableOn, + Servable, + await_converged_everywhere, + await_servable_everywhere, + first_lagging_replica, + converge_timeout_message, +) MODEL: Final = "gpt-under-test" TIMEOUT: Final = 10.0 INTERVAL: Final = 2.0 +RPM_BEFORE_UPDATE: Final = 100 +RPM_AFTER_UPDATE: Final = 200 @dataclass @@ -78,6 +95,91 @@ class TestAwaitServableEverywhere: assert _await(pollers) == Servable() +def _key_info(rpm_limit: int) -> Success[KeyInfoResponse]: + return Success(status_code=200, data=KeyInfoResponse(info=KeyInfo(rpm_limit=rpm_limit))) + + +def _reads(results: Iterable[Result[KeyInfoResponse]]) -> Poller[Result[KeyInfoResponse]]: + it: Final = iter(results) + return lambda: next(it) + + +def _updated(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, Success) and result.data.info.rpm_limit == RPM_AFTER_UPDATE + + +def _converge( + pollers: Mapping[str, Poller[Result[KeyInfoResponse]]], clock: FakeClock +) -> Mapping[str, ConvergeOutcome[Result[KeyInfoResponse]]]: + return await_converged_everywhere( + pollers, + converged=_updated, + timeout=TIMEOUT, + interval=INTERVAL, + now=clock.now, + sleep=clock.sleep, + ) + + +class TestAwaitConvergedEverywhere: + def test_waits_for_the_replica_that_lags_behind_the_write(self) -> None: + clock: Final = FakeClock() + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))), + "gateway-2": _reads( + chain(repeat(_key_info(RPM_BEFORE_UPDATE), 2), repeat(_key_info(RPM_AFTER_UPDATE))) + ), + } + ) + outcomes: Final = _converge(pollers, clock) + assert outcomes == { + "gateway-1": Converged(result=_key_info(RPM_AFTER_UPDATE)), + "gateway-2": Converged(result=_key_info(RPM_AFTER_UPDATE)), + } + assert first_lagging_replica(outcomes) is None + assert clock.elapsed == 2 * INTERVAL + + def test_names_the_replica_that_never_converges_with_its_last_read(self) -> None: + clock: Final = FakeClock() + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))), + "gateway-2": _reads(repeat(_key_info(RPM_BEFORE_UPDATE))), + } + ) + outcomes: Final = _converge(pollers, clock) + assert first_lagging_replica(outcomes) == ( + "gateway-2", + NotConverged(last_result=_key_info(RPM_BEFORE_UPDATE)), + ) + assert clock.elapsed == TIMEOUT + message: Final = converge_timeout_message( + what="GET /key/info", + replica="gateway-2", + timeout=TIMEOUT, + last_result=_key_info(RPM_BEFORE_UPDATE), + ) + assert "gateway-2" in message and "/key/info" in message and str(RPM_BEFORE_UPDATE) in message + + def test_each_replica_gets_its_own_full_budget(self) -> None: + """A replica that converges late must not eat into the next replica's budget: both + need most of the timeout here, so one shared deadline would starve the second.""" + clock: Final = FakeClock() + slow: Final = chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE))) + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(slow), + "gateway-2": _reads( + chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE))) + ), + } + ) + outcomes: Final = _converge(pollers, clock) + assert first_lagging_replica(outcomes) is None + assert clock.elapsed == 2 * 3 * INTERVAL + + class TestParseReplicaUrls: def test_splits_and_trims_the_gateway_addresses(self) -> None: raw: Final = " http://127.0.0.1:4010/, http://127.0.0.1:4011 " diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 108e50b52a1..36fd744efc0 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8017,6 +8017,11 @@ export interface paths { * Update Key Fn * @description Update an existing API key's parameters. * + * The body is a merge patch: a field left out keeps its stored value, and on the key's own columns + * an explicit null clears it. The metadata-backed fields below are the exception, merging into the + * stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself + * replaces the stored metadata wholesale. + * * Parameters: * - key: Optional[str] - The key to update. Either key or key_alias must be provided. * - key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases)