scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 233

fix(gateway): stop serving routes outside the hooks and return reservations of failed requests (#47355)

Browse source 
* fix(gateway): refuse spent keys on every route and return reservations of requests that fail before routing

Auth refuses a key or team whose stored spend has reached its budget again, so
/v1/responses, audio and OCR, which skip the hook pipeline, stop exhausted keys
as the proxy does. The budgets hook still enforces the live counters on
messages and chat.

A request whose body is rejected after the hooks ran, such as an invalid
`cache` option, now records its failure so the budget reservation and rate
limit leases are returned. A request dropped mid-route, as when the client
goes away, is recorded as interrupted with status 499 for the same reason.

* fix(gateway): stop serving responses, audio and ocr instead of refusing spent keys in auth

/v1/responses, /v1/audio/transcriptions and /v1/ocr never ran the hook
pipeline, so budgets, rate limits and spend logging did not apply to them.
They now answer 501 like the other routes the gateway does not serve, and the
stored-spend check in auth is dropped again: the budgets hook covers every
route that remains.

* chore(gateway): update Cargo.lock for the dropped route crates

* fix(gateway): keep the reservation of a request cut off after it reached the provider
This commit is contained in:
Nate Armstrong 2026-10-19 14:38:45 -07:00 • committed by GitHub
parent 8fb0643e35
commit 4d7802f169
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 211 additions and 820 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm-rust/Cargo.lock generated
View file

@ -3855,9 +3855,6 @@ dependencies = [
"litellm-inference",
"litellm-inference-chat",
"litellm-inference-messages",
"litellm-inference-ocr",
"litellm-inference-responses",
"litellm-inference-transcription",
"litellm-llms",
"litellm-llms-types",
"litellm-router",

9
litellm-rust/crates/gateway-hooks/src/facts.rs
View file

@ -276,6 +276,15 @@ impl Record {
(200..300).contains(&self.outcome.status) && !self.outcome.interrupted
}
/// The request reached a provider and was cut off before it reported usage, as when the
/// client goes away while the provider works. The provider may still bill it, so what it
/// cost is unknown and reservations should be kept rather than returned.
pub fn cut_off_after_dispatch(&self) -> bool {
self.outcome.interrupted
&& self.outcome.upstream_started_at.is_some()
&& self.outcome.usage.is_none()
}
pub fn duration(&self) -> Duration {
self.outcome
.ended_at

5
litellm-rust/crates/gateway-inference/Cargo.toml
View file

@ -7,18 +7,15 @@ repository.workspace = true
[dependencies]
litellm-cache-response.workspace = true
axum = { workspace = true, features = ["json", "multipart", "original-uri"] }
axum = { workspace = true, features = ["json", "original-uri"] }
base64.workspace = true
bytes.workspace = true
litellm-auth.workspace = true
litellm-gateway-auth.workspace = true
litellm-gateway-hooks.workspace = true
litellm-inference.workspace = true
litellm-inference-transcription.workspace = true
litellm-inference-responses.workspace = true
litellm-inference-messages.workspace = true
litellm-inference-chat.workspace = true
litellm-inference-ocr.workspace = true
litellm-host-http.workspace = true
litellm-host.workspace = true
litellm-http.workspace = true

61
litellm-rust/crates/gateway-inference/src/audio_transcription.rs
View file

@ -1,61 +0,0 @@
use litellm_gateway_auth::AuthenticatedRequest;
use std::{path::Path, sync::Arc};
use axum::{Json, extract::State, response::IntoResponse};
use base64::{Engine, engine::general_purpose::STANDARD};
use litellm_inference_transcription::types::AudioTranscriptionRequest;
use serde_json::{Value, json};
use crate::{
Error, Gateway,
request::{self, InferenceBody},
};
pub(crate) async fn create(
State(gateway): State<Arc<Gateway>>,
identity: AuthenticatedRequest,
body: InferenceBody,
) -> Result<impl IntoResponse, Error> {
handle(&gateway, &identity, body).await.map(Json)
}
async fn handle(
gateway: &Gateway,
identity: &AuthenticatedRequest,
InferenceBody {
fields: body,
upload,
}: InferenceBody,
) -> Result<Value, Error> {
gateway.check_body(&body)?;
let deployment = request::resolve_deployment(gateway, identity, &body)?;
request::authorize_model(identity, deployment, &body).await?;
let audio = match upload {
Some(upload) => {
let format = upload
.file_name
.as_deref()
.and_then(|name| Path::new(name).extension())
.and_then(|extension| extension.to_str())
.map(str::to_ascii_lowercase);
json!({"data": STANDARD.encode(upload.bytes), "format": format})
}
None => body.get("audio").cloned().unwrap_or_default(),
};
Ok(gateway
.audio_transcription
.execute(AudioTranscriptionRequest {
model: &deployment.model,
audio,
api_key: deployment.api_key.as_deref(),
api_base: deployment.api_base.as_deref(),
custom_llm_provider: deployment.custom_llm_provider.as_deref(),
extra_headers: None,
optional_params: body
.into_iter()
.filter(|(name, _)| !matches!(name.as_str(), "model" | "audio"))
.collect(),
timeout: deployment.timeout,
})
.await?)
}

10
litellm-rust/crates/gateway-inference/src/chat_completions.rs
View file

@ -74,7 +74,15 @@ async fn handle(
&name,
)
.await?;
let (mut body, cache_options) = crate::caching::prepare(identity, body)?;
let (mut body, cache_options) = match crate::caching::prepare(identity, body) {
Ok(prepared) => prepared,
Err(error) => {
if let Some(tracking) = tracking {
tracking.fail(&error);
}
return Err(error);
}
};
let edit = |body: &mut Map<String, Value>| {
if let Some(tracking) = &tracking {
tracking.edit_body(body);

110
litellm-rust/crates/gateway-inference/src/hooks.rs
View file

@ -48,13 +48,16 @@ impl Hooks {
}
}
/// One request's admission, and the deployment its last attempt went to.
/// One request's admission, and the deployment its last attempt went to. Dropped before it
/// finishes, as when the client goes away mid-route, it records the request as interrupted so
/// the hooks still settle: what they reserved comes back unless the request already reached a
/// provider, which may bill it.
pub(crate) struct Tracking {
pipeline: Arc<Pipeline>,
route: Route,
admission: Admission,
admission: Option<Admission>,
/// The deployment the latest attempt went to, and when it was sent.
deployment: Mutex<Option<(DeploymentFacts, SystemTime)>>,
deployment: Mutex<Option<(Option<DeploymentFacts>, SystemTime)>>,
started: Instant,
}
@ -101,7 +104,7 @@ pub(crate) async fn admit(
Ok(Some(Tracking {
pipeline: pipeline.clone(),
route,
admission,
admission: Some(admission),
deployment: Mutex::new(None),
started: Instant::now(),
}))
@ -168,7 +171,9 @@ fn end_user(body: &Map<String, Value>) -> Option<&str> {
impl Tracking {
/// Applies the hooks' body changes to an attempt's body.
pub(crate) fn edit_body(&self, body: &mut Map<String, Value>) {
self.admission.edit().apply_upstream_body(body);
if let Some(admission) = &self.admission {
admission.edit().apply_upstream_body(body);
}
}
/// The headers the gateway would forward, after the hooks' changes; `None` when no hook
@ -177,7 +182,7 @@ impl Tracking {
&self,
forwarded: Vec<(String, String)>,
) -> Option<Vec<(String, String)>> {
let edit = self.admission.edit();
let edit = self.admission.as_ref()?.edit();
edit.changes_upstream_headers().then(|| {
let mut forwarded = forwarded;
edit.apply_upstream_headers(&mut forwarded);
@ -186,53 +191,73 @@ impl Tracking {
}
/// Notes the deployment an attempt went to; the last one is the one that served.
/// Its send time is kept even when no hook reads the deployment, so settling can tell a
/// request that reached a provider from one that did not.
pub(crate) fn attempt(&self, name: &str, attempt: &Attempt<'_>) {
if !self.pipeline.reads(Facts::DEPLOYMENT) {
return;
}
let deployment: &Deployment = attempt.deployment;
let facts = DeploymentFacts {
id: attempt.id.into(),
model_group: name.into(),
model: deployment.model.clone(),
provider: deployment.custom_llm_provider.clone(),
api_base: deployment.api_base.clone(),
};
let facts = self
.pipeline
.reads(Facts::DEPLOYMENT)
.then(|| DeploymentFacts {
id: attempt.id.into(),
model_group: name.into(),
model: deployment.model.clone(),
provider: deployment.custom_llm_provider.clone(),
api_base: deployment.api_base.clone(),
});
*self.deployment.lock().unwrap_or_else(|e| e.into_inner()) =
Some((facts, SystemTime::now()));
}
/// Records a request that ended with an error before a response.
pub(crate) fn fail(mut self, error: &Error) {
self.record_failure(error.status().as_u16(), false);
}
fn record_failure(&mut self, status: u16, interrupted: bool) {
let Some(admission) = self.admission.take() else {
return;
};
let (deployment, upstream_started_at) = self.take_deployment();
let outcome = Outcome {
deployment,
status,
usage: None,
cost: None,
time_to_first_byte: None,
upstream_started_at,
ended_at: SystemTime::now(),
interrupted,
};
self.pipeline.finish(admission, outcome, None);
}
fn take_deployment(&self) -> (Option<DeploymentFacts>, Option<SystemTime>) {
self.deployment
.lock()
.unwrap_or_else(|e| e.into_inner())
.take()
.map_or((None, None), |(facts, sent_at)| (facts, Some(sent_at)))
}
/// Adds the hooks' response headers and follows the body to its end; an error is recorded
/// at once.
pub(crate) async fn finish(self, result: Result<Response, Error>) -> Result<Response, Error> {
let Self {
pipeline,
route,
admission,
deployment,
started,
} = self;
let (deployment, upstream_started_at) = deployment
.into_inner()
.unwrap_or_else(|e| e.into_inner())
.unzip();
pub(crate) async fn finish(
mut self,
result: Result<Response, Error>,
) -> Result<Response, Error> {
let mut response = match result {
Ok(response) => response,
Err(error) => {
let outcome = Outcome {
deployment,
status: error.status().as_u16(),
usage: None,
cost: None,
time_to_first_byte: None,
upstream_started_at,
ended_at: SystemTime::now(),
interrupted: false,
};
pipeline.finish(admission, outcome, None);
self.fail(&error);
return Err(error);
}
};
let Some(admission) = self.admission.take() else {
return Ok(response);
};
let (deployment, upstream_started_at) = self.take_deployment();
let (pipeline, route, started) = (self.pipeline.clone(), self.route, self.started);
if response.status().is_success() {
let headers = response.headers_mut();
for (name, value) in admission.edit().response_headers() {
@ -295,6 +320,13 @@ enum Meter {
Body(Vec<u8>),
}
impl Drop for Tracking {
fn drop(&mut self) {
// 499, the status the proxy's logs give a request the client closed.
self.record_failure(499, true);
}
}
/// Records the request when the response body ends or is dropped.
struct Finish {
pipeline: Arc<Pipeline>,

47
litellm-rust/crates/gateway-inference/src/lib.rs
View file

@ -3,7 +3,6 @@
//! Authentication, rate limiting and logging are the mounting server's layers; this crate
//! maps a public model name to its deployment and runs the core route.
mod audio_transcription;
mod caching;
mod chat_completions;
mod client_side;
@ -11,9 +10,7 @@ mod error;
mod hooks;
pub mod messages;
pub mod observe;
mod ocr;
mod request;
mod responses;
use std::sync::{Arc, OnceLock};
@ -25,13 +22,7 @@ use litellm_http::{
use litellm_inference::resources::CoreResources;
use litellm_inference_chat::ChatCompletionsRoute;
use litellm_inference_messages::MessagesRoute;
use litellm_inference_ocr::OcrRoute;
use litellm_inference_responses::ResponsesRoute;
use litellm_inference_transcription::AudioTranscriptionRoute;
use litellm_llms::base_llm::{
model_info::ModelInfo,
ocr::{handler::OcrClient, settings::OcrSettings},
};
use litellm_llms::base_llm::model_info::ModelInfo;
use litellm_secrets::source::SecretSource;
pub use error::Error;
@ -45,11 +36,8 @@ pub struct Gateway {
model_info: OnceLock<Arc<dyn ModelInfo>>,
drop_params: bool,
allow_client_side_credentials: bool,
pub audio_transcription: AudioTranscriptionRoute,
pub chat_completions: ChatCompletionsRoute,
pub messages: MessagesRoute,
pub ocr: OcrRoute,
pub responses: ResponsesRoute,
pub models: ModelRouter,
pub secrets: Arc<dyn SecretSource>,
pub resources: CoreResources,
@ -132,11 +120,6 @@ impl Gateway {
model_info: OnceLock::new(),
drop_params: false,
allow_client_side_credentials: false,
audio_transcription: AudioTranscriptionRoute::new(
provider.clone(),
auth.clone(),
secrets.clone(),
),
chat_completions: ChatCompletionsRoute::new(
provider.clone(),
auth.clone(),
@ -147,16 +130,7 @@ impl Gateway {
&http,
UrlPolicy::default(),
)?),
messages: MessagesRoute::new(provider.clone(), auth.clone(), secrets.clone()),
responses: ResponsesRoute::new(provider, auth.clone(), secrets.clone()),
ocr: OcrRoute::new(OcrClient::new(
&resources.pool,
&http,
UrlPolicy::default(),
auth,
OcrSettings::default(),
secrets.clone(),
)?),
messages: MessagesRoute::new(provider, auth, secrets.clone()),
models,
secrets,
resources,
@ -168,19 +142,18 @@ impl Gateway {
pub fn router(gateway: Arc<Gateway>) -> Router {
Router::new()
.route("/v1/messages", post(messages::create))
.route("/ocr", post(ocr::create))
.route("/v1/ocr", post(ocr::create))
.route("/chat/completions", post(chat_completions::create))
.route("/v1/chat/completions", post(chat_completions::create))
.nest("/engines/{model}", model_routes())
.nest("/openai/deployments/{model}", model_routes())
.route("/audio/transcriptions", post(audio_transcription::create))
.route(
"/v1/audio/transcriptions",
post(audio_transcription::create),
)
.route("/responses", post(responses::create))
.route("/v1/responses", post(responses::create))
// Not served: these routes skip the hooks (budgets, rate limits, spend), so a key could
// spend past its limits through them.
.route("/responses", post(request::unsupported))
.route("/v1/responses", post(request::unsupported))
.route("/audio/transcriptions", post(request::unsupported))
.route("/v1/audio/transcriptions", post(request::unsupported))
.route("/ocr", post(request::unsupported))
.route("/v1/ocr", post(request::unsupported))
.route("/embeddings", post(request::unsupported))
.route("/v1/embeddings", post(request::unsupported))
.route("/completions", post(request::unsupported))

10
litellm-rust/crates/gateway-inference/src/messages.rs
View file

@ -61,7 +61,15 @@ async fn handle(
&name,
)
.await?;
let (mut body, cache_options) = crate::caching::prepare(identity, object)?;
let (mut body, cache_options) = match crate::caching::prepare(identity, object) {
Ok(prepared) => prepared,
Err(error) => {
if let Some(tracking) = tracking {
tracking.fail(&error);
}
return Err(error);
}
};
let edit = |body: &mut Map<String, Value>| {
if let Some(tracking) = &tracking {
tracking.edit_body(body);

84
litellm-rust/crates/gateway-inference/src/ocr.rs
View file

@ -1,84 +0,0 @@
use litellm_gateway_auth::AuthenticatedRequest;
use std::sync::Arc;
use axum::{Json, extract::State, http::HeaderMap, response::IntoResponse};
use litellm_auth::SecretValue;
use litellm_inference_ocr::types::{LiteLLMOcrRequest, OcrConnectionInputs, OcrDocumentInput};
use litellm_llms::base_llm::ocr::transformation::decode_request_value;
use litellm_llms_types::formats::ocr::OcrDocument;
use serde_json::Value;
use crate::{
Error, Gateway,
request::{self, InferenceBody},
};
pub(crate) async fn create(
State(gateway): State<Arc<Gateway>>,
identity: AuthenticatedRequest,
headers: HeaderMap,
body: InferenceBody,
) -> Result<impl IntoResponse, Error> {
handle(&gateway, &identity, &headers, body).await.map(Json)
}
async fn handle(
gateway: &Gateway,
identity: &AuthenticatedRequest,
headers: &HeaderMap,
InferenceBody {
fields: body,
upload,
}: InferenceBody,
) -> Result<Value, Error> {
let header_format = headers
.get("x-req-format")
.and_then(|value| value.to_str().ok())
.map(str::to_owned);
gateway.check_body(&body)?;
let deployment = request::resolve_deployment(gateway, identity, &body)?;
request::authorize_model(identity, deployment, &body).await?;
let document = match upload {
Some(upload) => OcrDocumentInput::Bytes {
bytes: upload.bytes,
file_name: upload.file_name,
mime_type: upload.mime_type,
},
None => decode_request_value::<OcrDocument>(
body.get("document").cloned().unwrap_or_default(),
"document",
)?
.into(),
};
let format = body
.get("req_format")
.filter(|value| !value.is_null())
.cloned()
.or_else(|| header_format.map(Value::String));
let format = format.map(|value| match value {
Value::String(value) => Value::String(value.trim().to_ascii_lowercase()),
value => value,
});
let options = body
.into_iter()
.filter(|(name, _)| !matches!(name.as_str(), "model" | "document" | "req_format"))
.chain(format.map(|value| ("req_format".into(), value)))
.collect();
let call = LiteLLMOcrRequest::from_inputs(
deployment.model.clone(),
document,
deployment.custom_llm_provider.as_deref(),
options,
OcrConnectionInputs {
api_key: deployment.api_key.clone().map(SecretValue::new),
api_base: deployment.api_base.clone(),
timeout: deployment.timeout,
..Default::default()
},
)?;
let response = gateway.ocr.execute(call, &(), None).await?;
match response.provider_native_response {
Some(native) => Ok(Value::Object(native)),
None => Ok(response.into_json()),
}
}

127
litellm-rust/crates/gateway-inference/src/request.rs
View file

@ -1,6 +1,6 @@
use axum::{
body::{Body, Bytes, HttpBody},
extract::{FromRequest, FromRequestParts, Multipart, OriginalUri, Request},
extract::{FromRequest, FromRequestParts, OriginalUri, Request},
http::request::Parts,
response::{IntoResponse, Response},
};
@ -13,12 +13,6 @@ 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(crate) struct Upload {
pub bytes: Bytes,
pub file_name: Option<String>,
pub mime_type: Option<String>,
}
pub struct JsonObject(pub Map<String, Value>);
/// A JSON object body and the bytes it was parsed from, for a handler that parses it again.
@ -70,38 +64,6 @@ impl<S: Send + Sync> FromRequest<S> for JsonBody {
}
}
pub(crate) struct InferenceBody {
pub fields: Map<String, Value>,
pub upload: Option<Upload>,
}
impl<S: Send + Sync> FromRequest<S> for InferenceBody {
type Rejection = Error;
async fn from_request(request: Request, state: &S) -> Result<Self, Self::Rejection> {
let multipart = request
.headers()
.get("content-type")
.and_then(|header| header.to_str().ok())
.is_some_and(|value| {
value
.to_ascii_lowercase()
.starts_with("multipart/form-data")
});
if !multipart {
let JsonObject(fields) = JsonObject::from_request(request, state).await?;
return Ok(Self {
fields,
upload: None,
});
}
let multipart = Multipart::from_request(request, state)
.await
.map_err(|error| Error::InvalidBody(error.to_string()))?;
parse_multipart(multipart).await
}
}
pub(crate) fn object(body: &[u8]) -> Result<Map<String, Value>, Error> {
match serde_json::from_slice(body) {
Ok(Value::Object(body)) => Ok(body),
@ -110,14 +72,6 @@ pub(crate) fn object(body: &[u8]) -> Result<Map<String, Value>, Error> {
}
}
pub(crate) async fn authorize_model(
identity: &litellm_gateway_auth::AuthenticatedRequest,
deployment: &Deployment,
body: &Map<String, Value>,
) -> Result<(), Error> {
authorize(identity, model_name(body)?, deployment).await
}
/// Checks that the caller may use the public model `name` served by `deployment`.
pub(crate) async fn authorize(
identity: &litellm_gateway_auth::AuthenticatedRequest,
@ -163,27 +117,6 @@ pub(crate) fn usable(
|deployment| identity.may_use_deployment(deployment.team_id.as_deref())
}
/// The error for a group none of whose deployments the caller may use.
pub(crate) fn no_usable_deployment(body: &Map<String, Value>) -> Error {
Error::UnknownModel(model_name(body).unwrap_or_default().to_owned())
}
/// One deployment the caller may use of the group a request names, for endpoints served
/// without retries.
pub(crate) fn resolve_deployment<'a>(
gateway: &'a Gateway,
identity: &litellm_gateway_auth::AuthenticatedRequest,
body: &Map<String, Value>,
) -> Result<&'a Deployment, Error> {
let model = model_name(body)?;
let deployment = gateway
.models
.select_where(resolve_group(gateway, body)?, usable(identity))
.ok_or_else(|| no_usable_deployment(body))?;
crate::observe::record(model, deployment);
Ok(deployment)
}
/// A routed response that keeps its deployment marked in flight until its body is done.
/// A body of known length is already complete, so its mark is released at once.
pub(crate) fn hold(routed: Routed<Response>) -> Response {
@ -199,44 +132,6 @@ pub(crate) fn hold(routed: Routed<Response>) -> Response {
})
}
async fn parse_multipart(mut multipart: Multipart) -> Result<InferenceBody, Error> {
let mut fields = Map::new();
let mut upload = None;
while let Some(field) = multipart.next_field().await.map_err(multipart_error)? {
let name = field.name().unwrap_or_default().to_owned();
if name == "file" {
let file_name = field.file_name().map(str::to_owned);
let mime_type = field
.content_type()
.and_then(|value| value.split(';').next())
.map(str::trim)
.filter(|value| *value != "application/octet-stream")
.map(str::to_owned);
let bytes = field.bytes().await.map_err(multipart_error)?;
if bytes.len() > MAX_FILE_BYTES {
return Err(Error::BodyTooLarge);
}
upload = Some(Upload {
bytes,
file_name,
mime_type,
});
} else if name != "document" {
let text = field.text().await.map_err(multipart_error)?;
let value = serde_json::from_str(&text).unwrap_or(Value::String(text));
fields.insert(name, value);
}
}
Ok(InferenceBody { fields, upload })
}
fn multipart_error(error: axum::extract::multipart::MultipartError) -> Error {
if error.status() == axum::http::StatusCode::PAYLOAD_TOO_LARGE {
return Error::BodyTooLarge;
}
Error::InvalidBody(error.to_string())
}
pub(crate) async fn unsupported(OriginalUri(uri): OriginalUri) -> impl IntoResponse {
Error::Unsupported(uri.path().to_owned())
}
@ -255,28 +150,20 @@ mod tests {
use super::*;
#[fixture]
fn limited_uploads() -> Router {
fn limited_bodies() -> Router {
Router::new()
.route("/", post(|_: InferenceBody| async {}))
.route("/", post(|_: JsonBody| async {}))
.layer(DefaultBodyLimit::max(64))
}
#[rstest]
#[case::json("application/json", format!("{{\"text\":\"{}\"}}", "x".repeat(64)))]
#[case::multipart(
"multipart/form-data; boundary=test",
format!("--test\r\nContent-Disposition: form-data; name=\"file\"; filename=\"file.pdf\"\r\n\r\n{}\r\n--test--\r\n", "x".repeat(64)),
)]
#[tokio::test]
async fn uploads_respect_the_configured_body_limit(
limited_uploads: Router,
#[case] content_type: &str,
#[case] payload: String,
) {
let response = limited_uploads
async fn bodies_respect_the_configured_limit(limited_bodies: Router) {
let payload = format!("{{\"text\":\"{}\"}}", "x".repeat(64));
let response = limited_bodies
.oneshot(
Request::post("/")
.header("content-type", content_type)
.header("content-type", "application/json")
.body(Body::from(payload))
.unwrap(),
)

53
litellm-rust/crates/gateway-inference/src/responses.rs
View file

@ -1,53 +0,0 @@
use std::sync::Arc;
use axum::{Json, body::Bytes, extract::State, response::Response};
use litellm_gateway_auth::AuthenticatedRequest;
use litellm_host_http::Sse;
use litellm_inference_responses::{route::Responses, types::ResponsesCall};
use serde_json::json;
use crate::{Error, Gateway, JsonObject, request};
pub(crate) async fn create(
State(gateway): State<Arc<Gateway>>,
identity: AuthenticatedRequest,
JsonObject(body): JsonObject,
) -> Result<Response, Error> {
gateway.check_body(&body)?;
let deployment = request::resolve_deployment(&gateway, &identity, &body)?;
request::authorize_model(&identity, deployment, &body).await?;
let (body, cache_options) = crate::caching::prepare(&identity, body)?;
let route = gateway.responses.clone();
let route = match &gateway.cache {
Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new(
cache.clone(),
cache_options.scope.clone(),
)),
None => route,
};
let call = ResponsesCall {
model: deployment.model.clone(),
input: body.get("input").cloned().unwrap_or_default(),
optional_params: body
.into_iter()
.filter(|(name, _)| !matches!(name.as_str(), "model" | "input"))
.collect(),
api_key: deployment.api_key.clone(),
api_base: deployment.api_base.clone(),
custom_llm_provider: deployment.custom_llm_provider.clone(),
extra_headers: None,
timeout: deployment.timeout,
};
let machine = route.machine(call, cache_options.policy);
let stream = Sse::<Responses, _, _>::new(Json, |error| {
let error = Error::from(error);
Bytes::from(format!(
"event: error\ndata: {}\n\n",
json!({"type": "error", "code": error.status().as_u16().to_string(), "message": error.to_string(), "param": null})
))
});
let headers = crate::caching::CacheHeaders::default();
let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?;
Ok(headers.apply(response))
}

46
litellm-rust/crates/gateway-inference/tests/caching.rs
View file

@ -11,12 +11,14 @@ use rstest::rstest;
use serde_json::{Value, json};
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method};
fn message() -> Value {
json!({"id":"message-1", "model":"test-model", "type":"message", "role":"assistant", "content":[{"type":"text","text":"hello"}], "stop_reason":"end_turn", "usage":{"input_tokens":1,"output_tokens":1}})
}
#[rstest]
#[case::chat("/v1/chat/completions", "anthropic/test-model", false)]
#[case::messages("/v1/messages", "anthropic/test-model", false)]
#[case::responses("/v1/responses", "openai/test-model", false)]
#[case::messages_stream("/v1/messages", "anthropic/test-model", true)]
#[case::responses_stream("/v1/responses", "openai/test-model", true)]
#[tokio::test]
async fn all_inference_endpoints_share_native_cache(
#[case] path: &str,
@ -25,18 +27,8 @@ async fn all_inference_endpoints_share_native_cache(
#[values("s-maxage", "s-max-age")] max_age: &str,
) {
let upstream = MockServer::start().await;
let is_responses = path.ends_with("responses");
let provider_body = if is_responses {
json!({"id":"response-1", "model":"test-model", "status":"completed", "output":[]})
} else {
json!({"id":"message-1", "model":"test-model", "type":"message", "role":"assistant", "content":[{"type":"text","text":"hello"}], "stop_reason":"end_turn", "usage":{"input_tokens":1,"output_tokens":1}})
};
let terminal = if is_responses {
"response.completed"
} else {
"message_stop"
};
let events = format!("event: {terminal}\ndata: {{\"type\":\"{terminal}\"}}\n\n");
let provider_body = message();
let events = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n".to_owned();
let template = if stream {
ResponseTemplate::new(200).set_body_raw(events.clone(), "text/event-stream")
} else {
@ -58,11 +50,7 @@ async fn all_inference_endpoints_share_native_cache(
}),
);
let app = support::app_with_cache(model, &upstream.uri(), cache.clone());
let request = if is_responses {
json!({"model":"public/model", "input":"hello", "stream":stream, "cache":{(max_age):600}})
} else {
json!({"model":"public/model", "messages":[{"role":"user","content":"hello"}], "max_tokens":16, "stream":stream, "cache":{(max_age):600}})
};
let request = json!({"model":"public/model", "messages":[{"role":"user","content":"hello"}], "max_tokens":16, "stream":stream, "cache":{(max_age):600}});
let first = support::post(app.clone(), path, request.clone()).await;
assert_eq!(first.status(), 200);
assert!(!first.headers().contains_key("x-litellm-cache-key"));
@ -139,9 +127,7 @@ async fn authenticated_callers_do_not_share_cached_responses(
let upstream = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"id":"response-1", "model":"test-model", "status":"completed", "output":[]
})))
.respond_with(ResponseTemplate::new(200).set_body_json(message()))
.expect(2)
.mount(&upstream)
.await;
@ -149,34 +135,34 @@ async fn authenticated_callers_do_not_share_cached_responses(
InMemoryCache::new(Some(100), Some(Duration::from_secs(60))),
)));
let first_caller = support::app_with_cache_for_principal(
"openai/test-model",
"anthropic/test-model",
&upstream.uri(),
cache.clone(),
Principal::new("issuer".into(), "tenant-a".into(), PrincipalKind::Service),
);
let second_caller = support::app_with_cache_for_principal(
"openai/test-model",
"anthropic/test-model",
&upstream.uri(),
cache,
Principal::new(authority.into(), subject.into(), PrincipalKind::Service),
);
let body = json!({"model":"public/model","input":"same prompt"});
let first = support::post(first_caller.clone(), "/v1/responses", body.clone()).await;
let body = json!({"model":"public/model","messages":[{"role":"user","content":"same prompt"}],"max_tokens":16});
let first = support::post(first_caller.clone(), "/v1/messages", body.clone()).await;
assert_eq!(first.status(), 200);
assert!(!first.headers().contains_key("x-litellm-cache-key"));
let first_hit = support::post(first_caller.clone(), "/v1/responses", body.clone()).await;
let first_hit = support::post(first_caller.clone(), "/v1/messages", body.clone()).await;
assert_eq!(first_hit.status(), 200);
let first_key = first_hit.headers().get("x-litellm-cache-key").unwrap();
let second = support::post(second_caller.clone(), "/v1/responses", body.clone()).await;
let second = support::post(second_caller.clone(), "/v1/messages", body.clone()).await;
assert_eq!(second.status(), 200);
assert!(!second.headers().contains_key("x-litellm-cache-key"));
let second_hit = support::post(second_caller, "/v1/responses", body.clone()).await;
let second_hit = support::post(second_caller, "/v1/messages", body.clone()).await;
assert_eq!(second_hit.status(), 200);
assert_ne!(
second_hit.headers().get("x-litellm-cache-key").unwrap(),
first_key
);
let first_again = support::post(first_caller, "/v1/responses", body).await;
let first_again = support::post(first_caller, "/v1/messages", body).await;
assert_eq!(first_again.status(), 200);
assert_eq!(
first_again.headers().get("x-litellm-cache-key"),

2
litellm-rust/crates/gateway-inference/tests/messages.rs
View file

@ -178,8 +178,6 @@ async fn hosted_provider_failure_preserves_status_body_and_request_id() {
#[case::messages("/v1/messages")]
#[case::chat("/v1/chat/completions")]
#[case::model_path("/engines/path-model/chat/completions")]
#[case::ocr("/ocr")]
#[case::transcription("/audio/transcriptions")]
#[tokio::test]
async fn model_permissions_prevent_provider_calls(#[case] path: &str) {
let upstream = MockServer::start().await;

189
litellm-rust/crates/gateway-inference/tests/ocr.rs
View file

@ -1,189 +0,0 @@
mod support;
use axum::{body::Body, http::Request};
use litellm_gateway_inference::Error;
use litellm_llms::base_llm::ocr::{error::Error as OcrError, transformation::decode_request_value};
use litellm_llms_types::formats::ocr::OcrDocument;
use rstest::rstest;
use serde_json::{Value, json};
use tower::ServiceExt;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_json, header, method, path},
};
const DOCUMENT: &str = "data:application/pdf;base64,YWJj";
#[rstest]
#[case("/ocr", false)]
#[case("/v1/ocr", true)]
#[tokio::test]
async fn json_and_multipart_reach_ocr_with_the_deployment(
#[case] route: &str,
#[case] multipart: bool,
) {
let upstream = MockServer::start().await;
Mock::given(method("POST")).and(path("/v1/ocr"))
.and(header("authorization", "Bearer test-key"))
.and(body_json(json!({"model": "test-ocr", "document": {"type": "document_url", "document_url": DOCUMENT}, "pages": [0]})))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"pages": [{"index": 0, "markdown": "recognized text"}]})))
.expect(1).mount(&upstream).await;
let app = support::app("mistral/test-ocr", &upstream.uri());
let response = if multipart {
let body = "--boundary\r\nContent-Disposition: form-data; name=\"model\"\r\n\r\npublic/model\r\n--boundary\r\nContent-Disposition: form-data; name=\"pages\"\r\n\r\n[0]\r\n--boundary\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.pdf\"\r\nContent-Type: application/pdf\r\n\r\nabc\r\n--boundary--\r\n";
app.oneshot(
Request::post(route)
.header("content-type", "multipart/form-data; boundary=boundary")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap()
} else {
support::post(app, route, json!({"model": "public/model", "document": {"type": "document_url", "document_url": DOCUMENT}, "pages": [0]})).await
};
assert_eq!(response.status(), 200);
let body = support::json(response).await;
assert_eq!(body["pages"][0]["markdown"], "recognized text");
assert_eq!(body["model"], "test-ocr");
}
#[rstest]
#[case(None, true)]
#[case(Some("litellm"), false)]
#[tokio::test]
async fn native_format_header_is_used_unless_the_body_overrides_it(
#[case] format: Option<&str>,
#[case] native: bool,
) {
let upstream = MockServer::start().await;
let payload = json!({"pages": [{"index": 0, "markdown": "text"}], "provider_only": true});
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200).set_body_json(payload.clone()))
.expect(1)
.mount(&upstream)
.await;
let body = json!({"model": "public/model", "document": {"type": "document_url", "document_url": DOCUMENT}, "req_format": format});
let response = support::app("mistral/test-ocr", &upstream.uri())
.oneshot(
Request::post("/ocr")
.header("x-req-format", " Native ")
.body(Body::from(body.to_string()))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), 200);
let body = support::json(response).await;
if native {
assert_eq!(body, payload);
} else {
assert_eq!(body["object"], "ocr");
assert_eq!(
body["pages"][0]["markdown"],
payload["pages"][0]["markdown"]
);
}
}
#[tokio::test]
async fn ocr_keeps_upstream_status_in_an_openai_error_envelope() {
let upstream = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(429).set_body_json(json!({"message": "busy"})))
.expect(1)
.mount(&upstream)
.await;
let response = support::post(support::app("mistral/test-ocr", &upstream.uri()), "/ocr",
json!({"model": "public/model", "document": {"type": "document_url", "document_url": DOCUMENT}})).await;
assert_eq!(response.status(), 429);
let body = support::json(response).await;
assert_eq!(body["error"]["code"], 429);
assert!(body["error"]["message"].as_str().unwrap().contains("busy"));
}
#[rstest]
#[case(json!({"model": "public/model"}))]
#[case(json!({"model": "public/model", "document": "/etc/passwd"}))]
#[case(json!({"model": "missing", "document": {"type": "document_url", "document_url": DOCUMENT}}))]
#[tokio::test]
async fn invalid_ocr_requests_do_not_call_the_provider(#[case] body: Value) {
let upstream = MockServer::start().await;
let response = support::post(
support::app("mistral/test-ocr", &upstream.uri()),
"/ocr",
body,
)
.await;
assert_eq!(response.status(), 400);
assert!(support::json(response).await["error"]["message"].is_string());
assert!(upstream.received_requests().await.unwrap().is_empty());
}
#[rstest]
#[tokio::test]
async fn malformed_multipart_uses_an_openai_error_envelope(
#[values("/v1/ocr", "/v1/audio/transcriptions")] route: &str,
) {
let upstream = MockServer::start().await;
let response = support::app("mistral/test-ocr", &upstream.uri())
.oneshot(
Request::post(route)
.header("content-type", "multipart/form-data")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), 400);
let body = support::json(response).await;
assert_eq!(body["error"]["type"], "invalid_request_error");
assert_eq!(body["error"]["code"], 400);
let error_message = body["error"]["message"].as_str().unwrap();
assert!(!error_message.is_empty());
assert!(upstream.received_requests().await.unwrap().is_empty());
}
#[rstest]
#[case::missing_document(
"/v1/ocr", "mistral/test-ocr", "",
Error::Ocr(decode_request_value::<OcrDocument>(Value::Null, "document").unwrap_err()),
)]
#[case::empty_document(
"/v1/ocr",
"mistral/test-ocr",
"--test\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.pdf\"\r\n\r\n\r\n",
Error::Ocr(OcrError::EmptyFile)
)]
#[case::empty_audio(
"/v1/audio/transcriptions", "bedrock/test-model",
"--test\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.wav\"\r\n\r\n\r\n",
Error::Route(litellm_llms::Error::MissingField("audio.data").into()),
)]
#[tokio::test]
async fn upload_validation_errors_come_from_core(
#[case] route: &str,
#[case] model: &str,
#[case] file: &str,
#[case] error: Error,
) {
let upstream = MockServer::start().await;
let payload = format!(
"--test\r\nContent-Disposition: form-data; name=\"model\"\r\n\r\npublic/model\r\n{file}--test--\r\n"
);
let response = support::app(model, &upstream.uri())
.oneshot(
Request::post(route)
.header("content-type", "multipart/form-data; boundary=test")
.body(Body::from(payload))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), 400);
assert_eq!(
support::json(response).await,
support::json(error.openai_response()).await
);
assert!(upstream.received_requests().await.unwrap().is_empty());
}

118
litellm-rust/crates/gateway-inference/tests/responses.rs
View file

@ -1,118 +0,0 @@
mod support;
use axum::body::to_bytes;
use rstest::rstest;
use serde_json::json;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_json, header, method, path},
};
#[rstest]
#[case::completed(false)]
#[case::streaming(true)]
#[tokio::test]
async fn responses_aliases_run_the_core_route(
#[case] stream: bool,
#[values("/responses", "/v1/responses")] route: &str,
) {
let upstream = MockServer::start().await;
let completed =
json!({"id": "response-1", "model": "test-model", "output": [], "provider_extra": true});
let events = format!(
"event: response.completed\ndata: {}\n\n",
json!({"type": "response.completed", "response": completed})
);
let template = if stream {
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(&events)
} else {
ResponseTemplate::new(200).set_body_json(&completed)
};
Mock::given(method("POST"))
.and(path("/responses"))
.and(header("authorization", "Bearer test-key"))
.and(body_json(json!({"model": "test-model", "input": "hello", "stream": stream, "metadata": {"caller": "test"}})))
.respond_with(template)
.expect(1)
.mount(&upstream).await;
let response = support::post(support::app("openai/test-model", &upstream.uri()), route,
json!({"model": "public/model", "input": "hello", "stream": stream, "metadata": {"caller": "test"}})).await;
assert_eq!(response.status(), 200);
if stream {
assert_eq!(response.headers()["content-type"], "text/event-stream");
assert_eq!(
to_bytes(response.into_body(), 4096).await.unwrap().as_ref(),
events.as_bytes()
);
} else {
assert_eq!(support::json(response).await, completed);
}
}
#[rstest]
#[case::missing_input(json!({"model": "public/model"}), 400)]
#[case::invalid_stream(json!({"model": "public/model", "input": "hello", "stream": "yes"}), 400)]
#[case::unknown_model(json!({"model": "unknown", "input": "hello"}), 400)]
#[tokio::test]
async fn invalid_responses_requests_never_reach_the_provider(
#[case] body: serde_json::Value,
#[case] status: u16,
) {
let upstream = MockServer::start().await;
let response = support::post(
support::app("openai/test-model", &upstream.uri()),
"/v1/responses",
body,
)
.await;
assert_eq!(response.status(), status);
assert_eq!(support::json(response).await["error"]["code"], status);
assert!(upstream.received_requests().await.unwrap().is_empty());
}
#[rstest]
#[tokio::test]
async fn responses_preserve_upstream_errors_without_retry() {
let upstream = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(
ResponseTemplate::new(429).set_body_json(json!({"error": {"message": "slow down"}})),
)
.expect(1)
.mount(&upstream)
.await;
let response = support::post(
support::app("openai/test-model", &upstream.uri()),
"/v1/responses",
json!({"model": "public/model", "input": "hello"}),
)
.await;
assert_eq!(response.status(), 429);
assert!(
support::json(response).await["error"]["message"]
.as_str()
.unwrap()
.contains("slow down")
);
}
#[rstest]
#[tokio::test]
async fn responses_authorize_the_public_model_before_provider_execution() {
let upstream = MockServer::start().await;
let app = support::app_with_permissions(
"openai/test-model",
&upstream.uri(),
litellm_gateway_auth::Permissions::None,
);
let response = support::post(
app,
"/v1/responses",
json!({"model": "public/model", "input": "hello"}),
)
.await;
assert_eq!(response.status(), 403);
assert!(upstream.received_requests().await.unwrap().is_empty());
}

81
litellm-rust/crates/gateway-inference/tests/routes.rs
View file

@ -131,30 +131,17 @@ async fn chat_errors_come_from_core(
assert!(upstream.received_requests().await.unwrap().is_empty());
}
#[rstest]
#[case::audio("/v1/audio/transcriptions", "bedrock/test-model", "audio")]
#[case::document("/v1/ocr", "mistral/test-ocr", "document")]
#[tokio::test]
async fn missing_inference_fields_use_the_same_validation_as_null(
#[case] path: &str,
#[case] model: &str,
#[case] field: &str,
) {
let upstream = MockServer::start().await;
let app = support::app(model, &upstream.uri());
let missing = support::post(app.clone(), path, json!({"model": "public/model"})).await;
let null = support::post(app, path, json!({"model": "public/model", field: null})).await;
assert_eq!(missing.status(), 400);
assert_eq!(missing.status(), null.status());
assert_eq!(support::json(missing).await, support::json(null).await);
assert!(upstream.received_requests().await.unwrap().is_empty());
}
#[rstest]
#[case("/embeddings")]
#[case("/v1/embeddings")]
#[case("/completions")]
#[case("/v1/completions")]
#[case("/responses")]
#[case("/v1/responses")]
#[case("/audio/transcriptions")]
#[case("/v1/audio/transcriptions")]
#[case("/ocr")]
#[case("/v1/ocr")]
#[case("/engines/public%2Fmodel/embeddings")]
#[case("/openai/deployments/public%2Fmodel/completions")]
#[tokio::test]
@ -172,66 +159,10 @@ async fn unimplemented_routes_return_an_explicit_error(#[case] path: &str) {
assert!(message.contains(path));
}
#[rstest]
#[case("/audio/transcriptions")]
#[case("/v1/audio/transcriptions")]
#[tokio::test]
async fn transcription_aliases_reach_core_validation(#[case] path: &str) {
let response = support::post(
support::app("bedrock/test-model", "http://127.0.0.1:1"),
path,
json!({"model": "public/model", "audio": {"data": "YWJj", "format": "invalid"}}),
)
.await;
assert_eq!(response.status(), 400);
let body: Value = support::json(response).await;
assert!(
body["error"]["message"]
.as_str()
.unwrap()
.contains("audio.format")
);
}
#[rstest]
#[case::no_extension("test")]
#[case::unsupported_extension("test.invalid")]
#[tokio::test]
async fn upload_audio_format_validation_matches_core(#[case] filename: &str) {
let upstream = MockServer::start().await;
let app = support::app("bedrock/test-model", &upstream.uri());
let path = "/v1/audio/transcriptions";
let expected = support::post(
app.clone(),
path,
json!({"model": "public/model", "audio": {"data": "YWJj", "format": null}}),
)
.await;
let payload = format!(
"--test\r\nContent-Disposition: form-data; name=\"model\"\r\n\r\npublic/model\r\n\
--test\r\nContent-Disposition: form-data; name=\"file\"; filename=\"{filename}\"\r\n\r\nabc\r\n--test--\r\n"
);
let response = app
.oneshot(
Request::post(path)
.header("content-type", "multipart/form-data; boundary=test")
.body(Body::from(payload))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), 400);
assert_eq!(support::json(response).await, support::json(expected).await);
assert!(upstream.received_requests().await.unwrap().is_empty());
}
#[rstest]
#[case::chat("/v1/chat/completions", false)]
#[case::responses("/v1/responses", false)]
#[case::deployment("/openai/deployments/public%2Fmodel/chat/completions", false)]
#[case::messages("/v1/messages", true)]
#[case::ocr("/v1/ocr", false)]
#[case::transcription("/v1/audio/transcriptions", false)]
#[tokio::test]
async fn json_extraction_rejections_use_the_endpoint_error_envelope(
#[case] path: &str,

2
litellm-rust/crates/gateway/src/budgets.rs
View file

@ -952,7 +952,9 @@ impl Settle for Reservation {
};
let mut keys = Vec::new();
let mut changes = Vec::new();
let keep = record.cut_off_after_dispatch();
for (counter, charged) in self.plan.counters.iter().zip(&self.charged) {
let cost = if keep { *charged } else { cost };
let delta = cost - charged;
if delta == 0.0 && *charged == 0.0 {
continue;

5
litellm-rust/crates/gateway/src/rate_limits.rs
View file

@ -1416,12 +1416,13 @@ impl Settle for Held {
tokens,
} = *self;
let charge = tokens.and_then(|tokens| {
// The proxy counts every token but cache reads. Without usage, a success keeps its
// reservation and a failure returns it.
// The proxy counts every token but cache reads. Without usage, a success or a request
// cut off after it reached the provider keeps its reservation, and a failure returns it.
let actual = match (&record.outcome.usage, record.succeeded()) {
(Some(usage), _) => i64::try_from(usage.total_tokens() - usage.cache_read_tokens)
.unwrap_or(i64::MAX),
(None, true) => tokens.reserved,
(None, false) if record.cut_off_after_dispatch() => tokens.reserved,
(None, false) => 0,
};
let delta = actual - tokens.reserved;

69
litellm-rust/crates/gateway/tests/budgets.rs
View file

@ -41,10 +41,15 @@ async fn upstream() -> (String, Arc<AtomicUsize>) {
let counted = calls.clone();
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = Router::new().fallback(move |_: Request<Body>| {
let app = Router::new().fallback(move |request: Request<Body>| {
let counted = counted.clone();
async move {
counted.fetch_add(1, Ordering::SeqCst);
// A request whose prompt says "hold" stays with the provider until the client leaves.
let body = to_bytes(request.into_body(), 1 << 20).await.unwrap();
if body.windows(4).any(|window| window == b"hold") {
tokio::time::sleep(Duration::from_secs(60)).await;
}
([(CONTENT_TYPE, "application/json")], MESSAGE).into_response()
}
});
@ -160,6 +165,10 @@ async fn send_as(app: &Router, token: &str, model: &str, end_user: Option<&str>)
if let Some(end_user) = end_user {
body["metadata"] = json!({"user_id": end_user});
}
send_body(app, token, body).await
}
async fn send_body(app: &Router, token: &str, body: Value) -> (u16, Value) {
let response = app
.clone()
.oneshot(
@ -340,6 +349,64 @@ async fn reservations_hold_back_concurrent_requests() {
assert_eq!(statuses, [200, 422]);
}
#[rstest]
#[tokio::test]
async fn requests_refused_after_admission_return_their_reservation() {
// The reservation takes the whole budget; a body the gateway rejects after the hooks ran
// must hand it back, or the key stays blocked.
let key = key(
None,
KeyBudgets {
key: budget(0.0, 0.0001),
..KeyBudgets::default()
},
);
let gateway = gateway(vec![("sk-key", key)], None, "").await;
let invalid = json!({
"model": "claude",
"max_tokens": 8,
"messages": [{"role": "user", "content": "hi"}],
"cache": {"ttl": -1},
});
let (status, body) = send_body(&gateway.app, "sk-key", invalid).await;
assert_eq!(status, 400, "{body}");
settled().await;
let (status, body) = send(&gateway.app, "sk-key", "claude").await;
assert_eq!(status, 200, "{body}");
}
#[rstest]
#[tokio::test]
async fn requests_cut_off_after_reaching_the_provider_keep_their_reservation() {
// The provider may bill a request the client walked away from, so its reservation is kept
// rather than refunded: the next request finds the budget taken.
let key = key(
None,
KeyBudgets {
key: budget(0.0, 0.0001),
..KeyBudgets::default()
},
);
let gateway = gateway(vec![("sk-key", key)], None, "").await;
let held = json!({
"model": "claude",
"max_tokens": 8,
"messages": [{"role": "user", "content": "hold"}],
});
let request = tokio::spawn({
let app = gateway.app.clone();
async move { send_body(&app, "sk-key", held).await }
});
while gateway.calls.load(Ordering::SeqCst) == 0 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
request.abort();
let _ = request.await;
settled().await;
let (status, body) = send(&gateway.app, "sk-key", "claude").await;
assert_eq!(status, 422, "{body}");
}
fn redis_url() -> Option<String> {
std::env::var("LITELLM_TEST_REDIS_URL").ok()
}