From 8893c75eb365a2cdb6a591d8b2cb7c39fef07da8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 23 Oct 2026 11:13:59 -0700 Subject: [PATCH] feat(gateway): forward declined and unserved requests to Python The gateway acts on the core decline. A replay 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. The 501 stubs for embeddings and completions and the gateway's own Unsupported variant are gone, since a matched stub would never reach Python. LITELLM_PYTHON_UPSTREAM picks the Fallback at startup. Forward uses its own hyper client to the loopback proxy, strips hop-by-hop headers, adds x-litellm-gateway-declined with the reason and appends the client address to X-Forwarded-For. An upstream that is not an http origin is rejected with InvalidUpstream, which keeps the URI parse failure as its source. Error responses attach the decline in one place, so the messages handler builds its Anthropic envelope through Error::anthropic_response instead of a tuple that dropped it. Forwarding is a gateway concern, so the dispatch module lives in the gateway crate and replays declines for every mounted router. The replay layer is installed only when a Python upstream is set, so a standalone gateway does not buffer request bodies. Tests run the gateway and a stand-in Python on real listeners through litellm_gateway::serve, which main.rs uses too. The rust-error skill says where an unserved path and a forwarded decline go. Co-Authored-By: Claude Opus 5.5 --- .../rust-error/references/boundaries.md | 6 +- .../rust-error/references/classification.md | 2 +- .../skills/rust-error/references/gateway.md | 5 +- litellm-rust/Cargo.lock | 6 + .../crates/gateway-inference/Cargo.toml | 1 + .../crates/gateway-inference/src/error.rs | 34 +- .../crates/gateway-inference/src/lib.rs | 17 +- .../crates/gateway-inference/src/messages.rs | 2 +- .../crates/gateway-inference/src/request.rs | 9 +- .../crates/gateway-inference/tests/decline.rs | 72 ++++ .../crates/gateway-inference/tests/routes.rs | 10 +- litellm-rust/crates/gateway/Cargo.toml | 8 +- litellm-rust/crates/gateway/README.md | 8 + .../crates/gateway/src/dispatch/forward.rs | 179 +++++++++ .../crates/gateway/src/dispatch/mod.rs | 52 +++ .../crates/gateway/src/dispatch/replay.rs | 78 ++++ litellm-rust/crates/gateway/src/lib.rs | 39 +- litellm-rust/crates/gateway/src/main.rs | 13 +- litellm-rust/crates/gateway/src/settings.rs | 8 + litellm-rust/crates/gateway/tests/dispatch.rs | 363 ++++++++++++++++++ litellm-rust/crates/gateway/tests/mcp.rs | 3 +- litellm-rust/crates/gateway/tests/server.rs | 71 +++- 22 files changed, 925 insertions(+), 61 deletions(-) create mode 100644 litellm-rust/crates/gateway-inference/tests/decline.rs create mode 100644 litellm-rust/crates/gateway/src/dispatch/forward.rs create mode 100644 litellm-rust/crates/gateway/src/dispatch/mod.rs create mode 100644 litellm-rust/crates/gateway/src/dispatch/replay.rs create mode 100644 litellm-rust/crates/gateway/tests/dispatch.rs diff --git a/litellm-rust/.agents/skills/rust-error/references/boundaries.md b/litellm-rust/.agents/skills/rust-error/references/boundaries.md index d68da416363..2f0a108dcd0 100644 --- a/litellm-rust/.agents/skills/rust-error/references/boundaries.md +++ b/litellm-rust/.agents/skills/rust-error/references/boundaries.md @@ -16,7 +16,7 @@ The inference gateway and the Python bridge's route mapping both receive a neste | Predicate | Gateway | Bridge | |---|---|---| -| `declined()` | 501 | `RustBridgeDeclined`, without the route's public mapping | +| `declined()` | 501, or the request forwarded to Python when the gateway runs next to the proxy | `RustBridgeDeclined`, without the route's public mapping | | `upstream_response()` | the provider's status and body | `RustUpstreamError` | | `is_missing_credential()` | 401 | no class of its own, a `RuntimeError` because `is_request()` answers false | | `is_request()` | 400 | `ValueError` | @@ -24,11 +24,11 @@ The inference gateway and the Python bridge's route mapping both receive a neste If you find yourself adding an arm for your variant at a boundary, check whether a predicate is missing first -The two boundaries are not identical. The gateway has variants of its own for body parsing, an unknown model and an unsupported path, and the OCR bridge has Python parity arms such as `FileNotFoundError` for a missing file. Those arms match the boundary's own variants or the top level of a public enum, never a nested error's variants. A missing API key is 401 at the gateway and `RuntimeError` at both bridges, pinned by tests to what Python's own path does +The two boundaries are not identical. The gateway has variants of its own for body parsing, an invalid path parameter and an unknown model, and the OCR bridge has Python parity arms such as `FileNotFoundError` for a missing file. Those arms match the boundary's own variants or the top level of a public enum, never a nested error's variants. A missing API key is 401 at the gateway and `RuntimeError` at both bridges, pinned by tests to what Python's own path does ## Declines -A failure before anything reached the provider is a decline, and `declined()` on the route error names its reason. It means another host may serve the same request as it was received, so each boundary hands it to the host unchanged: the gateway answers 501 and the Python bridge raises `RustBridgeDeclined`, never a public exception or a `ValueError`. See [gateway](gateway.md) and [pyo3](pyo3.md) for each host's side +A failure before anything reached the provider is a decline, and `declined()` on the route error names its reason. It means another host may serve the same request as it was received, so each boundary hands it to the host unchanged: the gateway forwards the recorded request to Python, or answers 501 when it runs standalone, and the Python bridge raises `RustBridgeDeclined`, never a public exception or a `ValueError`. See [gateway](gateway.md) and [pyo3](pyo3.md) for each host's side Anything after the send is terminal, never a decline. `declined()` answers `None` for a response fault and for a post-call hook failure even when the wrapped error was a decline, and every route that transforms the provider's answer collapses its errors with `Error::after_send`, which keeps a provider answer and a response fault as they are and nests anything else as the source of an invalid response, so a decline can never describe a call that was billed. Ask `declined()`, do not match the variant behind it, because that after-send cut lives in the predicate diff --git a/litellm-rust/.agents/skills/rust-error/references/classification.md b/litellm-rust/.agents/skills/rust-error/references/classification.md index e32463847db..341a84a13c8 100644 --- a/litellm-rust/.agents/skills/rust-error/references/classification.md +++ b/litellm-rust/.agents/skills/rust-error/references/classification.md @@ -39,7 +39,7 @@ pub fn is_request(&self) -> bool { Add a predicate when a real caller needs a new distinction, and name it after the fact the caller needs, in the caller's words, not after the variant's spelling. An upper type's predicates call through to the nested error, as in `Self::Http(error) => error.is_request()` and `Self::Auth(error) => error.is_missing_credential()` -A public enum may be matched one level down by the type that owns the variant, such as the gateway matching its own `Unsupported` or the OCR bridge matching `FileRead`. Going a second level down, into a variant of a variant, is what to avoid, the gateway used to do that for `MissingApiKey` and the old `Unsupported` route variant +A public enum may be matched one level down by the type that owns the variant, such as the gateway matching its own `UnknownModel` or the OCR bridge matching `FileRead`. Going a second level down, into a variant of a variant, is what to avoid, the gateway used to do that for `MissingApiKey` and the old `Unsupported` route variant ## Facts, not policy diff --git a/litellm-rust/.agents/skills/rust-error/references/gateway.md b/litellm-rust/.agents/skills/rust-error/references/gateway.md index 5da34aa84b2..2c31767b078 100644 --- a/litellm-rust/.agents/skills/rust-error/references/gateway.md +++ b/litellm-rust/.agents/skills/rust-error/references/gateway.md @@ -21,7 +21,7 @@ Each hop is one `From` that nests the lower error in one variant, see [conversio | Answer | Status | |---|---| -| `declined()` | 501, for a route error, the gateway matched and parsed the request but core cannot serve it, nothing reached the provider | +| `declined()` | 501, for a route error, the gateway matched and parsed the request but core cannot serve it, nothing reached the provider. The response carries the `Decline` as an extension, so the dispatch layer can forward the request to Python instead | | `upstream_response()` | the provider's own status, 502 when it is not a valid code | | `is_missing_credential()` | 401 | | `is_request()` | 400, for a route error | @@ -46,7 +46,8 @@ The message of anything that is not a provider answer is the `litellm_error::Rep A request can fail before the handler runs, in the handler, or after the headers went out. All three end in the same `Error`, so they get the same classification and envelope - An extractor sets `type Rejection = Error`, so a bad body is the surface's envelope rather than axum's plain text -- A path the gateway knows but does not serve returns `Error::Unsupported(path)` from one handler, so the client gets a 501 in the surface's envelope rather than a bare 404 +- A path or method no native route serves never reaches an `Error`. The merged router's fallbacks in `crates/gateway/src/dispatch` forward it to Python, and a standalone gateway answers axum's 404 or 405 +- `openai_response()` and `anthropic_response()` attach a route error's `Decline` to the response, so a handler that builds its envelope through them never drops a decline - A stream that fails after its headers went out hands the surface's SSE frame of `Error::from(error)` to `litellm_host_http::Sse::new`. `litellm_host_http::Error` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the `From` maps the second to `Internal` The axum hooks behind each line are in [`rust-gateway`](../../rust-gateway/SKILL.md) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 9592091eb60..4f40c4e1f32 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3761,11 +3761,14 @@ version = "0.1.0" dependencies = [ "axum", "base64 0.22.1", + "bytes", "envy", "futures-util", "http-body-util", + "hyper-util", "litellm-auth-types", "litellm-config", + "litellm-core-utils", "litellm-error", "litellm-gateway-auth", "litellm-gateway-inference", @@ -3776,6 +3779,7 @@ dependencies = [ "litellm-llms", "litellm-secrets", "litellm-tracing", + "reqwest", "rstest", "serde", "serde_json", @@ -3787,6 +3791,7 @@ dependencies = [ "tower-sessions-moka-store", "tracing", "uuid", + "wiremock", ] [[package]] @@ -3822,6 +3827,7 @@ dependencies = [ "litellm-auth", "litellm-cache-memory", "litellm-cache-response", + "litellm-core-utils", "litellm-error", "litellm-gateway-auth", "litellm-host", diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 034c3a9826e..fdc550bc773 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -11,6 +11,7 @@ axum = { workspace = true, features = ["json", "multipart", "original-uri"] } base64.workspace = true bytes.workspace = true litellm-auth.workspace = true +litellm-core-utils.workspace = true litellm-error.workspace = true litellm-gateway-auth.workspace = true litellm-inference.workspace = true diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 1ab55bd7487..5439883deb1 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_inference::Error as InferenceError; use litellm_llms::base_llm::ocr::error::Error as OcrError; use serde_json::{Map, Value, json}; @@ -23,8 +24,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("unexpected HTTP host operation")] @@ -73,7 +72,6 @@ impl Error { match self { Self::Auth(error) => error.status(), Self::InvalidPath(rejection) => rejection.status(), - Self::Unsupported(_) => StatusCode::NOT_IMPLEMENTED, Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE, Self::InvalidBody(_) | Self::UnknownModel { .. } => StatusCode::BAD_REQUEST, Self::Internal => StatusCode::INTERNAL_SERVER_ERROR, @@ -82,8 +80,34 @@ impl Error { } } + pub fn declined(&self) -> Option<&Decline> { + match self { + Self::Route(error) => error.declined(), + Self::Auth(_) + | Self::InvalidBody(_) + | Self::InvalidPath(_) + | Self::UnknownModel { .. } + | Self::Ocr(_) + | Self::BodyTooLarge + | Self::Internal => None, + } + } + pub fn openai_response(self) -> Response { - (self.status(), Json(self.openai_body())).into_response() + self.respond(self.openai_body()) + } + + pub fn anthropic_response(self, request_id: Option<&str>) -> Response { + self.respond(self.body(request_id)) + } + + /// 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 OpenAI envelope as an `event: error` frame, for a stream that fails after its @@ -141,7 +165,6 @@ impl Error { | Self::UnknownModel { .. } | Self::Route(_) | Self::Ocr(_) - | Self::Unsupported(_) | Self::BodyTooLarge | Self::Internal => litellm_error::Report::new(self).to_string(), } @@ -334,7 +357,6 @@ mod tests { #[case::missing_model(BodyError::MissingModel.into(), StatusCode::BAD_REQUEST)] #[case::unknown_model(unknown_model(), StatusCode::BAD_REQUEST)] #[case::oversized_body(Error::BodyTooLarge, StatusCode::PAYLOAD_TOO_LARGE)] - #[case::unserved_path(Error::Unsupported("/v1/completions".into()), StatusCode::NOT_IMPLEMENTED)] #[case::host_protocol( litellm_host_http::Error::::Protocol.into(), StatusCode::INTERNAL_SERVER_ERROR diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 43966f2bee7..e9bef1752b6 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -27,7 +27,7 @@ use litellm_secrets::source::SecretSource; pub use error::{BodyError, Error}; pub use litellm_router::{Deployment, Router as ModelRouter}; -pub use request::{JsonObject, RequestId}; +pub use request::{JsonObject, MAX_BODY_BYTES, RequestId}; pub struct Gateway { cache: Option>, @@ -104,10 +104,6 @@ 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)) .layer(axum::extract::DefaultBodyLimit::max( request::MAX_BODY_BYTES, )) @@ -115,11 +111,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 58ecd40204d..d52b7c857f4 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 d41794242ff..20c1c63e705 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -1,15 +1,14 @@ 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}; use crate::{BodyError, Deployment, Error, Gateway}; pub(crate) const MAX_FILE_BYTES: usize = 50 * 1024 * 1024; -pub(crate) const MAX_BODY_BYTES: usize = MAX_FILE_BYTES + 1024 * 1024; +pub const MAX_BODY_BYTES: usize = MAX_FILE_BYTES + 1024 * 1024; pub(crate) struct Upload { pub bytes: Bytes, @@ -179,10 +178,6 @@ fn multipart_error(error: axum::extract::multipart::MultipartError) -> Error { BodyError::MultipartField(error).into() } -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/decline.rs b/litellm-rust/crates/gateway-inference/tests/decline.rs new file mode 100644 index 00000000000..65b0b694484 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/decline.rs @@ -0,0 +1,72 @@ +mod support; + +use litellm_core_utils::decline::Decline; +use rstest::rstest; +use serde_json::{Value, json}; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; + +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", + json!({"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", + json!({"model": "public/model", "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}), + "Bedrock invoke messages request shaping" +)] +#[case::responses_provider( + "gemini/test-model", + "/v1/responses", + json!({"model": "public/model", "input": "hi"}), + "native HTTP responses provider" +)] +#[tokio::test] +async fn a_core_decline_is_a_501_that_carries_the_decline( + #[case] model: &str, + #[case] path: &str, + #[case] body: Value, + #[case] reason: &'static str, +) { + let provider = provider().await; + let response = support::post(support::app(model, &provider.uri()), path, body).await; + + assert_eq!(response.status(), 501); + assert_eq!( + response.extensions().get::(), + Some(&Decline::new(reason)) + ); + let body = support::json(response).await; + assert!(body.to_string().contains(reason)); + 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_carries_no_decline(#[case] body: Value, #[case] status: u16) { + let provider = provider().await; + let response = support::post( + support::app("anthropic/test-model", &provider.uri()), + "/v1/chat/completions", + body, + ) + .await; + + assert_eq!(response.status(), status); + assert!(response.extensions().get::().is_none()); +} diff --git a/litellm-rust/crates/gateway-inference/tests/routes.rs b/litellm-rust/crates/gateway-inference/tests/routes.rs index 24b8c5cc694..720e9b83245 100644 --- a/litellm-rust/crates/gateway-inference/tests/routes.rs +++ b/litellm-rust/crates/gateway-inference/tests/routes.rs @@ -160,18 +160,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] @@ -291,7 +287,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/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index d50fa5d32ae..4135ee4b0ff 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -6,9 +6,12 @@ license.workspace = true repository.workspace = true [dependencies] -axum.workspace = true +axum = { workspace = true, features = ["original-uri"] } +bytes.workspace = true envy = "0.4.2" http-body-util = "0.1" +hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] } +litellm-core-utils.workspace = true litellm-inference.workspace = true litellm-gateway-inference.workspace = true litellm-gateway-mcp.workspace = true @@ -32,6 +35,9 @@ uuid.workspace = true tower-sessions-moka-store = "0.15.0" [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } +reqwest.workspace = true +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/dispatch/forward.rs b/litellm-rust/crates/gateway/src/dispatch/forward.rs new file mode 100644 index 00000000000..255df04a0a4 --- /dev/null +++ b/litellm-rust/crates/gateway/src/dispatch/forward.rs @@ -0,0 +1,179 @@ +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, InvalidUri, 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. +const DECLINED_HEADER: HeaderName = HeaderName::from_static("x-litellm-gateway-declined"); + +const X_FORWARDED_FOR: HeaderName = HeaderName::from_static("x-forwarded-for"); + +const KEEP_ALIVE: HeaderName = HeaderName::from_static("keep-alive"); + +const HOP_BY_HOP: [HeaderName; 8] = [ + CONNECTION, + 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 {upstream}")] +pub struct InvalidUpstream { + upstream: String, + #[source] + cause: Option, +} + +/// 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 = |cause| InvalidUpstream { + upstream: upstream.to_owned(), + cause, + }; + let uri: Uri = upstream.parse().map_err(|cause| invalid(Some(cause)))?; + 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(None)), + }; + Ok(Self { + client: Arc::new(Client::builder(TokioExecutor::new()).build(HttpConnector::new())), + authority, + }) + } + + pub(super) 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!(connect = error.is_connect(), "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); + } +} + +#[cfg(test)] +mod tests { + use axum::http::uri::InvalidUri; + use rstest::rstest; + + use super::Forward; + + #[rstest] + #[case::https("https://127.0.0.1:4000", false)] + #[case::with_path("http://127.0.0.1:4000/v1", false)] + #[case::no_scheme("127.0.0.1:4000", false)] + #[case::unparsable("http://exa mple", true)] + fn only_an_http_origin_is_a_python_upstream(#[case] upstream: &str, #[case] unparsable: bool) { + let Err(error) = Forward::new(upstream) else { + panic!("{upstream} must be rejected"); + }; + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(upstream), "{report}"); + assert_eq!( + litellm_error::find::(&error).is_some(), + unparsable + ); + } + + #[rstest] + #[case::bare("http://127.0.0.1:4000")] + #[case::root("http://127.0.0.1:4000/")] + fn an_http_origin_is_a_python_upstream(#[case] upstream: &str) { + assert!(Forward::new(upstream).is_ok()); + } +} diff --git a/litellm-rust/crates/gateway/src/dispatch/mod.rs b/litellm-rust/crates/gateway/src/dispatch/mod.rs new file mode 100644 index 00000000000..dab89790257 --- /dev/null +++ b/litellm-rust/crates/gateway/src/dispatch/mod.rs @@ -0,0 +1,52 @@ +//! What the gateway does with a request it does not serve natively. +//! +//! A mounted router declines by putting a [`Decline`] in its response's extensions. Nothing +//! else crosses this boundary: only this module records requests and talks to Python. + +mod forward; +mod replay; + +pub use forward::{Forward, InvalidUpstream}; + +use axum::{Router, extract::Request}; +use litellm_core_utils::decline::Decline; + +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 mounted router's own error. + Reject, + /// Next to the Python proxy: both go to Python unchanged. + Forward(Forward), +} + +/// Forwards declined requests on the routes `served` holds so far. `record_limit` must cover +/// the largest body those routes accept, or a declined request that size keeps its error. +pub(crate) fn replay_declines(served: Router, fallback: &Fallback, record_limit: usize) -> Router { + match fallback { + Fallback::Reject => served, + Fallback::Forward(forward) => served.route_layer(axum::middleware::from_fn_with_state( + replay::Replay { + forward: forward.clone(), + record_limit, + }, + replay::replay, + )), + } +} + +pub(crate) fn forward_unmatched(app: Router, fallback: Fallback) -> Router { + match fallback { + Fallback::Reject => app, + Fallback::Forward(forward) => { + let unmatched = forward.clone(); + app.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/dispatch/replay.rs b/litellm-rust/crates/gateway/src/dispatch/replay.rs new file mode 100644 index 00000000000..0d371cd7902 --- /dev/null +++ b/litellm-rust/crates/gateway/src/dispatch/replay.rs @@ -0,0 +1,78 @@ +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::Forward; + +#[derive(Clone)] +pub(super) struct Replay { + pub forward: Forward, + pub record_limit: usize, +} + +/// Records the request before any handler reads it, and forwards the recording when the +/// response says the request was declined. +pub(super) async fn replay( + State(Replay { + forward, + record_limit, + }): State, + request: Request, + next: Next, +) -> Response { + let (parts, body) = request.into_parts(); + let body = match buffer(body, record_limit).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; + }; + 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, limit: usize) -> 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() <= limit => { + 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/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index a24069f3f6d..63fb0bd7c42 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -1,6 +1,8 @@ mod auth; +mod dispatch; mod error; mod secrets; +pub use dispatch::{Fallback, Forward, InvalidUpstream}; pub use error::Error; use std::{sync::Arc, time::Instant}; @@ -118,20 +120,33 @@ 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) - .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, - } + let routes = litellm_gateway_inference::router(inference).merge(mcp.unwrap_or_default()); + let served = + dispatch::replay_declines(routes, &fallback, litellm_gateway_inference::MAX_BODY_BYTES) + .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)) + .merge(ui.unwrap_or_default()); + dispatch::forward_unmatched(served, fallback) +} + +pub async fn serve( + listener: tokio::net::TcpListener, + app: Router, + shutdown: impl Future + Send + 'static, +) -> std::io::Result<()> { + axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(shutdown) + .await } async fn log_request(request: Request, next: Next) -> Response { diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index f8150e6e860..08ad3664558 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -4,6 +4,7 @@ use std::{ }; use litellm_config::Config; +use litellm_gateway::{Fallback, Forward}; use litellm_gateway_auth::UiBackend; 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,13 +102,13 @@ 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( + let result = litellm_gateway::serve( listener, - litellm_gateway::router(inference, &config, ui, mcp_router), + litellm_gateway::router(inference, &config, ui, mcp_router, fallback), + shutdown.clone().cancelled_owned(), ) - .with_graceful_shutdown(shutdown.clone().cancelled_owned()) .await; shutdown.cancel(); if let Some(mcp) = mcp { 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/dispatch.rs b/litellm-rust/crates/gateway/tests/dispatch.rs new file mode 100644 index 00000000000..f4a00a185a0 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/dispatch.rs @@ -0,0 +1,363 @@ +use std::{ + convert::Infallible, + net::SocketAddr, + sync::{Arc, Mutex}, + time::Duration, +}; + +use axum::{ + Router, + body::{Body, Bytes, to_bytes}, + extract::{ConnectInfo, DefaultBodyLimit}, + http::{HeaderMap, Request, StatusCode}, + response::{IntoResponse, Response}, +}; +use futures_util::{StreamExt, stream}; +use litellm_config::Config; +use litellm_gateway::{Fallback, Forward}; +use litellm_gateway_inference::MAX_BODY_BYTES; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::{ + net::TcpListener, + sync::{mpsc, oneshot}, + time::timeout, +}; +use tower::ServiceExt; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; + +fn app(python: &str, model: &str, provider: &str) -> Router { + let config = Config::from_yaml(&format!( + "general_settings:\n master_key: gateway-key\nmodel_list:\n - model_name: public/model\n litellm_params:\n model: {model}\n api_base: {provider}\n api_key: test-key\n" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let fallback = Fallback::Forward(Forward::new(python).unwrap()); + litellm_gateway::router(inference, &config, None, None, fallback) +} + +async fn send(app: Router, path: &str, body: Value) -> Response { + app.oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer gateway-key") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap() +} + +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 = app(&python.uri(), model, &provider.uri()); + + let response = app + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer gateway-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 gateway-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] +#[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: Value, #[case] status: u16) { + let python = python().await; + let provider = provider().await; + let response = send( + app(&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") + .header("authorization", "Bearer gateway-key"); + 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(), + )); + + app(&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); +} + +async fn listen(app: Router) -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(litellm_gateway::serve( + listener, + app, + std::future::pending(), + )); + address +} + +async fn gateway_in_front_of(python: SocketAddr) -> SocketAddr { + listen(app( + &format!("http://{python}"), + "anthropic/test-model", + "http://127.0.0.1:1", + )) + .await +} + +async fn recording_python( + respond: fn() -> Response, +) -> (SocketAddr, mpsc::UnboundedReceiver<(HeaderMap, Bytes)>) { + let (sender, received) = mpsc::unbounded_channel(); + let python = Router::new() + .fallback(move |headers: HeaderMap, body: Bytes| async move { + sender.send((headers, body)).unwrap(); + respond() + }) + .layer(DefaultBodyLimit::disable()); + (listen(python).await, received) +} + +fn streaming_chat(content: &str) -> String { + json!({"model": "public/model", "stream": true, "messages": [{"role": "user", "content": content}]}) + .to_string() +} + +fn streaming_chat_of_size(size: usize) -> String { + streaming_chat(&"a".repeat(size - streaming_chat("").len())) +} + +fn chat(gateway: SocketAddr, body: String) -> reqwest::RequestBuilder { + litellm_http::Client::for_test() + .post(format!("http://{gateway}/v1/chat/completions")) + .header("content-type", "application/json") + .header("authorization", "Bearer gateway-key") + .timeout(Duration::from_secs(10)) + .body(body) +} + +#[rstest] +#[tokio::test] +async fn a_forwarded_stream_reaches_the_client_before_python_finishes_it() { + let (finish, finished) = oneshot::channel::<()>(); + let finished = Arc::new(Mutex::new(Some(finished))); + let python = listen(Router::new().fallback(move || async move { + let finished = finished.lock().unwrap().take().unwrap(); + let first = stream::once(async { Ok::<_, Infallible>(Bytes::from("data: first\n\n")) }); + let last = stream::once(async move { + let _ = finished.await; + Ok(Bytes::from("data: last\n\n")) + }); + ( + [("content-type", "text/event-stream")], + Body::from_stream(first.chain(last)), + ) + })) + .await; + let gateway = gateway_in_front_of(python).await; + + let response = chat(gateway, streaming_chat("hi")).send().await.unwrap(); + + assert_eq!(response.status(), 200); + assert_eq!(response.headers()["content-type"], "text/event-stream"); + let mut body = response.bytes_stream(); + let first = timeout(Duration::from_secs(5), body.next()) + .await + .expect("the first chunk arrives while Python holds the stream open"); + assert_eq!(first.unwrap().unwrap(), "data: first\n\n"); + finish.send(()).unwrap(); + let rest: Vec = body.map(Result::unwrap).collect().await; + assert_eq!(rest.concat(), b"data: last\n\n"); +} + +#[rstest] +#[tokio::test] +async fn python_answers_the_client_with_its_own_status_and_end_to_end_headers() { + let (python, _received) = recording_python(|| { + ( + StatusCode::TOO_MANY_REQUESTS, + [ + ("retry-after", "7"), + ("connection", "x-python-hop"), + ("x-python-hop", "internal"), + ], + "slow down", + ) + .into_response() + }) + .await; + let gateway = gateway_in_front_of(python).await; + let direct = litellm_http::Client::for_test() + .post(format!("http://{python}/")) + .send() + .await + .unwrap(); + assert_eq!(direct.headers()["x-python-hop"], "internal"); + + let response = chat(gateway, streaming_chat("hi")).send().await.unwrap(); + + assert_eq!(response.status(), 429); + assert_eq!(response.headers()["retry-after"], "7"); + assert!(!response.headers().contains_key("x-python-hop")); + assert_eq!(response.text().await.unwrap(), "slow down"); +} + +#[rstest] +#[tokio::test] +async fn python_sees_the_connecting_client_at_the_end_of_x_forwarded_for() { + let (python, mut received) = recording_python(|| "from python".into_response()).await; + let gateway = gateway_in_front_of(python).await; + + chat(gateway, streaming_chat("hi")) + .header("x-forwarded-for", "198.51.100.1") + .send() + .await + .unwrap(); + + let (headers, _) = received.try_recv().unwrap(); + assert_eq!(headers["x-forwarded-for"], "198.51.100.1, 127.0.0.1"); + assert_eq!(headers["x-litellm-gateway-declined"], "streaming"); +} + +#[rstest] +#[case::at_the_limit(MAX_BODY_BYTES, 200, true)] +#[case::one_byte_over(MAX_BODY_BYTES + 1, 413, false)] +#[tokio::test] +async fn a_declined_body_up_to_the_limit_reaches_python_byte_for_byte( + #[case] size: usize, + #[case] status: u16, + #[case] forwarded: bool, +) { + let (python, mut received) = recording_python(|| "from python".into_response()).await; + let gateway = gateway_in_front_of(python).await; + let body = streaming_chat_of_size(size); + + let response = chat(gateway, body.clone()).send().await.unwrap(); + + assert_eq!(response.status().as_u16(), status); + let received = received.try_recv().ok().map(|(_, received)| received); + assert_eq!(received.as_deref() == Some(body.as_bytes()), forwarded); + assert_eq!(received.is_some(), forwarded); +} + +#[rstest] +#[tokio::test] +async fn an_unreachable_python_is_a_bad_gateway() { + let closed = TcpListener::bind("127.0.0.1:0") + .await + .unwrap() + .local_addr() + .unwrap(); + let gateway = gateway_in_front_of(closed).await; + + let response = chat(gateway, streaming_chat("hi")).send().await.unwrap(); + + assert_eq!(response.status(), 502); +} diff --git a/litellm-rust/crates/gateway/tests/mcp.rs b/litellm-rust/crates/gateway/tests/mcp.rs index e16caa3e92e..023ff0f5c81 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::Fallback; use litellm_gateway_mcp::{ HttpConfig, rmcp::{ @@ -163,7 +164,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); Served { app, mcp, diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 41d04375580..1f86ff8825a 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -6,6 +6,7 @@ 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::{Fallback, Forward}; use litellm_gateway_inference::{Error, Gateway}; use litellm_http::ClientVariant; use litellm_tracing::{Logger, Metadata, Record, Sink}; @@ -74,7 +75,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; @@ -132,7 +133,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); + } +}