From f9fbeca60f85656058d99b26d562d5de65383bac Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 21 Oct 2026 07:32:23 -0700 Subject: [PATCH] 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 --- litellm-rust/Cargo.lock | 5 +- .../crates/gateway-inference/Cargo.toml | 4 +- .../gateway-inference/src/dispatch/mod.rs | 24 ------- .../crates/gateway-inference/src/lib.rs | 9 +-- .../crates/gateway-inference/src/request.rs | 2 +- .../crates/gateway-inference/tests/decline.rs | 72 +++++++++++++++++++ .../gateway-inference/tests/support/mod.rs | 24 +------ litellm-rust/crates/gateway/Cargo.toml | 5 +- .../src/dispatch/forward.rs | 4 +- .../crates/gateway/src/dispatch/mod.rs | 52 ++++++++++++++ .../src/dispatch/replay.rs | 27 ++++--- litellm-rust/crates/gateway/src/lib.rs | 40 ++++------- litellm-rust/crates/gateway/src/main.rs | 2 +- .../tests/dispatch.rs | 70 +++++++++--------- litellm-rust/crates/gateway/tests/mcp.rs | 2 +- litellm-rust/crates/gateway/tests/server.rs | 6 +- 16 files changed, 211 insertions(+), 137 deletions(-) delete mode 100644 litellm-rust/crates/gateway-inference/src/dispatch/mod.rs create mode 100644 litellm-rust/crates/gateway-inference/tests/decline.rs rename litellm-rust/crates/{gateway-inference => gateway}/src/dispatch/forward.rs (96%) create mode 100644 litellm-rust/crates/gateway/src/dispatch/mod.rs rename litellm-rust/crates/{gateway-inference => gateway}/src/dispatch/replay.rs (81%) rename litellm-rust/crates/{gateway-inference => gateway}/tests/dispatch.rs (74%) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4259ca1965b..8d30a5e0bd0 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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", ] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index a316ff58b50..ed8cdfd60d9 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -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"] } diff --git a/litellm-rust/crates/gateway-inference/src/dispatch/mod.rs b/litellm-rust/crates/gateway-inference/src/dispatch/mod.rs deleted file mode 100644 index 845394da8f9..00000000000 --- a/litellm-rust/crates/gateway-inference/src/dispatch/mod.rs +++ /dev/null @@ -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), -} diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index b980ea2d186..16f4d492133 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -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>, @@ -89,7 +88,7 @@ impl Gateway { } } -pub fn router(gateway: Arc, fallback: dispatch::Fallback) -> Router { +pub fn router(gateway: Arc) -> Router { Router::new() .route("/v1/messages", post(messages::create)) .route("/ocr", post(ocr::create)) @@ -105,10 +104,6 @@ pub fn router(gateway: Arc, 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, )) diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs index 293c37ec551..41f1aab3ebe 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -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, 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/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 5366cbd4118..0e37bb0aff9 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -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>, principal: Option, - 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, )) diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index e8edf2b114a..28eab1466a1 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 diff --git a/litellm-rust/crates/gateway-inference/src/dispatch/forward.rs b/litellm-rust/crates/gateway/src/dispatch/forward.rs similarity index 96% rename from litellm-rust/crates/gateway-inference/src/dispatch/forward.rs rename to litellm-rust/crates/gateway/src/dispatch/forward.rs index dc24f19b4a0..664c4bfe121 100644 --- a/litellm-rust/crates/gateway-inference/src/dispatch/forward.rs +++ b/litellm-rust/crates/gateway/src/dispatch/forward.rs @@ -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 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-inference/src/dispatch/replay.rs b/litellm-rust/crates/gateway/src/dispatch/replay.rs similarity index 81% rename from litellm-rust/crates/gateway-inference/src/dispatch/replay.rs rename to litellm-rust/crates/gateway/src/dispatch/replay.rs index 5c4ff545013..0d371cd7902 100644 --- a/litellm-rust/crates/gateway-inference/src/dispatch/replay.rs +++ b/litellm-rust/crates/gateway/src/dispatch/replay.rs @@ -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, +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).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::() 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 { +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() <= MAX_BODY_BYTES => { + Ok(chunk) if seen.len() + chunk.len() <= limit => { seen.extend_from_slice(&chunk); } failed => { diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 44de2ff9570..d74fc670c09 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}; @@ -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 { diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 81fc214bb4b..ff489448b85 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -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; diff --git a/litellm-rust/crates/gateway-inference/tests/dispatch.rs b/litellm-rust/crates/gateway/tests/dispatch.rs similarity index 74% rename from litellm-rust/crates/gateway-inference/tests/dispatch.rs rename to litellm-rust/crates/gateway/tests/dispatch.rs index b546a1a89a6..a063d81cf48 100644 --- a/litellm-rust/crates/gateway-inference/tests/dispatch.rs +++ b/litellm-rust/crates/gateway/tests/dispatch.rs @@ -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::().unwrap(), )); - support::app_forwarding_to(&python.uri(), "anthropic/test-model", &provider.uri()) + app(&python.uri(), "anthropic/test-model", &provider.uri()) .oneshot(request) .await .unwrap(); diff --git a/litellm-rust/crates/gateway/tests/mcp.rs b/litellm-rust/crates/gateway/tests/mcp.rs index a017a440efc..69212512f83 100644 --- a/litellm-rust/crates/gateway/tests/mcp.rs +++ b/litellm-rust/crates/gateway/tests/mcp.rs @@ -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::{ diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 24093640f5e..b4ccb90fa23 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -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};