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] 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())); + } +}