From 213740863b9fd1987f189707cfe3dd7f91a78d2c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 20 Oct 2026 19:42:44 -0700 Subject: [PATCH] feat(gateway): forward declined and unserved requests to Python The gateway now acts on the core decline. A replay route_layer records each native request before its handler runs. When the error response carries a Decline it forwards the recording to Python as the client sent it, or keeps the 501 when the gateway runs standalone. Paths and methods with no native route go to Python through the merged router's fallbacks, before authentication LITELLM_PYTHON_UPSTREAM picks the Fallback at startup. Forward uses its own hyper client to the loopback proxy, strips hop-by-hop headers and adds x-litellm-gateway-declined with the reason. It appends the client address to X-Forwarded-For, because Python otherwise sees every forwarded request coming from the gateway. Error responses attach the decline in one place, so the messages handler now builds its Anthropic envelope through Error::anthropic_response instead of a tuple that dropped it. The 501 stubs for embeddings and completions are gone, since a matched stub would never reach Python --- litellm-rust/Cargo.lock | 4 + .../crates/gateway-inference/Cargo.toml | 5 +- .../gateway-inference/src/dispatch/forward.rs | 138 ++++++++++++++ .../gateway-inference/src/dispatch/mod.rs | 24 +++ .../gateway-inference/src/dispatch/replay.rs | 73 ++++++++ .../crates/gateway-inference/src/error.rs | 46 +++-- .../crates/gateway-inference/src/lib.rs | 22 +-- .../crates/gateway-inference/src/messages.rs | 2 +- .../crates/gateway-inference/src/request.rs | 7 +- .../gateway-inference/tests/dispatch.rs | 177 ++++++++++++++++++ .../crates/gateway-inference/tests/routes.rs | 10 +- .../gateway-inference/tests/support/mod.rs | 24 ++- litellm-rust/crates/gateway/Cargo.toml | 1 + litellm-rust/crates/gateway/README.md | 8 + litellm-rust/crates/gateway/src/lib.rs | 26 ++- litellm-rust/crates/gateway/src/main.rs | 10 +- litellm-rust/crates/gateway/src/settings.rs | 8 + litellm-rust/crates/gateway/tests/mcp.rs | 3 +- litellm-rust/crates/gateway/tests/server.rs | 75 +++++++- 19 files changed, 606 insertions(+), 57 deletions(-) create mode 100644 litellm-rust/crates/gateway-inference/src/dispatch/forward.rs create mode 100644 litellm-rust/crates/gateway-inference/src/dispatch/mod.rs create mode 100644 litellm-rust/crates/gateway-inference/src/dispatch/replay.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/dispatch.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 98f36607fff..4259ca1965b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3757,6 +3757,7 @@ dependencies = [ "tower-sessions-moka-store", "tracing", "uuid", + "wiremock", ] [[package]] @@ -3788,9 +3789,11 @@ dependencies = [ "base64 0.22.1", "bytes", "futures-util", + "hyper-util", "litellm-auth", "litellm-cache-memory", "litellm-cache-response", + "litellm-core-utils", "litellm-gateway-auth", "litellm-host", "litellm-host-http", @@ -3811,6 +3814,7 @@ dependencies = [ "thiserror 2.0.19", "tokio", "tower", + "tracing", "wiremock", ] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 312d44b9a4b..a316ff58b50 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -10,7 +10,10 @@ litellm-cache-response.workspace = true axum = { workspace = true, features = ["json", "multipart", "original-uri"] } base64.workspace = true bytes.workspace = true +futures-util.workspace = true +hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] } litellm-auth.workspace = true +litellm-core-utils.workspace = true litellm-gateway-auth.workspace = true litellm-inference.workspace = true litellm-inference-transcription.workspace = true @@ -28,10 +31,10 @@ litellm-llms-types.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true +tracing.workspace = true [dev-dependencies] litellm-cache-memory.workspace = true -futures-util.workspace = true tokio = { workspace = true, features = ["io-util"] } rstest.workspace = true tower = { version = "0.5.3", features = ["util"] } diff --git a/litellm-rust/crates/gateway-inference/src/dispatch/forward.rs b/litellm-rust/crates/gateway-inference/src/dispatch/forward.rs new file mode 100644 index 00000000000..dc24f19b4a0 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/dispatch/forward.rs @@ -0,0 +1,138 @@ +use std::{net::SocketAddr, sync::Arc}; + +use axum::{ + body::Body, + extract::{ConnectInfo, OriginalUri, Request}, + http::{ + HeaderMap, HeaderName, HeaderValue, StatusCode, Uri, Version, + header::{ + CONNECTION, PROXY_AUTHENTICATE, PROXY_AUTHORIZATION, TE, TRAILER, TRANSFER_ENCODING, + UPGRADE, + }, + uri::{Authority, PathAndQuery, Scheme}, + }, + response::{IntoResponse, Response}, +}; +use hyper_util::{ + client::legacy::{Client, connect::HttpConnector}, + rt::TokioExecutor, +}; +use litellm_core_utils::decline::Decline; + +/// Python reads it to log why the gateway handed the request over. +pub const DECLINED_HEADER: HeaderName = HeaderName::from_static("x-litellm-gateway-declined"); + +const X_FORWARDED_FOR: HeaderName = HeaderName::from_static("x-forwarded-for"); + +const HOP_BY_HOP: [HeaderName; 8] = [ + CONNECTION, + HeaderName::from_static("keep-alive"), + PROXY_AUTHENTICATE, + PROXY_AUTHORIZATION, + TE, + TRAILER, + TRANSFER_ENCODING, + UPGRADE, +]; + +#[derive(Debug, thiserror::Error)] +#[error("the Python upstream must be an http:// origin, got {0}")] +pub struct InvalidUpstream(String); + +/// Reverse proxy to the Python proxy on loopback, with its own client: provider clients carry +/// proxy and TLS settings that must not apply here. +#[derive(Clone)] +pub struct Forward { + client: Arc>, + authority: Authority, +} + +impl Forward { + pub fn new(upstream: &str) -> Result { + let invalid = || InvalidUpstream(upstream.to_owned()); + let uri: Uri = upstream.parse().map_err(|_| invalid())?; + let origin_only = uri.path_and_query().is_none_or(|path| path.as_str() == "/"); + let authority = match (uri.scheme(), uri.authority()) { + (Some(scheme), Some(authority)) if *scheme == Scheme::HTTP && origin_only => { + authority.clone() + } + _ => return Err(invalid()), + }; + Ok(Self { + client: Arc::new(Client::builder(TokioExecutor::new()).build(HttpConnector::new())), + authority, + }) + } + + pub async fn send(&self, request: Request, decline: Decline) -> Response { + let (mut parts, body) = request.into_parts(); + let path = parts + .extensions + .get::() + .map_or(&parts.uri, |original| &original.0) + .path_and_query() + .cloned() + .unwrap_or_else(|| PathAndQuery::from_static("/")); + parts.uri = Uri::from_parts({ + let mut uri = axum::http::uri::Parts::default(); + uri.scheme = Some(Scheme::HTTP); + uri.authority = Some(self.authority.clone()); + uri.path_and_query = Some(path); + uri + }) + .expect("scheme, authority and path make an absolute uri"); + parts.version = Version::HTTP_11; + strip_hop_by_hop(&mut parts.headers); + if let Some(ConnectInfo(client)) = parts.extensions.get::>() { + append_forwarded_for(&mut parts.headers, *client); + } + parts + .headers + .insert(DECLINED_HEADER, HeaderValue::from_static(decline.reason)); + tracing::info!( + reason = decline.reason, + path = parts.uri.path(), + "forwarded to python" + ); + match self.client.request(Request::from_parts(parts, body)).await { + Ok(response) => { + let (mut parts, body) = response.into_parts(); + strip_hop_by_hop(&mut parts.headers); + Response::from_parts(parts, Body::new(body)) + } + Err(error) => { + tracing::warn!(%error, "python upstream unreachable"); + StatusCode::BAD_GATEWAY.into_response() + } + } + } +} + +/// Python sees the gateway as its peer, so the client's address goes on the end of the chain +/// any proxy in front of the gateway already started. +fn append_forwarded_for(headers: &mut HeaderMap, client: SocketAddr) { + let client = client.ip().to_string(); + let chain = headers + .get_all(&X_FORWARDED_FOR) + .iter() + .filter_map(|value| value.to_str().ok()) + .chain([client.as_str()]) + .collect::>() + .join(", "); + if let Ok(chain) = HeaderValue::from_str(&chain) { + headers.insert(X_FORWARDED_FOR, chain); + } +} + +fn strip_hop_by_hop(headers: &mut HeaderMap) { + let named: Vec = headers + .get_all(CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .filter_map(|name| HeaderName::try_from(name.trim()).ok()) + .collect(); + for name in HOP_BY_HOP.iter().chain(&named) { + headers.remove(name); + } +} diff --git a/litellm-rust/crates/gateway-inference/src/dispatch/mod.rs b/litellm-rust/crates/gateway-inference/src/dispatch/mod.rs new file mode 100644 index 00000000000..845394da8f9 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/dispatch/mod.rs @@ -0,0 +1,24 @@ +//! What the gateway does with a request Rust core does not serve. +//! +//! A core decline travels up as an error and rides on the error response as a [`Decline`] +//! extension. [`replay`] is the only place that acts on it: it hands the request, as the +//! client sent it, to the [`Fallback`]. + +mod forward; +mod replay; + +pub use forward::{Forward, InvalidUpstream}; +pub(crate) use replay::replay; + +use litellm_core_utils::decline::Decline; + +/// The reason a request with no native route carries to Python. +pub const NOT_NATIVE: Decline = Decline::new("route not native"); + +#[derive(Clone)] +pub enum Fallback { + /// Standalone: an unserved path is a 404, a decline is the gateway's own 501. + Reject, + /// Next to the Python proxy: both go to Python unchanged. + Forward(Forward), +} diff --git a/litellm-rust/crates/gateway-inference/src/dispatch/replay.rs b/litellm-rust/crates/gateway-inference/src/dispatch/replay.rs new file mode 100644 index 00000000000..5c4ff545013 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/src/dispatch/replay.rs @@ -0,0 +1,73 @@ +use std::net::SocketAddr; + +use axum::{ + body::{Body, Bytes}, + extract::{ConnectInfo, OriginalUri, Request, State}, + http::request::Parts, + middleware::Next, + response::Response, +}; +use bytes::BytesMut; +use futures_util::StreamExt; +use litellm_core_utils::decline::Decline; + +use super::Fallback; +use crate::request::MAX_BODY_BYTES; + +/// Records the request before any handler reads it, and forwards the recording when the +/// response says the request was declined. +pub(crate) async fn replay( + State(fallback): State, + request: Request, + next: Next, +) -> Response { + let (parts, body) = request.into_parts(); + let body = match buffer(body).await { + Ok(body) => body, + Err(body) => return next.run(Request::from_parts(parts, body)).await, + }; + let recorded = record(&parts, body.clone()); + let response = next.run(Request::from_parts(parts, Body::from(body))).await; + let Some(&decline) = response.extensions().get::() else { + return response; + }; + match fallback { + Fallback::Reject => response, + Fallback::Forward(forward) => forward.send(recorded.map(Body::from), decline).await, + } +} + +/// The whole body, or, when it is too large or unreadable, a body that yields the same bytes +/// and the same failure, so the handler reports it in its own error envelope. +async fn buffer(body: Body) -> Result { + let mut stream = body.into_data_stream(); + let mut seen = BytesMut::new(); + while let Some(chunk) = stream.next().await { + match chunk { + Ok(chunk) if seen.len() + chunk.len() <= MAX_BODY_BYTES => { + seen.extend_from_slice(&chunk); + } + failed => { + let replayed = futures_util::stream::iter([Ok(seen.freeze()), failed]); + return Err(Body::from_stream(replayed.chain(stream))); + } + } + } + Ok(seen.freeze()) +} + +fn record(parts: &Parts, body: Bytes) -> Request { + let uri = parts + .extensions + .get::() + .map_or_else(|| parts.uri.clone(), |original| original.0.clone()); + let mut recorded = Request::new(body); + *recorded.method_mut() = parts.method.clone(); + *recorded.uri_mut() = uri; + *recorded.version_mut() = parts.version; + *recorded.headers_mut() = parts.headers.clone(); + if let Some(&client) = parts.extensions.get::>() { + recorded.extensions_mut().insert(client); + } + recorded +} diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index e1470061190..9771c434be3 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -3,6 +3,7 @@ use axum::{ Json, response::{IntoResponse, Response}, }; +use litellm_core_utils::decline::Decline; use litellm_http::transport::Error as TransportError; use litellm_inference::Error as InferenceError; use litellm_llms::base_llm::ocr::error::Error as OcrError; @@ -22,8 +23,6 @@ pub enum Error { Route(#[from] InferenceError), #[error(transparent)] Ocr(#[from] OcrError), - #[error("{0} is not implemented by the Rust gateway")] - Unsupported(String), #[error("request body exceeds the size limit")] BodyTooLarge, #[error("{0}")] @@ -49,7 +48,6 @@ impl Error { pub fn status(&self) -> StatusCode { match self { Self::Auth(error) => error.status(), - Self::Unsupported(_) => StatusCode::NOT_IMPLEMENTED, Self::Route(error) if error.declined().is_some() => StatusCode::NOT_IMPLEMENTED, Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE, Self::Ocr( @@ -74,22 +72,44 @@ impl Error { } } + pub fn declined(&self) -> Option<&Decline> { + match self { + Self::Route(error) => error.declined(), + Self::Auth(_) + | Self::InvalidBody(_) + | Self::UnknownModel(_) + | Self::Ocr(_) + | Self::BodyTooLarge + | Self::Internal(_) => None, + } + } + pub fn openai_response(self) -> Response { let status = self.status(); let message = match &self { Self::UnknownModel(model) => format!("Invalid model name passed in model={model}"), _ => self.to_string(), }; - ( - status, - Json(json!({"error": { - "message": message, - "type": error_type(status), - "param": null, - "code": status.as_u16(), - }})), - ) - .into_response() + self.respond(json!({"error": { + "message": message, + "type": error_type(status), + "param": null, + "code": status.as_u16(), + }})) + } + + pub fn anthropic_response(self, request_id: Option<&str>) -> Response { + let body = self.body(request_id); + self.respond(body) + } + + /// A decline rides on the response, so the dispatch layer can hand the request on. + fn respond(&self, body: Value) -> Response { + let mut response = (self.status(), Json(body)).into_response(); + if let Some(decline) = self.declined() { + response.extensions_mut().insert(*decline); + } + response } /// The Anthropic error envelope Python's `AnthropicExceptionMapping` builds: an upstream diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 4663151f1de..b980ea2d186 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -6,6 +6,7 @@ mod audio_transcription; mod caching; mod chat_completions; +pub mod dispatch; mod error; pub mod messages; mod ocr; @@ -88,7 +89,7 @@ impl Gateway { } } -pub fn router(gateway: Arc) -> Router { +pub fn router(gateway: Arc, fallback: dispatch::Fallback) -> Router { Router::new() .route("/v1/messages", post(messages::create)) .route("/ocr", post(ocr::create)) @@ -104,10 +105,10 @@ pub fn router(gateway: Arc) -> Router { ) .route("/responses", post(responses::create)) .route("/v1/responses", post(responses::create)) - .route("/embeddings", post(request::unsupported)) - .route("/v1/embeddings", post(request::unsupported)) - .route("/completions", post(request::unsupported)) - .route("/v1/completions", post(request::unsupported)) + .route_layer(axum::middleware::from_fn_with_state( + fallback, + dispatch::replay, + )) .layer(axum::extract::DefaultBodyLimit::max( request::MAX_BODY_BYTES, )) @@ -115,11 +116,8 @@ pub fn router(gateway: Arc) -> Router { } fn model_routes() -> Router> { - Router::new() - .route( - "/chat/completions", - post(chat_completions::create_from_model_path), - ) - .route("/embeddings", post(request::unsupported)) - .route("/completions", post(request::unsupported)) + Router::new().route( + "/chat/completions", + post(chat_completions::create_from_model_path), + ) } diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index d612a042e03..bf2ed6469a2 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -37,7 +37,7 @@ pub async fn create( Ok(JsonObject(body)) => handle(&gateway, &identity, &headers, body).await, Err(error) => Err(error), }; - result.map_err(|error| (error.status(), Json(error.body(request_id.as_deref())))) + result.map_err(|error| error.anthropic_response(request_id.as_deref())) } async fn handle( diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs index b398da2013b..293c37ec551 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -1,8 +1,7 @@ use axum::{ body::Bytes, - extract::{FromRequest, FromRequestParts, Multipart, OriginalUri, Request}, + extract::{FromRequest, FromRequestParts, Multipart, Request}, http::request::Parts, - response::IntoResponse, }; use serde_json::{Map, Value}; @@ -162,10 +161,6 @@ fn multipart_error(error: axum::extract::multipart::MultipartError) -> Error { Error::InvalidBody(error.to_string()) } -pub(crate) async fn unsupported(OriginalUri(uri): OriginalUri) -> impl IntoResponse { - Error::Unsupported(uri.path().to_owned()) -} - #[cfg(test)] mod tests { use axum::{ diff --git a/litellm-rust/crates/gateway-inference/tests/dispatch.rs b/litellm-rust/crates/gateway-inference/tests/dispatch.rs new file mode 100644 index 00000000000..b546a1a89a6 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/dispatch.rs @@ -0,0 +1,177 @@ +mod support; + +use std::net::SocketAddr; + +use axum::{ + body::{Body, to_bytes}, + extract::ConnectInfo, + http::Request, +}; +use rstest::rstest; +use serde_json::json; +use tower::ServiceExt; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; + +async fn python() -> MockServer { + let python = MockServer::start().await; + Mock::given(any()) + .respond_with( + ResponseTemplate::new(200) + .insert_header("x-served-by", "python") + .set_body_string("from python"), + ) + .mount(&python) + .await; + python +} + +async fn provider() -> MockServer { + let provider = MockServer::start().await; + Mock::given(any()) + .respond_with(ResponseTemplate::new(500).set_body_string("provider failed")) + .mount(&provider) + .await; + provider +} + +#[rstest] +#[case::chat_streaming( + "anthropic/test-model", + "/v1/chat/completions?trace=1", + r#"{"model": "public/model", "messages": [{"role": "user", "content": "hi"}], "stream": true}"#, + "streaming" +)] +#[case::messages_bedrock_invoke( + "bedrock/invoke/anthropic.claude-3-5-sonnet-20240620-v1:0", + "/v1/messages", + r#"{"model":"public/model","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}"#, + "Bedrock invoke messages request shaping" +)] +#[case::responses_provider( + "gemini/test-model", + "/v1/responses", + r#"{"model":"public/model","input":"hi"}"#, + "native HTTP responses provider" +)] +#[tokio::test] +async fn a_core_decline_reaches_python_as_the_client_sent_it( + #[case] model: &str, + #[case] path: &str, + #[case] body: &'static str, + #[case] reason: &str, +) { + let python = python().await; + let provider = provider().await; + let app = support::app_forwarding_to(&python.uri(), model, &provider.uri()); + + let response = app + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer client-key") + .header("x-client-header", "kept") + .header("connection", "te") + .header("te", "trailers") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), 200); + assert_eq!(response.headers()["x-served-by"], "python"); + assert_eq!( + to_bytes(response.into_body(), 1024).await.unwrap(), + "from python" + ); + assert!(provider.received_requests().await.unwrap().is_empty()); + let forwarded = python.received_requests().await.unwrap(); + assert_eq!(forwarded.len(), 1); + let forwarded = &forwarded[0]; + assert_eq!(forwarded.method.as_str(), "POST"); + let forwarded_path = match forwarded.url.query() { + Some(query) => format!("{}?{query}", forwarded.url.path()), + None => forwarded.url.path().to_owned(), + }; + assert_eq!(forwarded_path, path); + assert_eq!(forwarded.body, body.as_bytes()); + assert_eq!(forwarded.headers["authorization"], "Bearer client-key"); + assert_eq!(forwarded.headers["x-client-header"], "kept"); + assert_eq!(forwarded.headers["x-litellm-gateway-declined"], reason); + assert!(!forwarded.headers.contains_key("te")); +} + +#[rstest] +#[tokio::test] +async fn a_standalone_gateway_answers_a_decline_with_its_own_501() { + let provider = provider().await; + let response = support::post( + support::app("gemini/test-model", &provider.uri()), + "/v1/responses", + json!({"model": "public/model", "input": "hi"}), + ) + .await; + + assert_eq!(response.status(), 501); + let body = support::json(response).await; + assert!( + body["error"]["message"] + .as_str() + .unwrap() + .contains("native HTTP responses provider") + ); + assert!(provider.received_requests().await.unwrap().is_empty()); +} + +#[rstest] +#[case::invalid_request(json!({"model": "public/model", "messages": "hi"}), 400)] +#[case::unknown_model(json!({"model": "other", "messages": [{"role": "user", "content": "hi"}]}), 400)] +#[case::provider_failure(json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), 500)] +#[tokio::test] +async fn a_terminal_error_is_never_forwarded(#[case] body: serde_json::Value, #[case] status: u16) { + let python = python().await; + let provider = provider().await; + let response = support::post( + support::app_forwarding_to(&python.uri(), "anthropic/test-model", &provider.uri()), + "/v1/chat/completions", + body, + ) + .await; + + assert_eq!(response.status(), status); + assert!(python.received_requests().await.unwrap().is_empty()); +} + +#[rstest] +#[case::direct_client(None, "203.0.113.7")] +#[case::behind_a_load_balancer(Some("198.51.100.1"), "198.51.100.1, 203.0.113.7")] +#[tokio::test] +async fn a_forwarded_decline_names_the_client_in_x_forwarded_for( + #[case] incoming: Option<&str>, + #[case] expected: &str, +) { + let python = python().await; + let provider = provider().await; + let request = Request::post("/v1/chat/completions").header("content-type", "application/json"); + let request = match incoming { + Some(chain) => request.header("x-forwarded-for", chain), + None => request, + }; + let mut request = request + .body(Body::from( + json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], "stream": true}) + .to_string(), + )) + .unwrap(); + request.extensions_mut().insert(ConnectInfo( + "203.0.113.7:50000".parse::().unwrap(), + )); + + support::app_forwarding_to(&python.uri(), "anthropic/test-model", &provider.uri()) + .oneshot(request) + .await + .unwrap(); + + let forwarded = python.received_requests().await.unwrap(); + assert_eq!(forwarded[0].headers["x-forwarded-for"], expected); +} diff --git a/litellm-rust/crates/gateway-inference/tests/routes.rs b/litellm-rust/crates/gateway-inference/tests/routes.rs index b0699ff10b5..a585a9f5700 100644 --- a/litellm-rust/crates/gateway-inference/tests/routes.rs +++ b/litellm-rust/crates/gateway-inference/tests/routes.rs @@ -157,18 +157,14 @@ async fn missing_inference_fields_use_the_same_validation_as_null( #[case("/engines/public%2Fmodel/embeddings")] #[case("/openai/deployments/public%2Fmodel/completions")] #[tokio::test] -async fn unimplemented_routes_return_an_explicit_error(#[case] path: &str) { +async fn routes_without_a_native_handler_are_left_to_the_mounting_server(#[case] path: &str) { let response = support::post( support::app("anthropic/test-model", "http://127.0.0.1:1"), path, json!({}), ) .await; - assert_eq!(response.status(), 501); - let body = support::json(response).await; - let message = body["error"]["message"].as_str().unwrap(); - assert!(message.contains("not implemented")); - assert!(message.contains(path)); + assert_eq!(response.status(), 404); } #[rstest] @@ -285,7 +281,7 @@ async fn json_extraction_rejections_use_the_endpoint_error_envelope( } #[rstest] -#[case::unsupported("/engines/public%2Fmodel/embeddings", 501)] +#[case::unserved("/engines/public%2Fmodel/embeddings", 404)] #[case::unknown("/engines/public%2Fmodel/unknown", 404)] #[case::unescaped_model("/engines/public/model/chat/completions", 404)] #[case::extra_segment("/engines/public%2Fmodel/extra/chat/completions", 404)] diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 0e37bb0aff9..5366cbd4118 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -10,7 +10,11 @@ use axum::{ response::Response, }; use futures_util::future::BoxFuture; -use litellm_gateway_inference::{Deployment, Gateway, router}; +use litellm_gateway_inference::{ + Deployment, Gateway, + dispatch::{Fallback, Forward}, + router, +}; use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; use litellm_inference::resources::CoreResources; use litellm_secrets::{SecretValue, source::SecretSource}; @@ -37,7 +41,18 @@ pub fn app_with_permissions( api_base: &str, permissions: litellm_gateway_auth::Permissions, ) -> Router { - configured_app(model, api_base, permissions, None, None) + configured_app(model, api_base, permissions, None, None, Fallback::Reject) +} + +pub fn app_forwarding_to(python: &str, model: &str, api_base: &str) -> Router { + configured_app( + model, + api_base, + litellm_gateway_auth::Permissions::All, + None, + None, + Fallback::Forward(Forward::new(python).unwrap()), + ) } pub fn app_with_cache( @@ -51,6 +66,7 @@ pub fn app_with_cache( litellm_gateway_auth::Permissions::All, Some(cache), None, + Fallback::Reject, ) } @@ -66,6 +82,7 @@ pub fn app_with_cache_for_principal( litellm_gateway_auth::Permissions::All, Some(cache), Some(principal), + Fallback::Reject, ) } @@ -75,6 +92,7 @@ fn configured_app( permissions: litellm_gateway_auth::Permissions, cache: Option>, principal: Option, + fallback: Fallback, ) -> Router { let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); let http = Resolution::from(&HttpSettings::default()).config; @@ -102,7 +120,7 @@ fn configured_app( Some(cache) => gateway.with_cache(cache), None => gateway, }; - router(Arc::new(gateway)).layer(axum::middleware::from_fn_with_state( + router(Arc::new(gateway), fallback).layer(axum::middleware::from_fn_with_state( (permissions, principal), test_identity, )) diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 387bcd43231..e8edf2b114a 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -31,6 +31,7 @@ uuid.workspace = true tower-sessions-moka-store = "0.15.0" [dev-dependencies] +wiremock.workspace = true base64.workspace = true rstest.workspace = true tempfile.workspace = true diff --git a/litellm-rust/crates/gateway/README.md b/litellm-rust/crates/gateway/README.md index 125c2b07448..2b770ffacd1 100644 --- a/litellm-rust/crates/gateway/README.md +++ b/litellm-rust/crates/gateway/README.md @@ -41,3 +41,11 @@ Each MCP entry supports HTTP or stdio transport, an optional pinned `server_id`, The gateway initializes upstreams before opening its listener. Startup fails when credentials cannot be resolved or an upstream cannot initialize. Shutdown cancels upstream connections and terminates owned stdio children The config schema preserves Python sections, but parsing a section does not implement its service. MCP OAuth, database-backed server management, guardrails, spend accounting, local Python tools, and legacy SSE upstream connections remain unimplemented. Unsupported MCP settings and configured guardrail policies fail at startup. Incoming legacy SSE clients remain supported + +## Python fallback + +Set `LITELLM_PYTHON_UPSTREAM` to the Python proxy's `http://` origin, usually a loopback address, and the gateway forwards every request it does not serve natively. That covers paths and methods with no native route, which reach Python without gateway authentication, and native requests that Rust core declines before calling the provider. A forwarded request keeps its method, path, query, headers and body bytes, and carries `x-litellm-gateway-declined` with the reason. Errors after the provider call are never forwarded + +The gateway appends the client address to `X-Forwarded-For`, since Python otherwise sees every forwarded request coming from the gateway. Set `general_settings.use_x_forwarded_for: true` on the Python proxy so `allowed_ips`, spend log IPs and the other IP checks read it, and include the gateway's address in its trusted proxy ranges + +Without it the gateway runs standalone. An unserved path is a 404 and a decline is a 501 that names the reason diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 0b3aba81181..44de2ff9570 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -16,7 +16,10 @@ use http_body_util::BodyExt; use litellm_config::Config; use litellm_gateway_auth::Auth; -use litellm_gateway_inference::{Gateway, ModelRouter}; +use litellm_gateway_inference::{ + Gateway, ModelRouter, + dispatch::{Fallback, NOT_NATIVE}, +}; use litellm_http::{ ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify, media::PublicDnsResolver, @@ -118,19 +121,30 @@ pub fn router( config: &Config, ui: Option, mcp: Option, + fallback: Fallback, ) -> Router { let auth = Auth::from_config(config, inference.secrets.clone()); - let inference = litellm_gateway_inference::router(inference) + let served = litellm_gateway_inference::router(inference, fallback.clone()) .merge(mcp.unwrap_or_default()) .route_layer(axum::middleware::from_fn(auth::bind_session_owner)) .route_layer(axum::middleware::from_fn_with_state( auth, litellm_gateway_auth::authenticate, )) - .layer(axum::middleware::from_fn(log_request)); - match ui { - Some(ui) => inference.merge(ui), - None => inference, + .layer(axum::middleware::from_fn(log_request)) + .merge(ui.unwrap_or_default()); + match fallback { + Fallback::Reject => served, + Fallback::Forward(forward) => { + let unmatched = forward.clone(); + served + .fallback(move |request: Request| async move { + unmatched.send(request, NOT_NATIVE).await + }) + .method_not_allowed_fallback(move |request: Request| async move { + forward.send(request, NOT_NATIVE).await + }) + } } } diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index f8150e6e860..81fc214bb4b 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -5,6 +5,7 @@ use std::{ use litellm_config::Config; use litellm_gateway_auth::UiBackend; +use litellm_gateway_inference::dispatch::{Fallback, Forward}; use litellm_tracing::{Level, Logger, Metadata, Record, Sink}; use serde_json::json; @@ -46,6 +47,10 @@ async fn main() -> Result<(), Box> { Logger::new(StderrSink { level }).install_global()?; let config = Config::load(settings.litellm_config)?; let inference = litellm_gateway::build_inference(&config)?; + let fallback = match settings.litellm_python_upstream.as_deref() { + Some(upstream) => Fallback::Forward(Forward::new(upstream)?), + None => Fallback::Reject, + }; let ui = match std::env::var_os("LITELLM_UI_PATH") { Some(directory) => { let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?; @@ -97,11 +102,12 @@ async fn main() -> Result<(), Box> { }); let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; - tracing::info!(address = %listener.local_addr()?, models = config.model_list.len(), log_level = %level, "gateway listening"); + tracing::info!(address = %listener.local_addr()?, models = config.model_list.len(), log_level = %level, python_upstream = settings.litellm_python_upstream.as_deref(), "gateway listening"); let result = axum::serve( listener, - litellm_gateway::router(inference, &config, ui, mcp_router), + litellm_gateway::router(inference, &config, ui, mcp_router, fallback) + .into_make_service_with_connect_info::(), ) .with_graceful_shutdown(shutdown.clone().cancelled_owned()) .await; diff --git a/litellm-rust/crates/gateway/src/settings.rs b/litellm-rust/crates/gateway/src/settings.rs index 2dd8872065f..7da04d13b47 100644 --- a/litellm-rust/crates/gateway/src/settings.rs +++ b/litellm-rust/crates/gateway/src/settings.rs @@ -10,6 +10,7 @@ pub(super) struct Settings { pub host: String, pub port: u16, pub litellm_config: String, + pub litellm_python_upstream: Option, rust_log: Option, } @@ -19,6 +20,7 @@ impl Default for Settings { host: "0.0.0.0".into(), port: 4000, litellm_config: "config.yaml".into(), + litellm_python_upstream: None, rust_log: None, } } @@ -81,6 +83,7 @@ mod tests { assert_eq!(settings.host, "0.0.0.0"); assert_eq!(settings.port, 4000); assert_eq!(settings.litellm_config, "config.yaml"); + assert_eq!(settings.litellm_python_upstream, None); assert_eq!(settings.log_level(), Level::INFO); } @@ -90,6 +93,7 @@ mod tests { ("HOST", "127.0.0.1"), ("PORT", "8080"), ("LITELLM_CONFIG", "/tmp/custom.yaml"), + ("LITELLM_PYTHON_UPSTREAM", "http://127.0.0.1:4001"), ("RUST_LOG", "debug"), ("LITELLM_UI_SECURE_COOKIES", "invalid-but-ui-disabled"), ]) @@ -98,6 +102,10 @@ mod tests { assert_eq!(settings.host, "127.0.0.1"); assert_eq!(settings.port, 8080); assert_eq!(settings.litellm_config, "/tmp/custom.yaml"); + assert_eq!( + settings.litellm_python_upstream.as_deref(), + Some("http://127.0.0.1:4001") + ); assert_eq!(settings.log_level(), Level::DEBUG); } diff --git a/litellm-rust/crates/gateway/tests/mcp.rs b/litellm-rust/crates/gateway/tests/mcp.rs index 593fbfeb129..a017a440efc 100644 --- a/litellm-rust/crates/gateway/tests/mcp.rs +++ b/litellm-rust/crates/gateway/tests/mcp.rs @@ -6,6 +6,7 @@ use axum::{ http::Request, }; use litellm_config::Config; +use litellm_gateway_inference::dispatch::Fallback; use litellm_gateway_mcp::{ HttpConfig, rmcp::{ @@ -145,7 +146,7 @@ mcp_servers: ..Default::default() }, ); - let app = litellm_gateway::router(inference, &config, None, Some(mcp_router)); + let app = litellm_gateway::router(inference, &config, None, Some(mcp_router), Fallback::Reject); for (path, method) in [ ("/mcp", "POST"), diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 71edd3c3650..24093640f5e 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -6,7 +6,10 @@ use std::{ use axum::{body::Body, http::Request}; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; use litellm_config::Config; -use litellm_gateway_inference::{Error, Gateway}; +use litellm_gateway_inference::{ + Error, Gateway, + dispatch::{Fallback, Forward}, +}; use litellm_http::ClientVariant; use litellm_tracing::{Logger, Metadata, Record, Sink}; use rstest::{fixture, rstest}; @@ -74,7 +77,7 @@ async fn authenticates_before_serving_mounted_inference_routes( let server = tokio::spawn(async move { axum::serve( listener, - litellm_gateway::router(inference, &config, None, None), + litellm_gateway::router(inference, &config, None, None, Fallback::Reject), ) .with_graceful_shutdown(async move { let _ = stopped.await; @@ -129,7 +132,10 @@ async fn logs_request_outcome_without_credentials_or_query(inference: Arc, + #[case] method: &str, + #[case] path: &str, + #[case] authorization: Option<&str>, + #[case] forwarded: bool, +) { + let python = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::any()) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_string("from python")) + .mount(&python) + .await; + let config = + Config::from_yaml("model_list: []\ngeneral_settings:\n master_key: gateway-key\n") + .unwrap(); + let fallback = Fallback::Forward(Forward::new(&python.uri()).unwrap()); + let request = Request::builder().method(method).uri(path); + let request = match authorization { + Some(value) => request.header("authorization", value), + None => request, + }; + + let mut request = request.body(Body::from("{}")).unwrap(); + request.extensions_mut().insert(axum::extract::ConnectInfo( + "203.0.113.7:50000".parse::().unwrap(), + )); + + let response = litellm_gateway::router(inference, &config, None, None, fallback) + .oneshot(request) + .await + .unwrap(); + + let received = python.received_requests().await.unwrap(); + assert_eq!(received.len(), usize::from(forwarded)); + if forwarded { + assert_eq!(response.status().as_u16(), 200); + assert_eq!(received[0].method.as_str(), method); + let url = &received[0].url; + let forwarded_path = match url.query() { + Some(query) => format!("{}?{query}", url.path()), + None => url.path().to_owned(), + }; + assert_eq!(forwarded_path, path); + assert_eq!( + received[0].headers["x-litellm-gateway-declined"], + "route not native" + ); + assert_eq!(received[0].headers["x-forwarded-for"], "203.0.113.7"); + } else { + assert_eq!(response.status().as_u16(), 401); + } +}