mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
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:
parent
658f50663d
commit
c79d1d12ae
7 changed files with 84 additions and 7 deletions
1
litellm-rust/Cargo.lock
generated
1
litellm-rust/Cargo.lock
generated
|
|
@ -1415,6 +1415,7 @@ dependencies = [
|
|||
"litellm-core",
|
||||
"pyo3",
|
||||
"reqwest",
|
||||
"rustls 0.23.42",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -3,3 +3,4 @@ pub mod ocr;
|
|||
pub mod realtime;
|
||||
pub mod realtime_pool;
|
||||
pub mod responses_ws;
|
||||
pub(crate) mod tls;
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
67
litellm-rust/crates/ai-gateway/src/io/tls.rs
Normal file
67
litellm-rust/crates/ai-gateway/src/io/tls.rs
Normal 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()));
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue