refactor(gateway): own Python forwarding in the gateway crate

Forwarding declined and unserved requests to Python is a gateway concern,
not an inference one. The dispatch module moves into the gateway crate and
replays declines for every mounted router, so MCP and future ported crates
decline the same way inference does, by attaching a core Decline to their
error response

gateway-inference keeps only its side of that contract and drops its hyper
client. The replay layer is installed only when a Python upstream is set,
so a standalone gateway no longer buffers request bodies

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-10-21 07:32:23 -07:00
parent 213740863b
commit f9fbeca60f
16 changed files with 211 additions and 137 deletions

View file

@ -3732,11 +3732,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-gateway-auth",
"litellm-gateway-inference",
"litellm-gateway-mcp",
@ -3789,7 +3792,6 @@ dependencies = [
"base64 0.22.1",
"bytes",
"futures-util",
"hyper-util",
"litellm-auth",
"litellm-cache-memory",
"litellm-cache-response",
@ -3814,7 +3816,6 @@ dependencies = [
"thiserror 2.0.19",
"tokio",
"tower",
"tracing",
"wiremock",
]

View file

@ -10,8 +10,6 @@ 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
@ -31,10 +29,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"] }

View file

@ -1,24 +0,0 @@
//! 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),
}

View file

@ -6,7 +6,6 @@
mod audio_transcription;
mod caching;
mod chat_completions;
pub mod dispatch;
mod error;
pub mod messages;
mod ocr;
@ -28,7 +27,7 @@ use litellm_secrets::source::SecretSource;
pub use error::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<Arc<dyn litellm_cache_response::ResponseCacheService>>,
@ -89,7 +88,7 @@ impl Gateway {
}
}
pub fn router(gateway: Arc<Gateway>, fallback: dispatch::Fallback) -> Router {
pub fn router(gateway: Arc<Gateway>) -> Router {
Router::new()
.route("/v1/messages", post(messages::create))
.route("/ocr", post(ocr::create))
@ -105,10 +104,6 @@ pub fn router(gateway: Arc<Gateway>, fallback: dispatch::Fallback) -> Router {
)
.route("/responses", post(responses::create))
.route("/v1/responses", post(responses::create))
.route_layer(axum::middleware::from_fn_with_state(
fallback,
dispatch::replay,
))
.layer(axum::extract::DefaultBodyLimit::max(
request::MAX_BODY_BYTES,
))

View file

@ -8,7 +8,7 @@ use serde_json::{Map, Value};
use crate::{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,

View file

@ -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::<Decline>(),
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::<Decline>().is_none());
}

View file

@ -10,11 +10,7 @@ use axum::{
response::Response,
};
use futures_util::future::BoxFuture;
use litellm_gateway_inference::{
Deployment, Gateway,
dispatch::{Fallback, Forward},
router,
};
use litellm_gateway_inference::{Deployment, Gateway, router};
use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver};
use litellm_inference::resources::CoreResources;
use litellm_secrets::{SecretValue, source::SecretSource};
@ -41,18 +37,7 @@ pub fn app_with_permissions(
api_base: &str,
permissions: litellm_gateway_auth::Permissions,
) -> Router {
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()),
)
configured_app(model, api_base, permissions, None, None)
}
pub fn app_with_cache(
@ -66,7 +51,6 @@ pub fn app_with_cache(
litellm_gateway_auth::Permissions::All,
Some(cache),
None,
Fallback::Reject,
)
}
@ -82,7 +66,6 @@ pub fn app_with_cache_for_principal(
litellm_gateway_auth::Permissions::All,
Some(cache),
Some(principal),
Fallback::Reject,
)
}
@ -92,7 +75,6 @@ fn configured_app(
permissions: litellm_gateway_auth::Permissions,
cache: Option<Arc<dyn litellm_cache_response::ResponseCacheService>>,
principal: Option<litellm_gateway_auth::Principal>,
fallback: Fallback,
) -> Router {
let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver)));
let http = Resolution::from(&HttpSettings::default()).config;
@ -120,7 +102,7 @@ fn configured_app(
Some(cache) => gateway.with_cache(cache),
None => gateway,
};
router(Arc::new(gateway), fallback).layer(axum::middleware::from_fn_with_state(
router(Arc::new(gateway)).layer(axum::middleware::from_fn_with_state(
(permissions, principal),
test_identity,
))

View file

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

View file

@ -20,7 +20,7 @@ use hyper_util::{
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 DECLINED_HEADER: HeaderName = HeaderName::from_static("x-litellm-gateway-declined");
const X_FORWARDED_FOR: HeaderName = HeaderName::from_static("x-forwarded-for");
@ -64,7 +64,7 @@ impl Forward {
})
}
pub async fn send(&self, request: Request, decline: Decline) -> Response {
pub(super) async fn send(&self, request: Request, decline: Decline) -> Response {
let (mut parts, body) = request.into_parts();
let path = parts
.extensions

View file

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

View file

@ -11,18 +11,26 @@ use bytes::BytesMut;
use futures_util::StreamExt;
use litellm_core_utils::decline::Decline;
use super::Fallback;
use crate::request::MAX_BODY_BYTES;
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(crate) async fn replay(
State(fallback): State<Fallback>,
pub(super) async fn replay(
State(Replay {
forward,
record_limit,
}): State<Replay>,
request: Request,
next: Next,
) -> Response {
let (parts, body) = request.into_parts();
let body = match buffer(body).await {
let body = match buffer(body, record_limit).await {
Ok(body) => body,
Err(body) => return next.run(Request::from_parts(parts, body)).await,
};
@ -31,20 +39,17 @@ pub(crate) async fn replay(
let Some(&decline) = response.extensions().get::<Decline>() else {
return response;
};
match fallback {
Fallback::Reject => response,
Fallback::Forward(forward) => forward.send(recorded.map(Body::from), decline).await,
}
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<Bytes, Body> {
async fn buffer(body: Body, limit: usize) -> Result<Bytes, Body> {
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 => {
Ok(chunk) if seen.len() + chunk.len() <= limit => {
seen.extend_from_slice(&chunk);
}
failed => {

View file

@ -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};
@ -16,10 +18,7 @@ use http_body_util::BodyExt;
use litellm_config::Config;
use litellm_gateway_auth::Auth;
use litellm_gateway_inference::{
Gateway, ModelRouter,
dispatch::{Fallback, NOT_NATIVE},
};
use litellm_gateway_inference::{Gateway, ModelRouter};
use litellm_http::{
ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify,
media::PublicDnsResolver,
@ -124,28 +123,17 @@ pub fn router(
fallback: Fallback,
) -> Router {
let auth = Auth::from_config(config, inference.secrets.clone());
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))
.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
})
}
}
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)
}
async fn log_request(request: Request, next: Next) -> Response {

View file

@ -4,8 +4,8 @@ use std::{
};
use litellm_config::Config;
use litellm_gateway::{Fallback, Forward};
use litellm_gateway_auth::UiBackend;
use litellm_gateway_inference::dispatch::{Fallback, Forward};
use litellm_tracing::{Level, Logger, Metadata, Record, Sink};
use serde_json::json;

View file

@ -1,17 +1,41 @@
mod support;
use std::net::SocketAddr;
use axum::{
Router,
body::{Body, to_bytes},
extract::ConnectInfo,
http::Request,
response::Response,
};
use litellm_config::Config;
use litellm_gateway::{Fallback, Forward};
use rstest::rstest;
use serde_json::json;
use serde_json::{Value, json};
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())
@ -62,13 +86,13 @@ async fn a_core_decline_reaches_python_as_the_client_sent_it(
) {
let python = python().await;
let provider = provider().await;
let app = support::app_forwarding_to(&python.uri(), model, &provider.uri());
let app = app(&python.uri(), model, &provider.uri());
let response = app
.oneshot(
Request::post(path)
.header("content-type", "application/json")
.header("authorization", "Bearer client-key")
.header("authorization", "Bearer gateway-key")
.header("x-client-header", "kept")
.header("connection", "te")
.header("te", "trailers")
@ -95,44 +119,22 @@ async fn a_core_decline_reaches_python_as_the_client_sent_it(
};
assert_eq!(forwarded_path, path);
assert_eq!(forwarded.body, body.as_bytes());
assert_eq!(forwarded.headers["authorization"], "Bearer client-key");
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]
#[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) {
async fn a_terminal_error_is_never_forwarded(#[case] body: 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()),
let response = send(
app(&python.uri(), "anthropic/test-model", &provider.uri()),
"/v1/chat/completions",
body,
)
@ -152,7 +154,9 @@ async fn a_forwarded_decline_names_the_client_in_x_forwarded_for(
) {
let python = python().await;
let provider = provider().await;
let request = Request::post("/v1/chat/completions").header("content-type", "application/json");
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,
@ -167,7 +171,7 @@ async fn a_forwarded_decline_names_the_client_in_x_forwarded_for(
"203.0.113.7:50000".parse::<SocketAddr>().unwrap(),
));
support::app_forwarding_to(&python.uri(), "anthropic/test-model", &provider.uri())
app(&python.uri(), "anthropic/test-model", &provider.uri())
.oneshot(request)
.await
.unwrap();

View file

@ -6,7 +6,7 @@ use axum::{
http::Request,
};
use litellm_config::Config;
use litellm_gateway_inference::dispatch::Fallback;
use litellm_gateway::Fallback;
use litellm_gateway_mcp::{
HttpConfig,
rmcp::{

View file

@ -6,10 +6,8 @@ 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,
dispatch::{Fallback, Forward},
};
use litellm_gateway::{Fallback, Forward};
use litellm_gateway_inference::{Error, Gateway};
use litellm_http::ClientVariant;
use litellm_tracing::{Logger, Metadata, Record, Sink};
use rstest::{fixture, rstest};