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.
This commit is contained in:
mateo-berri 2026-09-03 02:15:57 -07:00
parent 658f50663d
commit c79d1d12ae
7 changed files with 84 additions and 7 deletions

View file

@ -1415,6 +1415,7 @@ dependencies = [
"litellm-core",
"pyo3",
"reqwest",
"rustls 0.23.42",
"serde",
"serde_json",
"sha2 0.10.9",

View file

@ -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"

View file

@ -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

View file

@ -3,3 +3,4 @@ pub mod ocr;
pub mod realtime;
pub mod realtime_pool;
pub mod responses_ws;
pub(crate) mod tls;

View file

@ -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)

View file

@ -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(),

View file

@ -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<R>(
request: R,
) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), Box<Error>>
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()));
}
}