diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5c0401c0895..e52e979e792 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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", diff --git a/litellm-rust/crates/gateway-hooks/src/facts.rs b/litellm-rust/crates/gateway-hooks/src/facts.rs index b0ed1c58fe2..90f0e005541 100644 --- a/litellm-rust/crates/gateway-hooks/src/facts.rs +++ b/litellm-rust/crates/gateway-hooks/src/facts.rs @@ -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 diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index e589eb75719..5b112f791ec 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs deleted file mode 100644 index 2640f1bf800..00000000000 --- a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs +++ /dev/null @@ -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>, - identity: AuthenticatedRequest, - body: InferenceBody, -) -> Result { - handle(&gateway, &identity, body).await.map(Json) -} - -async fn handle( - gateway: &Gateway, - identity: &AuthenticatedRequest, - InferenceBody { - fields: body, - upload, - }: InferenceBody, -) -> Result { - 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?) -} diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 87380ead495..95d911cca29 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -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| { if let Some(tracking) = &tracking { tracking.edit_body(body); diff --git a/litellm-rust/crates/gateway-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs index d03cc2b83c4..649505d08e5 100644 --- a/litellm-rust/crates/gateway-inference/src/hooks.rs +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -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, route: Route, - admission: Admission, + admission: Option, /// The deployment the latest attempt went to, and when it was sent. - deployment: Mutex>, + deployment: Mutex, 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) -> Option<&str> { impl Tracking { /// Applies the hooks' body changes to an attempt's body. pub(crate) fn edit_body(&self, body: &mut Map) { - 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> { - 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, Option) { + 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) -> Result { - 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, + ) -> Result { 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), } +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, diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 010e92cfbe1..2b210e23e9b 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -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>, 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, 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) -> 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)) diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index 800cabe7659..c0919f41673 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -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| { if let Some(tracking) = &tracking { tracking.edit_body(body); diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs deleted file mode 100644 index c12275b8469..00000000000 --- a/litellm-rust/crates/gateway-inference/src/ocr.rs +++ /dev/null @@ -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>, - identity: AuthenticatedRequest, - headers: HeaderMap, - body: InferenceBody, -) -> Result { - handle(&gateway, &identity, &headers, body).await.map(Json) -} - -async fn handle( - gateway: &Gateway, - identity: &AuthenticatedRequest, - headers: &HeaderMap, - InferenceBody { - fields: body, - upload, - }: InferenceBody, -) -> Result { - 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::( - 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()), - } -} diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs index c1fa1aa7ca1..07730b8d7bb 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -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, - pub mime_type: Option, -} - pub struct JsonObject(pub Map); /// A JSON object body and the bytes it was parsed from, for a handler that parses it again. @@ -70,38 +64,6 @@ impl FromRequest for JsonBody { } } -pub(crate) struct InferenceBody { - pub fields: Map, - pub upload: Option, -} - -impl FromRequest for InferenceBody { - type Rejection = Error; - - async fn from_request(request: Request, state: &S) -> Result { - 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, Error> { match serde_json::from_slice(body) { Ok(Value::Object(body)) => Ok(body), @@ -110,14 +72,6 @@ pub(crate) fn object(body: &[u8]) -> Result, Error> { } } -pub(crate) async fn authorize_model( - identity: &litellm_gateway_auth::AuthenticatedRequest, - deployment: &Deployment, - body: &Map, -) -> 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) -> 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, -) -> 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 { @@ -199,44 +132,6 @@ pub(crate) fn hold(routed: Routed) -> Response { }) } -async fn parse_multipart(mut multipart: Multipart) -> Result { - 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(), ) diff --git a/litellm-rust/crates/gateway-inference/src/responses.rs b/litellm-rust/crates/gateway-inference/src/responses.rs deleted file mode 100644 index c26c9ed877c..00000000000 --- a/litellm-rust/crates/gateway-inference/src/responses.rs +++ /dev/null @@ -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>, - identity: AuthenticatedRequest, - JsonObject(body): JsonObject, -) -> Result { - 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::::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)) -} diff --git a/litellm-rust/crates/gateway-inference/tests/caching.rs b/litellm-rust/crates/gateway-inference/tests/caching.rs index 143bc1ba6db..a8879731586 100644 --- a/litellm-rust/crates/gateway-inference/tests/caching.rs +++ b/litellm-rust/crates/gateway-inference/tests/caching.rs @@ -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"), diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs index bf11e67e0fc..78cb342bc86 100644 --- a/litellm-rust/crates/gateway-inference/tests/messages.rs +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -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; diff --git a/litellm-rust/crates/gateway-inference/tests/ocr.rs b/litellm-rust/crates/gateway-inference/tests/ocr.rs deleted file mode 100644 index fbe523addf0..00000000000 --- a/litellm-rust/crates/gateway-inference/tests/ocr.rs +++ /dev/null @@ -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::(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()); -} diff --git a/litellm-rust/crates/gateway-inference/tests/responses.rs b/litellm-rust/crates/gateway-inference/tests/responses.rs deleted file mode 100644 index ded7d73c5f3..00000000000 --- a/litellm-rust/crates/gateway-inference/tests/responses.rs +++ /dev/null @@ -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()); -} diff --git a/litellm-rust/crates/gateway-inference/tests/routes.rs b/litellm-rust/crates/gateway-inference/tests/routes.rs index 7452306338e..6d37f57746f 100644 --- a/litellm-rust/crates/gateway-inference/tests/routes.rs +++ b/litellm-rust/crates/gateway-inference/tests/routes.rs @@ -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, diff --git a/litellm-rust/crates/gateway/src/budgets.rs b/litellm-rust/crates/gateway/src/budgets.rs index cf9b64db401..f817dcfe745 100644 --- a/litellm-rust/crates/gateway/src/budgets.rs +++ b/litellm-rust/crates/gateway/src/budgets.rs @@ -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; diff --git a/litellm-rust/crates/gateway/src/rate_limits.rs b/litellm-rust/crates/gateway/src/rate_limits.rs index 7ae2ab54e71..3d742edef93 100644 --- a/litellm-rust/crates/gateway/src/rate_limits.rs +++ b/litellm-rust/crates/gateway/src/rate_limits.rs @@ -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; diff --git a/litellm-rust/crates/gateway/tests/budgets.rs b/litellm-rust/crates/gateway/tests/budgets.rs index 83a7778c732..1277c067927 100644 --- a/litellm-rust/crates/gateway/tests/budgets.rs +++ b/litellm-rust/crates/gateway/tests/budgets.rs @@ -41,10 +41,15 @@ async fn upstream() -> (String, Arc) { 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| { + let app = Router::new().fallback(move |request: Request| { 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 { std::env::var("LITELLM_TEST_REDIS_URL").ok() }