diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 20bb696c46c..5605bcd7271 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3754,6 +3754,7 @@ dependencies = [ "futures-util", "http 1.4.2", "http-body-util", + "litellm-auth-aws", "litellm-auth-types", "litellm-config", "litellm-db", @@ -3770,11 +3771,13 @@ dependencies = [ "litellm-router", "litellm-secrets", "litellm-tracing", + "md-5", "rand 0.8.7", "redis", "rstest", "serde", "serde_json", + "sha2 0.10.9", "sqlx", "strum", "tempfile", @@ -3785,6 +3788,7 @@ dependencies = [ "tower-sessions-moka-store", "tracing", "uuid", + "wiremock", ] [[package]] diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 8fdbba21711..5a17e182a72 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -8,7 +8,8 @@ use sha2::{Digest, Sha256}; use aws_credential_types::Credentials; use aws_credential_types::provider::ProvideCredentials; use aws_sigv4::http_request::{ - SignableBody, SignableRequest, SigningParams, SigningSettings, sign, + PercentEncodingMode, SignableBody, SignableRequest, SigningParams, SigningSettings, + UriPathNormalizationMode, sign, }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; @@ -464,6 +465,58 @@ pub fn sign_post( service: &str, credentials: &Credentials, signing_time: SystemTime, +) -> Result, Error> { + sign_request( + "POST", + url, + body, + headers, + region, + service, + credentials, + signing_time, + SigningSettings::default(), + ) +} + +/// Signs an S3 request as botocore's `S3SigV4Auth` does: the path is signed as sent, encoded +/// once and not normalized, and every header given is covered. +pub fn sign_s3( + method: &str, + url: &str, + body: &[u8], + headers: &BTreeMap, + region: &str, + credentials: &Credentials, + signing_time: SystemTime, +) -> Result, Error> { + let mut settings = SigningSettings::default(); + settings.percent_encoding_mode = PercentEncodingMode::Single; + settings.uri_path_normalization_mode = UriPathNormalizationMode::Disabled; + sign_request( + method, + url, + body, + headers, + region, + "s3", + credentials, + signing_time, + settings, + ) +} + +#[allow(clippy::too_many_arguments)] +fn sign_request( + method: &str, + url: &str, + body: &[u8], + headers: &BTreeMap, + region: &str, + service: &str, + credentials: &Credentials, + signing_time: SystemTime, + settings: SigningSettings, ) -> Result, Error> { let identity: Identity = credentials.clone().into(); let params = v4::SigningParams::builder() @@ -471,14 +524,14 @@ pub fn sign_post( .region(region) .name(service) .time(signing_time) - .settings(SigningSettings::default()) + .settings(settings) .build() .map(SigningParams::from) .map_err(|error| Error::AwsSigningParameters(error.to_string()))?; let header_refs = headers .iter() .map(|(name, value)| (name.as_str(), value.as_str())); - let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body)) + let request = SignableRequest::new(method, url, header_refs, SignableBody::Bytes(body)) .map_err(|error| Error::AwsSignableRequest(error.to_string()))?; let (instructions, _) = sign(request, ¶ms) .map_err(|error| Error::AwsSigning(error.to_string()))? @@ -935,6 +988,49 @@ mod tests { ); } + /// The PUT Object example of the S3 SigV4 guide, "Signature Calculations for the + /// Authorization Header: Transferring Payload in a Single Chunk" (checked 2026-10). + #[test] + fn s3_requests_are_signed_like_the_s3_guide_example() { + let headers = BTreeMap::from([ + ( + "Date".to_string(), + "Fri, 24 May 2013 00:00:00 GMT".to_string(), + ), + ( + "x-amz-content-sha256".to_string(), + "44ce7dd67c959e0d3524ffac1771dfbba87d2b6b4b4e99e42034a8b803f8b072".to_string(), + ), + ( + "x-amz-storage-class".to_string(), + "REDUCED_REDUNDANCY".to_string(), + ), + ]); + let signed = sign_s3( + "PUT", + "https://examplebucket.s3.amazonaws.com/test%24file.text", + b"Welcome to Amazon S3.", + &headers, + "us-east-1", + &Credentials::new( + "AKIAIOSFODNN7EXAMPLE", + "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + None, + None, + "test", + ), + UNIX_EPOCH + Duration::from_secs(1_369_353_600), + ) + .unwrap(); + + assert_eq!( + signed["Authorization"], + "AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/20130524/us-east-1/s3/aws4_request, \ + SignedHeaders=date;host;x-amz-content-sha256;x-amz-date;x-amz-storage-class, \ + Signature=98ad721746da40c64f1a55b78f14c238d841ea1380cd77a1b5971af0ece108bd" + ); + } + #[test] fn same_role_comparison_matches_partition_account_and_role() { assert!(same_role_arns( diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index d7e2e59bb53..e1ca8682cab 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -21,7 +21,11 @@ litellm-gateway-auth.workspace = true litellm-gateway-hooks.workspace = true litellm-model-catalog.workspace = true litellm-gateway-ui.workspace = true +litellm-auth-aws.workspace = true litellm-auth-types.workspace = true +base64.workspace = true +md-5 = "0.11.0" +sha2.workspace = true litellm-config.workspace = true litellm-db.workspace = true litellm-http = { workspace = true, features = ["mcp"] } @@ -39,7 +43,7 @@ uuid.workspace = true tower-sessions-moka-store = "0.15.0" [dev-dependencies] -base64.workspace = true +wiremock.workspace = true litellm-db-testing.workspace = true sqlx.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index fbab39ef426..742a036e69d 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -22,6 +22,8 @@ pub enum Error { Secret(#[from] litellm_secrets::Error), #[error("litellm_settings.hooks: {0}")] Hooks(String), + #[error("litellm_settings.s3_callback_params: {0}")] + S3Log(String), #[error("coordination Redis: {0}")] Redis(String), #[error("MCP requires a configured master key")] diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index e7bc29372e4..d351afb2061 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -13,6 +13,7 @@ pub mod model_capacity; pub mod prometheus; pub mod rate_limits; pub mod request_hooks; +pub mod s3_log; mod secrets; pub mod spend; pub use error::Error; diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index bdf8ffef1da..71482524afc 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -144,6 +144,9 @@ async fn main() -> Result<(), Box> { if let Some(prometheus) = &prometheus { observers.push(prometheus.clone()); } + if let Some(s3_log) = litellm_gateway::s3_log::S3Log::from_config(&config, &inference).await? { + observers.push(s3_log); + } litellm_gateway::hooks::install(&inference, &config, hooks, observers).await?; let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; let address = listener.local_addr()?; diff --git a/litellm-rust/crates/gateway/src/prometheus.rs b/litellm-rust/crates/gateway/src/prometheus.rs index 1a397351fdb..28cf7ef0dcd 100644 --- a/litellm-rust/crates/gateway/src/prometheus.rs +++ b/litellm-rust/crates/gateway/src/prometheus.rs @@ -678,16 +678,31 @@ impl<'a> Values<'a> { /// reports for them; a provider error is named the way the proxy maps its status, prefixed by /// the provider. fn exception_class(record: &Record) -> String { - let status = record.outcome.status; + let class = error_class(record); let Some(deployment) = &record.outcome.deployment else { + return class.to_owned(); + }; + match deployment.provider() { + Some(provider) if !provider.is_empty() => { + let mut letters = provider.chars(); + let first = letters.next().map(|c| c.to_uppercase().collect::()); + format!("{}{}.{class}", first.unwrap_or_default(), letters.as_str()) + } + _ => class.to_owned(), + } +} + +/// The class of the exception the proxy raises for a failed request, without the provider. +pub(crate) fn error_class(record: &Record) -> &'static str { + let status = record.outcome.status; + if record.outcome.deployment.is_none() { return match status { 400 => "BudgetExceededError", 429 | 503 => "HTTPException", _ => "ProxyException", - } - .to_owned(); - }; - let class = match status { + }; + } + match status { 400 => "BadRequestError", 401 => "AuthenticationError", 403 => "PermissionDeniedError", @@ -698,14 +713,6 @@ fn exception_class(record: &Record) -> String { 500 => "InternalServerError", 503 => "ServiceUnavailableError", _ => "APIError", - }; - match deployment.provider() { - Some(provider) if !provider.is_empty() => { - let mut letters = provider.chars(); - let first = letters.next().map(|c| c.to_uppercase().collect::()); - format!("{}{}.{class}", first.unwrap_or_default(), letters.as_str()) - } - _ => class.to_owned(), } } diff --git a/litellm-rust/crates/gateway/src/s3_log/config.rs b/litellm-rust/crates/gateway/src/s3_log/config.rs new file mode 100644 index 00000000000..cbd99d9f5c2 --- /dev/null +++ b/litellm-rust/crates/gateway/src/s3_log/config.rs @@ -0,0 +1,676 @@ +//! `litellm_settings.s3_callback_params`, read the way the proxy's `S3Logger` reads it: values +//! may be `os.environ/` references, and a number or flag the proxy would not parse falls back to +//! its default with a warning instead of stopping startup. + +use std::time::Duration; + +use litellm_auth_aws::AwsAuthConfig; +use litellm_config::{Config, Value}; +use litellm_http::SslVerify; +use litellm_secrets::source::SecretSource; + +use super::key::Granularity; +use crate::Error; + +/// The names a config's callback lists turn this observer on with. +const NAMES: [&str; 2] = ["s3_v2", "s3_log"]; + +const DEFAULT_FLUSH_INTERVAL_SECONDS: u64 = 10; +const DEFAULT_MAX_CONCURRENT_UPLOADS: usize = 16; +const DEFAULT_MAX_QUEUE_SIZE: usize = 50_000; +const DEFAULT_MAX_RETRY_AGE_SECONDS: u64 = 3600; +const DEFAULT_UPLOAD_ATTEMPT_TIMEOUT_SECONDS: u64 = 30; +/// Where the proxy signs non-LLM AWS calls when no region is configured. +const DEFAULT_REGION: &str = "us-west-2"; + +/// A configured value once its `os.environ/` reference is resolved. +#[derive(Clone, Debug, PartialEq)] +enum Param { + Bool(bool), + Int(i64), + Float(f64), + Text(String), +} + +impl Param { + /// Python's `bool(value)`, which the proxy applies to the prefix and style flags: any + /// non-empty string is true, `"false"` included. + fn truthy(&self) -> bool { + match self { + Self::Bool(flag) => *flag, + Self::Int(number) => *number != 0, + Self::Float(number) => *number != 0.0, + Self::Text(text) => !text.is_empty(), + } + } + + /// Pydantic's lax `bool`. + fn boolean(&self) -> Option { + match self { + Self::Bool(flag) => Some(*flag), + Self::Int(0) => Some(false), + Self::Int(1) => Some(true), + Self::Float(number) if *number == 0.0 => Some(false), + Self::Float(number) if *number == 1.0 => Some(true), + Self::Text(text) => match text.trim().to_ascii_lowercase().as_str() { + "0" | "off" | "f" | "false" | "n" | "no" => Some(false), + "1" | "on" | "t" | "true" | "y" | "yes" => Some(true), + _ => None, + }, + _ => None, + } + } + + /// Pydantic's lax `int`. + fn integer(&self) -> Option { + match self { + Self::Bool(flag) => Some(i64::from(*flag)), + Self::Int(number) => Some(*number), + Self::Float(number) if number.fract() == 0.0 && number.is_finite() => { + Some(*number as i64) + } + Self::Text(text) => text.trim().parse().ok(), + Self::Float(_) => None, + } + } + + fn is_blank(&self) -> bool { + matches!(self, Self::Text(text) if text.is_empty()) + } +} + +/// The resolved `s3_callback_params`. +struct Params(Vec<(String, Param)>); + +impl Params { + fn get(&self, name: &str) -> Option<&Param> { + self.0 + .iter() + .find(|(key, _)| key == name) + .map(|(_, value)| value) + } + + /// `params.get(name) or None`: an empty or false value counts as unset. + fn text(&self, name: &str) -> Result, Error> { + match self.get(name) { + None => Ok(None), + Some(Param::Text(text)) => Ok(Some(text.clone()).filter(|text| !text.is_empty())), + Some(Param::Int(number)) => Ok(Some(number.to_string())), + Some(Param::Float(number)) => Ok(Some(number.to_string())), + Some(Param::Bool(_)) => Err(Error::S3Log(format!("{name} must be a string"))), + } + } + + fn truthy(&self, name: &str) -> bool { + self.get(name).is_some_and(Param::truthy) + } + + /// A flag the proxy parses with pydantic; an invalid value warns and takes `invalid`. + fn flag(&self, name: &str, default: bool, invalid: bool) -> bool { + match self.get(name) { + None => default, + Some(value) if value.is_blank() => default, + Some(value) => value.boolean().unwrap_or_else(|| { + tracing::warn!( + setting = name, + ?value, + "s3 logging: not a boolean, using {invalid}" + ); + invalid + }), + } + } + + /// A positive integer; anything else warns and takes `default`. + fn positive(&self, name: &str, default: usize, reject_bool: bool) -> usize { + let Some(value) = self.get(name).filter(|value| !value.is_blank()) else { + return default; + }; + let parsed = match value { + Param::Bool(_) if reject_bool => None, + value => value.integer(), + }; + match parsed.and_then(|number| usize::try_from(number).ok()) { + Some(number) if number >= 1 => number, + _ => { + tracing::warn!( + setting = name, + ?value, + "s3 logging: not a positive integer, using {default}" + ); + default + } + } + } +} + +/// How each request is logged and where it goes. +#[derive(Clone, Debug, PartialEq)] +pub(super) struct Settings { + pub bucket: String, + /// `s3_path`, the prefix every key starts with. + pub path: String, + /// The region requests are signed for and the default endpoint names. + pub region: String, + pub endpoint_url: Option, + pub virtual_hosted_style: bool, + /// `s3_verify`; unset keeps the gateway's own TLS settings. + pub verify: Option, + pub auth: AwsAuthConfig, + pub team_prefix: bool, + pub key_prefix: bool, + pub strip_base64_files: bool, + pub prompts_only: bool, + pub granularity: Granularity, + pub server_side_encryption: Option, + pub sse_kms_key_id: Option, + pub max_concurrent_uploads: usize, + pub max_queue_size: usize, + /// How long an upload that keeps failing is retried; `None` retries it until the queue + /// overflows. + pub max_retry_age: Option, + pub drop_on_terminal_error: bool, + pub batch_file_upload: bool, + pub flush_interval: Duration, + pub attempt_timeout: Duration, + pub log_success: bool, + pub log_failure: bool, + pub turn_off_message_logging: bool, + pub redact_user_api_key_info: bool, + pub user_agent_tags: bool, +} + +impl Settings { + /// `Ok(None)` when no callback list names `s3_v2` or `s3_log`. A config the proxy could + /// never upload with stops the gateway. + pub async fn from_config( + config: &Config, + secrets: &dyn SecretSource, + env: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, Error> { + let litellm = &config.litellm_settings; + let listed = |list: &Option>| { + list.iter().flatten().any(|callback| { + callback + .as_str() + .is_some_and(|name| NAMES.contains(&name.trim())) + }) + }; + let everywhere = listed(&litellm.callbacks); + let log_success = everywhere || listed(&litellm.success_callback); + let log_failure = everywhere || listed(&litellm.failure_callback); + if !log_success && !log_failure { + return Ok(None); + } + let params = resolve(litellm.additional_fields.get("s3_callback_params"), secrets).await?; + for (name, why) in [ + ( + "s3_flush_interval", + "the proxy reads DEFAULT_S3_FLUSH_INTERVAL_SECONDS instead", + ), + ("s3_batch_size", "the proxy flushes on its interval alone"), + ( + "s3_adaptive_concurrency", + "this gateway uploads s3_max_concurrent_uploads at a time", + ), + ] { + if params.get(name).is_some() { + tracing::warn!(setting = name, "s3 logging: ignored, {why}"); + } + } + let bucket = params + .text("s3_bucket_name")? + .ok_or_else(|| Error::S3Log("s3_bucket_name is required".into()))?; + let configured_region = params.text("s3_region_name")?; + let region = configured_region + .clone() + .or_else(|| env("AWS_REGION")) + .or_else(|| env("AWS_REGION_NAME")) + .unwrap_or_else(|| DEFAULT_REGION.into()); + let (server_side_encryption, sse_kms_key_id) = sse( + params.text("s3_server_side_encryption")?, + params.text("s3_sse_kms_key_id")?, + ); + let verify = match params.get("s3_verify") { + None => None, + Some(Param::Bool(true)) => Some(SslVerify::Enabled), + Some(Param::Bool(false)) => Some(SslVerify::Disabled), + Some(Param::Text(text)) => Some(SslVerify::parse(text)), + Some(other) => { + return Err(Error::S3Log(format!( + "s3_verify must be a boolean or a CA bundle path, not {other:?}" + ))); + } + }; + let env_number = |name: &str, default: u64| { + env(name) + .and_then(|value| value.trim().parse::().ok()) + .unwrap_or(default) + }; + let max_retry_age = match params.get("s3_max_retry_age_seconds") { + None => Some(DEFAULT_MAX_RETRY_AGE_SECONDS), + Some(value) if value.is_blank() => Some(DEFAULT_MAX_RETRY_AGE_SECONDS), + Some(value) => match value { + Param::Bool(_) => None, + value => value.integer(), + } + .filter(|seconds| *seconds >= 0) + .map(|seconds| seconds as u64) + .or_else(|| { + tracing::warn!(?value, "s3 logging: s3_max_retry_age_seconds is not a non-negative integer, using {DEFAULT_MAX_RETRY_AGE_SECONDS}"); + Some(DEFAULT_MAX_RETRY_AGE_SECONDS) + }), + } + .filter(|seconds| *seconds > 0) + .map(Duration::from_secs); + let prompts_only = match params.get("s3_log_prompts_only").cloned() { + None => env("S3_LOG_PROMPTS_ONLY").map(Param::Text), + value => value, + } + .filter(|value| !value.is_blank()) + .is_some_and(|value| { + value.boolean().unwrap_or_else(|| { + tracing::warn!( + ?value, + "s3 logging: s3_log_prompts_only is not a boolean, logging prompts only" + ); + true + }) + }); + let granularity = match params.get("s3_partition_granularity").cloned() { + None => env("S3_PARTITION_GRANULARITY").map(Param::Text), + value => value, + }; + let granularity = match granularity { + Some(Param::Text(text)) if text == "hour" => Granularity::Hour, + None => Granularity::Day, + Some(Param::Text(text)) if text.is_empty() || text == "day" => Granularity::Day, + Some(value) => { + tracing::warn!( + ?value, + "s3 logging: s3_partition_granularity is not one of day, hour, using day" + ); + Granularity::Day + } + }; + let litellm_flag = + |name: &str| matches!(litellm.additional_fields.get(name), Some(Value::Bool(true))); + Ok(Some(Self { + bucket, + path: params.text("s3_path")?.unwrap_or_default(), + region: region.clone(), + endpoint_url: params.text("s3_endpoint_url")?, + virtual_hosted_style: params.truthy("s3_use_virtual_hosted_style"), + verify, + auth: AwsAuthConfig { + access_key_id: params.text("s3_aws_access_key_id")?, + secret_access_key: params.text("s3_aws_secret_access_key")?, + session_token: params.text("s3_aws_session_token")?, + region_name: Some(region), + session_name: params.text("s3_aws_session_name")?, + profile_name: params.text("s3_aws_profile_name")?, + role_name: params.text("s3_aws_role_name")?, + web_identity_token: params.text("s3_aws_web_identity_token")?, + sts_endpoint: params.text("s3_aws_sts_endpoint")?, + external_id: None, + }, + team_prefix: params.truthy("s3_use_team_prefix"), + key_prefix: params.truthy("s3_use_key_prefix"), + strip_base64_files: params.truthy("s3_strip_base64_files"), + prompts_only, + granularity, + server_side_encryption, + sse_kms_key_id, + max_concurrent_uploads: params.positive( + "s3_max_concurrent_uploads", + env_number( + "DEFAULT_S3_MAX_CONCURRENT_UPLOADS", + DEFAULT_MAX_CONCURRENT_UPLOADS as u64, + ) as usize, + false, + ), + max_queue_size: params.positive("s3_max_queue_size", DEFAULT_MAX_QUEUE_SIZE, true), + max_retry_age, + drop_on_terminal_error: params.flag("s3_drop_on_terminal_error", true, true), + batch_file_upload: params.flag("s3_batch_file_upload", false, false), + flush_interval: Duration::from_secs( + env_number( + "DEFAULT_S3_FLUSH_INTERVAL_SECONDS", + DEFAULT_FLUSH_INTERVAL_SECONDS, + ) + .max(1), + ), + attempt_timeout: Duration::from_secs( + env_number( + "S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS", + DEFAULT_UPLOAD_ATTEMPT_TIMEOUT_SECONDS, + ) + .clamp(1, 3600), + ), + log_success, + log_failure, + turn_off_message_logging: litellm_flag("turn_off_message_logging"), + redact_user_api_key_info: litellm_flag("redact_user_api_key_info"), + user_agent_tags: !litellm_flag("disable_add_user_agent_to_request_tags"), + })) + } +} + +/// The encryption headers to send: a KMS key alone means `aws:kms`, and a key next to another +/// algorithm is dropped, as the proxy's `resolve_sse_params` does. +fn sse(algorithm: Option, key_id: Option) -> (Option, Option) { + let algorithm = algorithm.or_else(|| key_id.as_ref().map(|_| "aws:kms".into())); + match (algorithm, key_id) { + (None, _) => (None, None), + (Some(algorithm), Some(_)) if !algorithm.starts_with("aws:kms") => { + tracing::warn!(%algorithm, "s3 logging: ignoring s3_sse_kms_key_id, set s3_server_side_encryption to aws:kms to use it"); + (Some(algorithm), None) + } + (algorithm, key_id) => (algorithm, key_id), + } +} + +async fn resolve(value: Option<&Value>, secrets: &dyn SecretSource) -> Result { + let entries = match value { + None | Some(Value::Null) => return Ok(Params(Vec::new())), + Some(Value::Mapping(entries)) => entries, + Some(_) => return Err(Error::S3Log("s3_callback_params must be a mapping".into())), + }; + let mut params = Vec::with_capacity(entries.len()); + for (key, value) in entries { + let Some(name) = key.as_str() else { + return Err(Error::S3Log( + "s3_callback_params keys must be strings".into(), + )); + }; + let param = match value { + Value::Null => continue, + Value::Bool(flag) => Param::Bool(*flag), + Value::Number(number) => match number.as_i64() { + Some(number) => Param::Int(number), + None => Param::Float(number.as_f64().unwrap_or(f64::NAN)), + }, + Value::String(text) => match text.strip_prefix("os.environ/") { + Some(variable) => match secrets.get_secret_str(variable).await? { + Some(secret) => Param::Text(secret.expose().to_owned()), + None => continue, + }, + None => Param::Text(text.clone()), + }, + _ => return Err(Error::S3Log(format!("{name} must be a scalar"))), + }; + params.push((name.to_owned(), param)); + } + Ok(Params(params)) +} + +#[cfg(test)] +pub(super) mod tests { + use std::collections::HashMap; + + use futures_util::future::BoxFuture; + use litellm_auth_types::SecretValue; + use rstest::rstest; + + use super::*; + + struct Secrets(HashMap<&'static str, &'static str>); + + impl SecretSource for Secrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async move { Ok(self.0.get(name).map(|value| SecretValue::new(*value))) }) + } + } + + pub(crate) async fn settings_with( + yaml: &str, + environment: &[(&'static str, &'static str)], + ) -> Result, String> { + let config = Config::from_yaml(&format!("model_list: []\n{yaml}")).unwrap(); + let environment: HashMap<&str, &str> = environment.iter().copied().collect(); + let env = |name: &str| environment.get(name).map(|value| (*value).to_owned()); + Settings::from_config(&config, &Secrets(environment.clone()), &env) + .await + .map_err(|error| error.to_string()) + } + + async fn settings(params: &str) -> Settings { + settings_with( + &format!( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}" + ), + &[], + ) + .await + .unwrap() + .unwrap() + } + + #[rstest] + #[case::none("", None)] + #[case::other_callback(" success_callback: [prometheus]\n", None)] + #[case::callbacks(" callbacks: [s3_v2]\n", Some((true, true)))] + #[case::gateway_name(" callbacks: [s3_log]\n", Some((true, true)))] + #[case::success_only(" success_callback: [s3_v2]\n", Some((true, false)))] + #[case::failure_only(" failure_callback: s3_v2\n", Some((false, true)))] + #[tokio::test] + async fn on_for_the_lists_that_name_it( + #[case] lists: &str, + #[case] expected: Option<(bool, bool)>, + ) { + let yaml = + format!("litellm_settings:\n{lists} s3_callback_params:\n s3_bucket_name: logs\n"); + let settings = settings_with(&yaml, &[]).await.unwrap(); + assert_eq!( + settings.map(|settings| (settings.log_success, settings.log_failure)), + expected + ); + } + + #[tokio::test] + async fn defaults_are_the_proxys() { + let settings = settings("").await; + assert_eq!(settings.path, ""); + assert_eq!(settings.region, "us-west-2"); + assert_eq!(settings.granularity, Granularity::Day); + assert_eq!(settings.max_concurrent_uploads, 16); + assert_eq!(settings.max_queue_size, 50_000); + assert_eq!(settings.max_retry_age, Some(Duration::from_secs(3600))); + assert!(settings.drop_on_terminal_error); + assert!(!settings.batch_file_upload); + assert!(!settings.prompts_only); + assert_eq!(settings.flush_interval, Duration::from_secs(10)); + assert_eq!(settings.attempt_timeout, Duration::from_secs(30)); + assert_eq!(settings.verify, None); + assert_eq!(settings.server_side_encryption, None); + } + + #[tokio::test] + async fn references_are_resolved_from_the_secret_source() { + let settings = settings_with( + "litellm_settings: + callbacks: [s3_v2] + s3_callback_params: + s3_bucket_name: os.environ/BUCKET + s3_region_name: eu-west-1 + s3_aws_access_key_id: os.environ/KEY_ID + s3_aws_secret_access_key: os.environ/SECRET + s3_aws_session_token: os.environ/UNSET + s3_use_team_prefix: os.environ/TEAM_PREFIX +", + &[ + ("BUCKET", "audit"), + ("KEY_ID", "AKID"), + ("SECRET", "shh"), + ("TEAM_PREFIX", "false"), + ], + ) + .await + .unwrap() + .unwrap(); + assert_eq!(settings.bucket, "audit"); + assert_eq!(settings.region, "eu-west-1"); + assert_eq!(settings.auth.access_key_id.as_deref(), Some("AKID")); + assert_eq!(settings.auth.secret_access_key.as_deref(), Some("shh")); + assert_eq!(settings.auth.session_token, None); + assert_eq!(settings.auth.region_name.as_deref(), Some("eu-west-1")); + // The proxy takes `bool()` of the string, so any non-empty text turns it on. + assert!(settings.team_prefix); + } + + #[rstest] + #[case::configured(" s3_region_name: ap-south-1\n", &[("AWS_REGION", "us-east-2")], "ap-south-1")] + #[case::aws_region_wins("", &[("AWS_REGION", "us-east-2"), ("AWS_REGION_NAME", "eu-west-3")], "us-east-2")] + #[case::aws_region_name("", &[("AWS_REGION_NAME", "eu-west-3")], "eu-west-3")] + #[tokio::test] + async fn the_signing_region_falls_back_like_the_proxy( + #[case] params: &str, + #[case] environment: &[(&'static str, &'static str)], + #[case] expected: &str, + ) { + let settings = settings_with( + &format!( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}" + ), + environment, + ) + .await + .unwrap() + .unwrap(); + assert_eq!(settings.region, expected); + } + + #[rstest] + #[case::kms_key_alone(" s3_sse_kms_key_id: key-1\n", Some("aws:kms"), Some("key-1"))] + #[case::aes(" s3_server_side_encryption: AES256\n", Some("AES256"), None)] + #[case::key_with_aes( + " s3_server_side_encryption: AES256\n s3_sse_kms_key_id: key-1\n", + Some("AES256"), + None + )] + #[case::dsse( + " s3_server_side_encryption: aws:kms:dsse\n s3_sse_kms_key_id: key-1\n", + Some("aws:kms:dsse"), + Some("key-1") + )] + #[tokio::test] + async fn encryption_settings_combine_like_the_proxy( + #[case] params: &str, + #[case] algorithm: Option<&str>, + #[case] key: Option<&str>, + ) { + let settings = settings(params).await; + assert_eq!(settings.server_side_encryption.as_deref(), algorithm); + assert_eq!(settings.sse_kms_key_id.as_deref(), key); + } + + #[rstest] + #[case::text_number(" s3_max_concurrent_uploads: ' 4 '\n", 4, 50_000)] + #[case::zero_falls_back(" s3_max_concurrent_uploads: 0\n", 16, 50_000)] + #[case::garbage_falls_back(" s3_max_queue_size: lots\n", 16, 50_000)] + #[case::bool_queue_falls_back(" s3_max_queue_size: true\n", 16, 50_000)] + #[case::bool_width_is_one(" s3_max_concurrent_uploads: true\n", 1, 50_000)] + #[case::queue(" s3_max_queue_size: 10\n", 16, 10)] + #[tokio::test] + async fn numbers_parse_or_fall_back( + #[case] params: &str, + #[case] uploads: usize, + #[case] queue: usize, + ) { + let settings = settings(params).await; + assert_eq!(settings.max_concurrent_uploads, uploads); + assert_eq!(settings.max_queue_size, queue); + } + + #[rstest] + #[case::zero_disables(" s3_max_retry_age_seconds: 0\n", None)] + #[case::seconds(" s3_max_retry_age_seconds: '60'\n", Some(60))] + #[case::negative_falls_back(" s3_max_retry_age_seconds: -1\n", Some(3600))] + #[case::blank_is_default(" s3_max_retry_age_seconds: ''\n", Some(3600))] + #[tokio::test] + async fn retry_age(#[case] params: &str, #[case] seconds: Option) { + assert_eq!( + settings(params).await.max_retry_age, + seconds.map(Duration::from_secs) + ); + } + + #[rstest] + #[case::configured(" s3_log_prompts_only: 'yes'\n", &[], true)] + #[case::invalid_logs_prompts_only(" s3_log_prompts_only: sometimes\n", &[], true)] + #[case::environment("", &[("S3_LOG_PROMPTS_ONLY", "true")], true)] + #[case::configured_beats_environment(" s3_log_prompts_only: false\n", &[("S3_LOG_PROMPTS_ONLY", "true")], false)] + #[case::unset("", &[], false)] + #[tokio::test] + async fn prompts_only( + #[case] params: &str, + #[case] environment: &[(&'static str, &'static str)], + #[case] expected: bool, + ) { + let settings = settings_with( + &format!( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}" + ), + environment, + ) + .await + .unwrap() + .unwrap(); + assert_eq!(settings.prompts_only, expected); + } + + #[rstest] + #[case::hour(" s3_partition_granularity: hour\n", &[], Granularity::Hour)] + #[case::environment("", &[("S3_PARTITION_GRANULARITY", "hour")], Granularity::Hour)] + #[case::unknown_is_day(" s3_partition_granularity: minute\n", &[], Granularity::Day)] + #[tokio::test] + async fn partition_granularity( + #[case] params: &str, + #[case] environment: &[(&'static str, &'static str)], + #[case] expected: Granularity, + ) { + let settings = settings_with( + &format!( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}" + ), + environment, + ) + .await + .unwrap() + .unwrap(); + assert_eq!(settings.granularity, expected); + } + + #[rstest] + #[case::no_bucket( + "litellm_settings:\n callbacks: [s3_v2]\n", + "s3_bucket_name is required" + )] + #[case::unset_bucket_reference( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: os.environ/NOPE\n", + "s3_bucket_name is required" + )] + #[case::not_a_mapping( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params: [a]\n", + "s3_callback_params must be a mapping" + )] + #[case::nested_value( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n s3_path: [a, b]\n", + "s3_path must be a scalar" + )] + #[case::bad_verify( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n s3_verify: 3\n", + "s3_verify must be a boolean or a CA bundle path" + )] + #[tokio::test] + async fn a_config_that_cannot_upload_stops_the_gateway( + #[case] yaml: &str, + #[case] message: &str, + ) { + let error = settings_with(yaml, &[]).await.unwrap_err(); + assert!(error.contains(message), "{error}"); + } +} diff --git a/litellm-rust/crates/gateway/src/s3_log/key.rs b/litellm-rust/crates/gateway/src/s3_log/key.rs new file mode 100644 index 00000000000..19bc6be2d34 --- /dev/null +++ b/litellm-rust/crates/gateway/src/s3_log/key.rs @@ -0,0 +1,389 @@ +//! Object keys, download names and URLs, laid out as the proxy's `get_s3_object_key` and +//! `S3Logger._build_object_url` lay them out, length bounds included. + +use std::time::{SystemTime, UNIX_EPOCH}; + +use sha2::{Digest, Sha256}; + +use super::config::Settings; +use crate::spend::civil; + +/// S3's limit on a key, in UTF-8 bytes. +const MAX_KEY_BYTES: usize = 1024; +/// The proxy's limit on the last segment of a key. +const MAX_FILE_NAME_BYTES: usize = 255; +/// The proxy's limit on the `Content-Disposition` file name. +const MAX_DOWNLOAD_NAME_BYTES: usize = 1024; +/// How much of an over-long name is kept before its digest. +const BOUNDED_HEAD_BYTES: i64 = 64; +const MIN_BOUNDED_FILE_NAME_BYTES: usize = 64; +const PREFIX_DIGEST_CHARS: usize = 16; +const EXTENSION: &str = ".json"; + +/// `s3_partition_granularity`: a folder per day, or per hour within it. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(super) enum Granularity { + #[default] + Day, + Hour, +} + +/// A wall-clock time, UTC, split the way the key names it. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct Utc { + year: i64, + month: u32, + day: u32, + hour: u64, + minute: u64, + second: u64, + micros: u32, +} + +impl Utc { + pub fn new(time: SystemTime) -> Self { + let elapsed = time.duration_since(UNIX_EPOCH).unwrap_or_default(); + let seconds = elapsed.as_secs(); + let (year, month, day) = civil((seconds / 86_400) as i64); + let of_day = seconds % 86_400; + Self { + year, + month, + day, + hour: of_day / 3600, + minute: of_day / 60 % 60, + second: of_day % 60, + micros: elapsed.subsec_micros(), + } + } + + fn date(&self) -> String { + format!("{:04}-{:02}-{:02}", self.year, self.month, self.day) + } + + /// `%H-%M-%S`. + pub fn clock(&self) -> String { + format!("{:02}-{:02}-{:02}", self.hour, self.minute, self.second) + } + + /// `%H-%M-%S-%f`. + fn clock_micros(&self) -> String { + format!("{}-{:06}", self.clock(), self.micros) + } +} + +/// The proxy's `get_logging_id`: the request's time and the payload id. +pub(super) fn logging_id(start: &Utc, id: &str) -> String { + format!("time-{}_{id}", start.clock_micros()) +} + +/// `{s3_path}/[prefix]YYYY-MM-DD[/HH]/{file_name}.json`, shortened as the proxy shortens it when +/// it would pass S3's limits: the file name first, the configured prefix only if that is what +/// does not fit. +pub(super) fn object_key( + path: &str, + prefix: &str, + start: &Utc, + file_name: &str, + granularity: Granularity, +) -> String { + let sanitized = file_name.replace(['/', ':'], "_"); + let configured = match path.is_empty() { + true => prefix.to_owned(), + false => format!("{}/{prefix}", path.trim_end_matches('/')), + }; + let date = match granularity { + Granularity::Day => format!("{}/", start.date()), + Granularity::Hour => format!("{}/{:02}/", start.date(), start.hour), + }; + let file_name_budget = MAX_FILE_NAME_BYTES - EXTENSION.len(); + let key = format!("{configured}{date}{sanitized}{EXTENSION}"); + if key.len() <= MAX_KEY_BYTES && sanitized.len() <= file_name_budget { + return key; + } + let budget = MAX_KEY_BYTES - date.len() - EXTENSION.len(); + if configured.len() + MIN_BOUNDED_FILE_NAME_BYTES <= budget { + let bounded = bounded_name( + file_name, + &sanitized, + file_name_budget.min(budget - configured.len()), + ); + return format!("{configured}{date}{bounded}{EXTENSION}"); + } + let shortest = bounded_name(file_name, &sanitized, MIN_BOUNDED_FILE_NAME_BYTES); + let prefix = bounded_prefix(&configured, budget.saturating_sub(shortest.len())); + format!("{prefix}{date}{shortest}{EXTENSION}") +} + +/// The proxy's `get_s3_object_download_filename`. +pub(super) fn download_name(start: &Utc, id: &str) -> String { + let stamp = format!("time-{}T{}", start.date(), start.clock_micros()); + let name = format!("{stamp}_{id}"); + let sanitized = format!("{stamp}_{}", id.replace(['/', '"'], "_")); + let budget = MAX_DOWNLOAD_NAME_BYTES - EXTENSION.len(); + match sanitized.len() <= budget { + true => format!("{sanitized}{EXTENSION}"), + false => format!("{}{EXTENSION}", bounded_name(&name, &sanitized, budget)), + } +} + +/// As much of `sanitized` as fits, then the sha256 of the unsanitized name. +fn bounded_name(name: &str, sanitized: &str, max_bytes: usize) -> String { + let digest = hex_sha256(name); + let head_budget = BOUNDED_HEAD_BYTES.min(max_bytes as i64 - digest.len() as i64 - 1); + let head = truncate(sanitized, head_budget); + match head.is_empty() { + true => digest, + false => format!("{head}_{digest}"), + } +} + +/// As much of the configured prefix as fits, then a digest segment naming all of it. +fn bounded_prefix(prefix: &str, max_bytes: usize) -> String { + let digest = format!("{}/", &hex_sha256(prefix)[..PREFIX_DIGEST_CHARS]); + if max_bytes < digest.len() { + return String::new(); + } + let head = truncate(prefix, max_bytes as i64 - digest.len() as i64 - 1).trim_end_matches('/'); + match head.is_empty() { + true => digest, + false => format!("{head}/{digest}"), + } +} + +/// `value` cut to at most `max_bytes` of UTF-8, never inside a character. +fn truncate(value: &str, max_bytes: i64) -> &str { + let Ok(max_bytes) = usize::try_from(max_bytes) else { + return ""; + }; + if value.len() <= max_bytes { + return value; + } + let end = (0..=max_bytes) + .rev() + .find(|end| value.is_char_boundary(*end)) + .unwrap_or(0); + &value[..end] +} + +fn hex_sha256(value: &str) -> String { + Sha256::digest(value.as_bytes()) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +/// Everything before the last `/`, the folder batch files are grouped by. +pub(super) fn parent(key: &str) -> &str { + key.rsplit_once('/').map_or("", |(parent, _)| parent) +} + +/// The URL that is signed and sent, with the key percent-encoded once, as the proxy builds it. +pub(super) fn object_url(settings: &Settings, key: &str) -> String { + let key = encode(key); + let bucket = &settings.bucket; + match &settings.endpoint_url { + Some(endpoint) if settings.virtual_hosted_style => { + let (scheme, host) = match endpoint.strip_prefix("https://") { + Some(host) => ("https://", host), + None => ("http://", endpoint.trim_start_matches("http://")), + }; + format!("{scheme}{bucket}.{host}/{key}") + } + Some(endpoint) => format!("{endpoint}/{bucket}/{key}"), + None => format!( + "https://{bucket}.s3.{region}.{suffix}/{key}", + region = settings.region, + suffix = dns_suffix(&settings.region) + ), + } +} + +/// Python's `quote(key, safe="/")`. +fn encode(key: &str) -> String { + key.bytes() + .map(|byte| match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' | b'/' => { + char::from(byte).to_string() + } + byte => format!("%{byte:02X}"), + }) + .collect() +} + +/// The AWS partition's domain for a region. +fn dns_suffix(region: &str) -> &'static str { + [ + ("cn-", "amazonaws.com.cn"), + ("us-isob-", "sc2s.sgov.gov"), + ("us-isof-", "csp.hci.ic.gov"), + ("us-iso-", "c2s.ic.gov"), + ("eu-isoe-", "cloud.adc-e.uk"), + ] + .into_iter() + .find(|(prefix, _)| region.starts_with(prefix)) + .map_or("amazonaws.com", |(_, suffix)| suffix) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use rstest::rstest; + + use super::*; + + /// 2026-03-04 05:06:07.000089 UTC. + fn start() -> Utc { + Utc::new(UNIX_EPOCH + Duration::from_micros(1_772_600_767_000_089)) + } + + #[rstest] + #[case::bare( + "", + "", + Granularity::Day, + "2026-03-04/time-05-06-07-000089_chatcmpl-1.json" + )] + #[case::path( + "logs/", + "", + Granularity::Day, + "logs/2026-03-04/time-05-06-07-000089_chatcmpl-1.json" + )] + #[case::prefixes( + "logs", + "research/team-key/", + Granularity::Day, + "logs/research/team-key/2026-03-04/time-05-06-07-000089_chatcmpl-1.json" + )] + #[case::hour( + "logs", + "", + Granularity::Hour, + "logs/2026-03-04/05/time-05-06-07-000089_chatcmpl-1.json" + )] + fn keys_are_laid_out_like_the_proxy( + #[case] path: &str, + #[case] prefix: &str, + #[case] granularity: Granularity, + #[case] expected: &str, + ) { + let file_name = logging_id(&start(), "chatcmpl-1"); + assert_eq!( + object_key(path, prefix, &start(), &file_name, granularity), + expected + ); + } + + #[test] + fn slashes_and_colons_in_the_id_stay_in_one_segment() { + let file_name = logging_id(&start(), "a/b:c"); + assert_eq!( + object_key("", "", &start(), &file_name, Granularity::Day), + "2026-03-04/time-05-06-07-000089_a_b_c.json" + ); + assert_eq!( + download_name(&start(), "a/b\"c"), + "time-2026-03-04T05-06-07-000089_a_b_c.json" + ); + } + + #[test] + fn a_long_id_is_shortened_to_a_digest_within_s3_limits() { + let id = "é".repeat(400); + let file_name = logging_id(&start(), &id); + let key = object_key("logs", "", &start(), &file_name, Granularity::Day); + let (folder, last) = key.rsplit_once('/').unwrap(); + assert_eq!(folder, "logs/2026-03-04"); + assert!(last.len() <= MAX_FILE_NAME_BYTES, "{last}"); + assert!(last.ends_with(&format!("_{}.json", hex_sha256(&file_name)))); + assert!(last.starts_with("time-05-06-07-000089_é")); + } + + #[test] + fn a_long_prefix_keeps_its_head_and_a_digest_of_all_of_it() { + let path = "p".repeat(2000); + let file_name = logging_id(&start(), "chatcmpl-1"); + let key = object_key(&path, "", &start(), &file_name, Granularity::Day); + assert!(key.len() <= MAX_KEY_BYTES, "{}", key.len()); + let digest = &hex_sha256(&format!("{path}/"))[..PREFIX_DIGEST_CHARS]; + assert!(key.starts_with("ppp"), "{key}"); + assert!(key.contains(&format!("/{digest}/2026-03-04/")), "{key}"); + } + + fn settings(endpoint: Option<&str>, virtual_hosted: bool, region: &str) -> Settings { + Settings { + bucket: "logs".into(), + path: String::new(), + region: region.into(), + endpoint_url: endpoint.map(str::to_owned), + virtual_hosted_style: virtual_hosted, + verify: None, + auth: Default::default(), + team_prefix: false, + key_prefix: false, + strip_base64_files: false, + prompts_only: false, + granularity: Granularity::Day, + server_side_encryption: None, + sse_kms_key_id: None, + max_concurrent_uploads: 1, + max_queue_size: 1, + max_retry_age: None, + drop_on_terminal_error: true, + batch_file_upload: false, + flush_interval: Duration::from_secs(1), + attempt_timeout: Duration::from_secs(1), + log_success: true, + log_failure: true, + turn_off_message_logging: false, + redact_user_api_key_info: false, + user_agent_tags: true, + } + } + + #[rstest] + #[case::aws( + None, + false, + "eu-west-1", + "https://logs.s3.eu-west-1.amazonaws.com/a/b%20c%3D%2B.json" + )] + #[case::china( + None, + false, + "cn-north-1", + "https://logs.s3.cn-north-1.amazonaws.com.cn/a/b%20c%3D%2B.json" + )] + #[case::path_style( + Some("http://minio:9000"), + false, + "us-east-1", + "http://minio:9000/logs/a/b%20c%3D%2B.json" + )] + #[case::virtual_hosted( + Some("https://s3.example"), + true, + "us-east-1", + "https://logs.s3.example/a/b%20c%3D%2B.json" + )] + fn urls_are_built_like_the_proxy( + #[case] endpoint: Option<&str>, + #[case] virtual_hosted: bool, + #[case] region: &str, + #[case] expected: &str, + ) { + assert_eq!( + object_url(&settings(endpoint, virtual_hosted, region), "a/b c=+.json"), + expected + ); + } + + #[rstest] + #[case("a/b/c.json", "a/b")] + #[case("c.json", "")] + fn parents(#[case] key: &str, #[case] expected: &str) { + assert_eq!(parent(key), expected); + } +} diff --git a/litellm-rust/crates/gateway/src/s3_log/mod.rs b/litellm-rust/crates/gateway/src/s3_log/mod.rs new file mode 100644 index 00000000000..73437dbdcfe --- /dev/null +++ b/litellm-rust/crates/gateway/src/s3_log/mod.rs @@ -0,0 +1,506 @@ +//! The `s3_v2` callback: every request's `StandardLoggingPayload` written to S3 as a JSON +//! object, under the proxy's key layout, as the proxy's `S3Logger` writes it. +//! +//! It is an observer, so nothing here runs on the request path. Each finished request is turned +//! into its object and queued; every flush interval the queue is uploaded with bounded +//! concurrency, or written as one NDJSON object per folder with `s3_batch_file_upload`. What +//! fails for now goes back on the queue for the next flush, until it is older than +//! `s3_max_retry_age_seconds` or the queue overflows. + +mod config; +mod key; +mod payload; +mod upload; + +use std::{ + collections::{BTreeMap, VecDeque}, + sync::{Arc, Mutex, Weak}, + time::{Instant, SystemTime}, +}; + +use axum::body::Bytes; +use futures_util::{StreamExt, future::BoxFuture}; +use litellm_config::Config; +use litellm_gateway_hooks::{Facts, Observer, Record}; +use litellm_gateway_inference::Gateway; +use serde_json::Value; + +use self::{ + config::Settings, + key::{Utc, download_name, logging_id, object_key, parent}, + payload::{Redaction, prompts_only, standard_logging_payload, strip_base64_files}, + upload::{Upload, Uploader}, +}; +use crate::Error; + +const JSON: &str = "application/json"; +const NDJSON: &str = "application/x-ndjson"; + +/// One object waiting to be written. +#[derive(Clone, Debug)] +struct Element { + key: String, + download_name: String, + body: Bytes, + content_type: &'static str, + /// A batch file, already holding many requests; never grouped again. + batch: bool, + /// When this object first failed to upload. + retrying_since: Option, +} + +#[derive(Default)] +struct Queue { + elements: VecDeque, + /// How many at the front were put back by the last failed flush. + requeued: usize, +} + +/// The `s3_v2` observer. +pub struct S3Log { + settings: Arc, + uploader: Uploader, + queue: Mutex, + /// Serializes flushes, so the periodic one and the one at shutdown never interleave. + flushing: tokio::sync::Mutex<()>, +} + +impl S3Log { + /// Starts the observer and its flush timer, or `Ok(None)` when no callback list names + /// `s3_v2`. Must be called inside a tokio runtime. + pub async fn from_config( + config: &Config, + inference: &Gateway, + ) -> Result>, Error> { + let env = |name: &str| std::env::var(name).ok(); + let Some(settings) = + Settings::from_config(config, inference.secrets.as_ref(), &env).await? + else { + return Ok(None); + }; + let settings = Arc::new(settings); + let log = Arc::new(Self { + uploader: Uploader::new( + settings.clone(), + inference.resources.pool.clone(), + &inference.http, + ), + settings, + queue: Mutex::default(), + flushing: tokio::sync::Mutex::new(()), + }); + tokio::spawn(run(Arc::downgrade(&log))); + Ok(Some(log)) + } + + fn element(&self, record: &Record) -> Element { + let settings = &self.settings; + let redaction = Redaction { + messages: settings.turn_off_message_logging, + user_api_key_info: settings.redact_user_api_key_info, + user_agent_tags: settings.user_agent_tags, + }; + let mut payload = standard_logging_payload(record, redaction); + if settings.strip_base64_files { + strip_base64_files(&mut payload); + } + let metadata = |field: &str| { + payload["metadata"][field] + .as_str() + .filter(|value| !value.is_empty()) + .map(str::to_owned) + }; + let prefix: String = [ + settings + .team_prefix + .then(|| metadata("user_api_key_team_alias")), + settings.key_prefix.then(|| metadata("user_api_key_alias")), + ] + .into_iter() + .flatten() + .flatten() + .map(|segment| format!("{segment}/")) + .collect(); + let id = payload["id"].as_str().unwrap_or_default().to_owned(); + let start = Utc::new(record.request.started_at); + if settings.prompts_only { + prompts_only(&mut payload); + } + Element { + key: object_key( + &settings.path, + &prefix, + &start, + &logging_id(&start, &id), + settings.granularity, + ), + download_name: download_name(&start, &id), + body: Bytes::from(Value::Object(payload).to_string()), + content_type: JSON, + batch: false, + retrying_since: None, + } + } + + fn enqueue(&self, element: Element) { + let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner()); + queue.elements.push_back(element); + let excess = queue + .elements + .len() + .saturating_sub(self.settings.max_queue_size); + if excess > 0 { + queue.elements.drain(..excess); + queue.requeued = queue.requeued.saturating_sub(excess); + tracing::warn!( + dropped = excess, + max_queue_size = self.settings.max_queue_size, + "s3 logging: queue full; dropped the oldest events" + ); + } + } + + /// Uploads everything queued; what fails for now is put back at the front of the queue. + pub async fn flush(&self) { + let _flushing = self.flushing.lock().await; + let (batch, requeued) = { + let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner()); + let requeued = std::mem::take(&mut queue.requeued); + (Vec::from(std::mem::take(&mut queue.elements)), requeued) + }; + if batch.is_empty() { + return; + } + let batch_len = batch.len(); + let uploads = match self.settings.batch_file_upload { + true => batch_files(batch), + false => batch, + }; + // What failed last time goes last, so a backlog never starves new requests. + let stale = match uploads.len() == batch_len { + true => requeued.min(uploads.len()), + false => 0, + }; + let order: Vec = (stale..uploads.len()).chain(0..stale).collect(); + let mut outcomes = vec![Upload::Retry; uploads.len()]; + let mut finished = futures_util::stream::iter(order) + .map(|index| { + let uploads = &uploads; + async move { (index, self.uploader.upload(&uploads[index]).await) } + }) + .buffer_unordered(self.settings.max_concurrent_uploads); + while let Some((index, outcome)) = finished.next().await { + outcomes[index] = outcome; + } + drop(finished); + let delivered = outcomes + .iter() + .filter(|outcome| **outcome == Upload::Delivered) + .count(); + let total = uploads.len(); + let failed = total - delivered; + if failed == 0 { + return; + } + let now = Instant::now(); + let max_age = self.settings.max_retry_age; + let failures = uploads + .into_iter() + .zip(outcomes) + .filter(|(_, outcome)| *outcome != Upload::Delivered); + // When nothing got through the bucket itself is failing: keep everything, unaged. + let kept: Vec = match delivered { + 0 => failures.map(|(element, _)| element).collect(), + _ => failures + .filter(|(element, outcome)| { + *outcome == Upload::Retry + && !matches!( + (max_age, element.retrying_since), + (Some(max_age), Some(since)) if now.duration_since(since) > max_age + ) + }) + .map(|(element, _)| Element { + retrying_since: element.retrying_since.or(max_age.map(|_| now)), + ..element + }) + .collect(), + }; + tracing::warn!( + failed, + kept = kept.len(), + total, + "s3 logging: uploads failed this flush; the kept ones are retried on the next, the rest \ + were rejected by S3 or retried longer than s3_max_retry_age_seconds" + ); + if kept.is_empty() { + return; + } + let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner()); + let arrivals = std::mem::take(&mut queue.elements); + let kept_len = kept.len(); + let mut elements: VecDeque = kept.into_iter().chain(arrivals).collect(); + let overflow = elements.len().saturating_sub(self.settings.max_queue_size); + if overflow > 0 { + elements.drain(..overflow); + tracing::warn!( + dropped = overflow, + max_queue_size = self.settings.max_queue_size, + "s3 logging: queue exceeded max_queue_size after a failed flush; dropped the oldest events" + ); + } + queue.elements = elements; + queue.requeued = kept_len.saturating_sub(overflow); + } +} + +/// `s3_batch_file_upload`: one NDJSON object per folder, named for the flush. +fn batch_files(batch: Vec) -> Vec { + let now = Utc::new(SystemTime::now()); + let (ready, singles): (Vec, Vec) = + batch.into_iter().partition(|element| element.batch); + let mut folders: BTreeMap> = BTreeMap::new(); + for element in singles { + folders + .entry(parent(&element.key).to_owned()) + .or_default() + .push(element); + } + let files = folders.into_iter().map(|(folder, elements)| { + let name = format!("batch_{}_{}", now.clock(), uuid::Uuid::new_v4().simple()); + let lines: Vec<&[u8]> = elements.iter().map(|element| &element.body[..]).collect(); + Element { + key: match folder.is_empty() { + true => format!("{name}.jsonl"), + false => format!("{folder}/{name}.jsonl"), + }, + download_name: format!("{name}.jsonl"), + body: Bytes::from(lines.join(&b'\n')), + content_type: NDJSON, + batch: true, + retrying_since: elements + .iter() + .filter_map(|element| element.retrying_since) + .min(), + } + }); + ready.into_iter().chain(files).collect() +} + +impl Observer for S3Log { + fn name(&self) -> &'static str { + "s3_v2" + } + + fn reads(&self) -> Facts { + Facts::KEY + | Facts::END_USER + | Facts::CLIENT_HEADERS + | Facts::DEPLOYMENT + | Facts::USAGE + | Facts::COST + | Facts::TIMINGS + | Facts::ANNOTATIONS + | Facts::REQUEST_BODY + | Facts::RESPONSE_BODY + } + + fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> { + let logged = match record.succeeded() { + true => self.settings.log_success, + false => self.settings.log_failure, + }; + if logged { + self.enqueue(self.element(record)); + } + Box::pin(async {}) + } + + fn flush(&self) -> BoxFuture<'_, ()> { + Box::pin(S3Log::flush(self)) + } +} + +/// Flushes every flush interval until the observer is dropped. +async fn run(log: Weak) { + let interval = match log.upgrade() { + Some(log) => log.settings.flush_interval, + None => return, + }; + loop { + tokio::time::sleep(interval).await; + let Some(log) = log.upgrade() else { + return; + }; + log.flush().await; + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use litellm_gateway_hooks::Route; + use serde_json::json; + + use super::*; + + fn element(key: &str, body: &str) -> Element { + Element { + key: key.into(), + download_name: "d.json".into(), + body: Bytes::from(body.to_owned()), + content_type: JSON, + batch: false, + retrying_since: None, + } + } + + #[test] + fn batch_files_hold_one_line_per_request_per_folder() { + let earlier = Instant::now(); + let mut retried = element("logs/2026-03-04/b.json", "{\"b\":1}"); + retried.retrying_since = Some(earlier); + let files = batch_files(vec![ + element("logs/2026-03-04/a.json", "{\"a\":1}"), + element("logs/team/2026-03-04/c.json", "{\"c\":1}"), + retried, + ]); + assert_eq!(files.len(), 2); + let day = &files[0]; + let (folder, name) = day.key.rsplit_once('/').unwrap(); + assert_eq!(folder, "logs/2026-03-04"); + assert!( + name.starts_with("batch_") && name.ends_with(".jsonl"), + "{name}" + ); + assert_eq!(day.download_name, name); + assert_eq!(&day.body[..], b"{\"a\":1}\n{\"b\":1}"); + assert_eq!(day.content_type, NDJSON); + assert_eq!(day.retrying_since, Some(earlier)); + assert!(files[1].key.starts_with("logs/team/2026-03-04/batch_")); + // A batch file kept for a retry is sent again as it is. + let again = batch_files(files.clone()); + assert_eq!( + again.iter().map(|file| &file.key).collect::>(), + files.iter().map(|file| &file.key).collect::>() + ); + } + + async fn log(params: &str) -> S3Log { + let settings = config::tests::settings_with( + &format!( + "litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n s3_path: audit\n{params}" + ), + &[], + ) + .await + .unwrap() + .unwrap(); + let settings = Arc::new(settings); + S3Log { + uploader: Uploader::new( + settings.clone(), + Arc::new(litellm_http::HttpClientPool::new(Arc::new( + litellm_http::media::PublicDnsResolver, + ))), + &litellm_http::Resolution::from(&litellm_http::HttpSettings::default()).config, + ), + settings, + queue: Mutex::default(), + flushing: tokio::sync::Mutex::new(()), + } + } + + fn record() -> Record { + payload::tests::record( + Route::ChatCompletions, + 200, + false, + json!({"messages": [{"role": "user", "content": "hi"}]}), + Some(payload::tests::COMPLETION), + ) + } + + #[rstest::rstest] + #[case::plain("", "audit/1970-01-01/time-00-16-40-000000_chatcmpl-1.json")] + #[case::team( + " s3_use_team_prefix: true\n", + "audit/research/1970-01-01/time-00-16-40-000000_chatcmpl-1.json" + )] + #[case::team_and_key( + " s3_use_team_prefix: true\n s3_use_key_prefix: true\n s3_partition_granularity: hour\n", + "audit/research/team-key/1970-01-01/00/time-00-16-40-000000_chatcmpl-1.json" + )] + #[tokio::test] + async fn objects_are_keyed_by_team_key_and_time(#[case] params: &str, #[case] expected: &str) { + let element = log(params).await.element(&record()); + assert_eq!(element.key, expected); + assert_eq!( + element.download_name, + "time-1970-01-01T00-16-40-000000_chatcmpl-1.json" + ); + } + + #[tokio::test] + async fn prompts_only_objects_have_no_response() { + let element = log(" s3_log_prompts_only: true\n") + .await + .element(&record()); + let body: Value = serde_json::from_slice(&element.body).unwrap(); + assert_eq!(body["response"], Value::Null); + assert_eq!(body["messages"], json!([{"role": "user", "content": "hi"}])); + assert!(element.key.ends_with("_chatcmpl-1.json")); + } + + #[tokio::test] + async fn only_the_listed_outcomes_are_queued() { + let settings = config::tests::settings_with( + "litellm_settings:\n failure_callback: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n", + &[], + ) + .await + .unwrap() + .unwrap(); + let mut log = log("").await; + log.settings = Arc::new(settings); + log.observe(&record()).await; + assert!(log.queue.lock().unwrap().elements.is_empty()); + let mut failed = record(); + failed.outcome.status = 500; + log.observe(&failed).await; + assert_eq!(log.queue.lock().unwrap().elements.len(), 1); + } + + #[tokio::test] + async fn a_full_queue_drops_the_oldest() { + let log = log(" s3_max_queue_size: 2\n").await; + for key in ["a", "b", "c"] { + log.enqueue(element(key, "{}")); + } + let queue = log.queue.lock().unwrap(); + let keys: Vec<_> = queue + .elements + .iter() + .map(|element| element.key.as_str()) + .collect(); + assert_eq!(keys, ["b", "c"]); + } + + #[tokio::test] + async fn an_unreachable_bucket_keeps_everything_for_the_next_flush() { + let log = log(" s3_endpoint_url: http://127.0.0.1:9\n s3_aws_access_key_id: k\n s3_aws_secret_access_key: s\n s3_max_retry_age_seconds: 1\n").await; + let mut old = element("a", "{}"); + old.retrying_since = Some(Instant::now() - Duration::from_secs(60)); + log.enqueue(old); + log.enqueue(element("b", "{}")); + log.flush().await; + let queue = log.queue.lock().unwrap(); + let kept: Vec<_> = queue + .elements + .iter() + .map(|element| element.key.as_str()) + .collect(); + assert_eq!(kept, ["a", "b"]); + assert_eq!(queue.requeued, 2); + } +} diff --git a/litellm-rust/crates/gateway/src/s3_log/payload.rs b/litellm-rust/crates/gateway/src/s3_log/payload.rs new file mode 100644 index 00000000000..8b67803dfb5 --- /dev/null +++ b/litellm-rust/crates/gateway/src/s3_log/payload.rs @@ -0,0 +1,1204 @@ +//! The proxy's `StandardLoggingPayload` for a request, from what the gateway knows of it, and +//! the transforms the proxy applies before logging it: base64 truncation, message redaction, +//! `s3_strip_base64_files` and `s3_log_prompts_only`. + +use std::{ + collections::BTreeMap, + time::{SystemTime, UNIX_EPOCH}, +}; + +use litellm_gateway_hooks::{Record, Route}; +use serde_json::{Map, Value, json}; + +use crate::{ + prometheus::error_class, + spend::{logged_model, spend_logs_metadata, usage_values}, +}; + +const REDACTED: &str = "redacted-by-litellm"; +/// The proxy's `MAX_BASE64_LENGTH_FOR_LOGGING`: longer data URIs are logged as their size. +const MAX_BASE64_LENGTH: usize = 64; +const MAX_TRUNCATION_DEPTH: usize = 20; +/// `DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER`, the depth `s3_strip_base64_files` walks. +const MAX_STRIP_DEPTH: usize = 10; + +/// The `StandardLoggingMetadata` keys, every one present and null when unset. +const METADATA_KEYS: [&str; 35] = [ + "user_api_key_hash", + "user_api_key_alias", + "user_api_key_spend", + "user_api_key_max_budget", + "user_api_key_budget_reset_at", + "user_api_key_user_spend", + "user_api_key_user_max_budget", + "user_api_key_team_spend", + "user_api_key_team_max_budget", + "user_api_key_team_id", + "user_api_key_org_id", + "user_api_key_org_alias", + "user_api_key_project_id", + "user_api_key_project_alias", + "user_api_key_user_id", + "user_api_key_team_alias", + "user_api_key_user_email", + "user_api_key_end_user_id", + "user_api_key_request_route", + "spend_logs_metadata", + "requester_ip_address", + "user_agent", + "requester_metadata", + "prompt_management_metadata", + "applied_guardrails", + "mcp_tool_call_metadata", + "vector_store_request_metadata", + "routing_decision", + "usage_object", + "requester_custom_headers", + "cold_storage_object_key", + "user_api_key_auth_metadata", + "team_alias", + "team_id", + "used_client_oauth_token", +]; + +/// What the proxy's settings change about every payload. +#[derive(Clone, Copy, Debug, Default)] +pub(super) struct Redaction { + /// `litellm_settings.turn_off_message_logging`. + pub messages: bool, + /// `litellm_settings.redact_user_api_key_info`. + pub user_api_key_info: bool, + /// The proxy's request tags carry the client's user agent unless + /// `disable_add_user_agent_to_request_tags` is set. + pub user_agent_tags: bool, +} + +/// The request's `StandardLoggingPayload`. +pub(super) fn standard_logging_payload( + record: &Record, + redaction: Redaction, +) -> Map { + let request = &record.request; + let outcome = &record.outcome; + let deployment = outcome.deployment.as_ref(); + let succeeded = record.succeeded(); + let body: Map = serde_json::from_slice(request.body()).unwrap_or_default(); + let response = match succeeded { + true => response(record), + false => Value::Object(Map::new()), + }; + let id = response + .get("id") + .and_then(Value::as_str) + .map_or_else(|| request.request_id.to_string(), str::to_owned); + let usage = outcome.usage.unwrap_or_default(); + let start = seconds(request.started_at); + let end = seconds(outcome.ended_at); + let completion_start = match (request.stream, outcome.time_to_first_byte) { + (true, Some(first_byte)) => seconds(request.started_at + first_byte), + _ => end, + }; + let status = match succeeded { + true => "success", + false => "failure", + }; + let model = logged_model(record); + let api_base = deployment + .and_then(|deployment| deployment.api_base.as_deref()) + .map(strip_trailing_slash) + .unwrap_or_default(); + let user_agent = request + .client_headers + .get(http::header::USER_AGENT) + .and_then(|value| value.to_str().ok()); + let error_message = (!succeeded) + .then(|| error_message(record)) + .flatten() + .unwrap_or_default(); + let messages = body.get("messages").map(|messages| { + let messages = with_system_prompt(messages.clone(), body.get("system")); + truncate_base64(messages) + }); + let hidden_params = json!({ + "model_id": deployment.map(|deployment| &*deployment.id), + "cache_key": null, + "api_base": deployment.and_then(|deployment| deployment.api_base.as_deref()), + "response_cost": outcome.cost, + "litellm_overhead_time_ms": null, + "additional_headers": null, + "batch_models": null, + "batch_successful_requests": null, + "batch_failed_requests": null, + "litellm_model_name": deployment.map(|_| &model), + "usage_object": null, + }); + let error_information = json!({ + "error_code": if succeeded { String::new() } else { outcome.status.to_string() }, + "error_class": if succeeded { "" } else { error_class(record) }, + "llm_provider": deployment + .filter(|_| !succeeded) + .and_then(|deployment| deployment.provider()) + .unwrap_or_default(), + "traceback": "", + "error_message": &error_message, + }); + let mut payload = json!({ + "id": id, + "litellm_call_id": &*request.request_id, + "trace_id": &*request.request_id, + "session_id": "", + "call_type": request.route.call_type(), + "cache_hit": if request.stream { Value::Bool(false) } else { Value::Null }, + "stream": if request.stream && succeeded { Value::Bool(true) } else { Value::Null }, + "status": status, + "status_fields": {"llm_api_status": status, "guardrail_status": "not_run"}, + "custom_llm_provider": deployment.and_then(|deployment| deployment.provider()), + "saved_cache_cost": 0.0, + "startTime": start, + "endTime": end, + "completionStartTime": completion_start, + "response_time": match request.stream { + true => completion_start - start, + false => end - start, + }, + "model": &model, + "metadata": metadata(record, &body, user_agent, redaction), + "cache_key": null, + "response_cost": outcome.cost.unwrap_or(0.0), + "cost_breakdown": null, + "autorouter_savings": null, + "autorouter_savings_estimate": null, + "autorouter_baseline_observation": null, + }); + let rest = json!({ + "total_tokens": usage.total_tokens(), + "prompt_tokens": usage.prompt_tokens(), + "completion_tokens": usage.output_tokens, + "request_tags": request_tags(&body, user_agent.filter(|_| redaction.user_agent_tags)), + "request_model_access_groups": [], + "end_user": request.end_user.as_deref(), + "api_base": api_base, + "model_group": deployment.map_or("", |deployment| &deployment.model_group), + "model_id": deployment.map_or("", |deployment| &deployment.id), + "requester_ip_address": null, + "user_agent": user_agent, + "messages": messages, + "response": response, + "model_parameters": model_parameters(&body), + "hidden_params": hidden_params, + "model_map_information": {"model_map_key": &model, "model_map_value": null}, + "error_str": (!succeeded && !error_message.is_empty()).then_some(&error_message), + "error_information": error_information, + "response_cost_failure_debug_info": null, + "zero_cost_diagnostic": null, + "guardrail_information": null, + "standard_built_in_tools_params": null, + }); + if let (Value::Object(payload), Value::Object(rest)) = (&mut payload, rest) { + payload.extend(rest); + } + if redaction.messages && message_logging_off(record) { + redact(&mut payload); + } + match payload { + Value::Object(payload) => payload, + _ => unreachable!("the payload is built as an object"), + } +} + +fn seconds(time: SystemTime) -> f64 { + time.duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64() +} + +fn strip_trailing_slash(api_base: &str) -> &str { + match api_base.ends_with("//") { + true => api_base.trim_end_matches('/'), + false => api_base.strip_suffix('/').unwrap_or(api_base), + } +} + +/// The proxy's `append_system_prompt_messages`: a string `system` becomes the first message. +fn with_system_prompt(messages: Value, system: Option<&Value>) -> Value { + let Some(Value::String(system)) = system else { + return messages; + }; + let system_message = json!({"role": "system", "content": system}); + match messages { + Value::Null => Value::Array(vec![system_message]), + Value::Array(messages) => { + if messages.first().is_some_and(|first| { + first.get("role") == Some(&json!("system")) + && first.get("content") == Some(&json!(system)) + }) { + return Value::Array(messages); + } + Value::Array(std::iter::once(system_message).chain(messages).collect()) + } + Value::String(text) => json!([system_message, {"role": "user", "content": text}]), + other => other, + } +} + +/// The proxy redacts unless the client opts out with `litellm-disable-message-redaction`. +fn message_logging_off(record: &Record) -> bool { + !record + .request + .client_headers + .get("litellm-disable-message-redaction") + .is_some_and(|value| !value.is_empty()) +} + +/// The proxy's `_redact_standard_logging_object`. +fn redact(payload: &mut Value) { + if !payload["messages"].is_null() { + payload["messages"] = json!([{"role": "user", "content": REDACTED}]); + } + let response = &mut payload["response"]; + match response { + Value::Null => {} + Value::Object(fields) if fields.get("choices").is_some() => { + if let Some(Value::Array(choices)) = fields.get_mut("choices") { + choices.iter_mut().for_each(redact_choice); + } + } + Value::String(_) => *response = Value::from(REDACTED), + _ => *response = json!({"text": REDACTED}), + } +} + +fn redact_choice(choice: &mut Value) { + let Value::Object(choice) = choice else { + return; + }; + let message = match (choice.contains_key("message"), choice.contains_key("delta")) { + (true, _) => choice.get_mut("message"), + (false, true) => choice.get_mut("delta"), + (false, false) => { + if choice.get("text").is_some_and(|text| !text.is_null()) { + choice.insert("text".into(), Value::from(REDACTED)); + } + return; + } + }; + let Some(Value::Object(message)) = message else { + return; + }; + for field in ["content", "reasoning_content"] { + if message.get(field).is_some_and(|value| !value.is_null()) { + message.insert(field.into(), Value::from(REDACTED)); + } + } + for field in ["thinking_blocks", "audio"] { + if message.contains_key(field) { + message.insert(field.into(), Value::Null); + } + } + if let Some(Value::Array(calls)) = message.get_mut("tool_calls") { + for call in calls { + if let Some(Value::Object(function)) = call.get_mut("function") { + function.insert("arguments".into(), Value::from(REDACTED)); + } + } + } + if let Some(Value::Object(function)) = message.get_mut("function_call") + && function.contains_key("arguments") + { + function.insert("arguments".into(), Value::from(REDACTED)); + } +} + +fn metadata( + record: &Record, + body: &Map, + user_agent: Option<&str>, + redaction: Redaction, +) -> Value { + let mut metadata: Map = METADATA_KEYS + .iter() + .map(|key| ((*key).to_owned(), Value::Null)) + .collect(); + let mut set = |key: &str, value: Value| { + metadata.insert(key.to_owned(), value); + }; + let request = &record.request; + if let Some(key) = record.key().filter(|_| !redaction.user_api_key_info) { + let budgets = &key.budgets; + set("user_api_key_hash", json!(key.hash)); + set("user_api_key_alias", json!(key.alias)); + set("user_api_key_spend", json!(budgets.key.spend)); + set("user_api_key_max_budget", json!(budgets.key.max_budget)); + set( + "user_api_key_user_spend", + json!(budgets.user.map(|user| user.spend)), + ); + set( + "user_api_key_user_max_budget", + json!(budgets.user.and_then(|user| user.max_budget)), + ); + set( + "user_api_key_team_spend", + json!(budgets.team.map(|team| team.spend)), + ); + set( + "user_api_key_team_max_budget", + json!(budgets.team.and_then(|team| team.max_budget)), + ); + set("user_api_key_team_id", json!(key.team_id)); + set("user_api_key_org_id", json!(key.org_id)); + set("user_api_key_user_id", json!(key.user_id)); + set("user_api_key_team_alias", json!(key.team_alias)); + } + if !redaction.user_api_key_info { + set( + "user_api_key_end_user_id", + json!(request.end_user.as_deref()), + ); + set( + "user_api_key_request_route", + json!(request.route.endpoint()), + ); + } + set("spend_logs_metadata", spend_logs_metadata(record)); + set("user_agent", json!(user_agent)); + if let Some(requester) = body.get("metadata").filter(|value| value.is_object()) { + set("requester_metadata", requester.clone()); + } + set( + "usage_object", + record + .outcome + .usage + .map_or_else(|| json!({}), |usage| usage_values(&usage).0), + ); + Value::Object(metadata) +} + +/// The request's `metadata.tags`, then the proxy's user agent tags. +fn request_tags(body: &Map, user_agent: Option<&str>) -> Vec { + let tags = body + .get("metadata") + .and_then(|metadata| metadata.get("tags")) + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let agent = user_agent.into_iter().flat_map(|agent| { + agent + .split_once('/') + .map(|(product, _)| format!("User-Agent: {product}")) + .into_iter() + .chain(std::iter::once(format!("User-Agent: {agent}"))) + .map(Value::from) + }); + tags.into_iter().chain(agent).collect() +} + +/// The request's model parameters: everything but the model, the prompt and the metadata. +fn model_parameters(body: &Map) -> Map { + body.iter() + .filter(|(key, _)| { + !matches!( + key.as_str(), + "model" | "messages" | "system" | "prompt" | "input" | "metadata" + ) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() +} + +/// The provider's or the gateway's error message, from an error body when there is one. +fn error_message(record: &Record) -> Option { + let body: Value = serde_json::from_slice(record.response_body.as_deref()?).ok()?; + let error = body.get("error").unwrap_or(&body); + error + .get("message") + .and_then(Value::as_str) + .or_else(|| error.as_str()) + .map(str::to_owned) +} + +/// The response the client got: the body itself, or the stream put back together the way the +/// proxy assembles it. +fn response(record: &Record) -> Value { + let Some(body) = record.response_body.as_deref() else { + return Value::Object(Map::new()); + }; + if let Ok(value @ Value::Object(_)) = serde_json::from_slice::(body) { + return value; + } + let Ok(text) = std::str::from_utf8(body) else { + return Value::Object(Map::new()); + }; + let events = text + .lines() + .filter_map(|line| line.strip_prefix("data:")) + .map(str::trim) + .filter_map(|data| serde_json::from_str::(data).ok()); + let assembled = match record.request.route { + Route::ChatCompletions => chat_completion(events), + Route::Messages => anthropic_message(events), + }; + assembled.unwrap_or_else(|| Value::Object(Map::new())) +} + +#[derive(Default)] +struct ChoiceParts { + role: Option, + content: Option, + reasoning: Option, + finish_reason: Value, + tool_calls: BTreeMap, +} + +/// A chat completion from its chunks. +fn chat_completion(events: impl Iterator) -> Option { + let mut head: Option = None; + let mut usage = Value::Null; + let mut choices: BTreeMap = BTreeMap::new(); + for chunk in events { + if head.is_none() { + head = Some(json!({ + "id": chunk.get("id"), + "created": chunk.get("created"), + "model": chunk.get("model"), + "system_fingerprint": chunk.get("system_fingerprint"), + })); + } + if chunk.get("usage").is_some_and(|usage| !usage.is_null()) { + usage = chunk["usage"].clone(); + } + for choice in chunk["choices"].as_array().into_iter().flatten() { + let parts = choices + .entry(choice["index"].as_u64().unwrap_or(0)) + .or_default(); + let delta = &choice["delta"]; + if let Some(role) = delta.get("role").filter(|role| !role.is_null()) { + parts.role = Some(role.clone()); + } + if let Some(text) = delta["content"].as_str() { + parts.content.get_or_insert_default().push_str(text); + } + if let Some(text) = delta["reasoning_content"].as_str() { + parts.reasoning.get_or_insert_default().push_str(text); + } + if !choice["finish_reason"].is_null() { + parts.finish_reason = choice["finish_reason"].clone(); + } + for call in delta["tool_calls"].as_array().into_iter().flatten() { + let slot = parts + .tool_calls + .entry(call["index"].as_u64().unwrap_or(0)) + .or_insert_with(|| { + json!({"id": null, "type": "function", "function": {"name": null, "arguments": ""}}) + }); + for field in ["id", "type"] { + if let Some(value) = call.get(field).filter(|value| !value.is_null()) { + slot[field] = value.clone(); + } + } + if let Some(name) = call["function"]["name"].as_str() { + slot["function"]["name"] = Value::from(name); + } + if let Some(arguments) = call["function"]["arguments"].as_str() { + let joined = format!( + "{}{arguments}", + slot["function"]["arguments"].as_str().unwrap_or_default() + ); + slot["function"]["arguments"] = Value::from(joined); + } + } + } + } + let mut completion = head?; + completion["object"] = Value::from("chat.completion"); + completion["choices"] = choices + .into_iter() + .map(|(index, parts)| { + let tool_calls: Vec = parts.tool_calls.into_values().collect(); + let mut message = json!({ + "content": parts.content, + "role": parts.role.unwrap_or_else(|| Value::from("assistant")), + "tool_calls": if tool_calls.is_empty() { Value::Null } else { Value::Array(tool_calls) }, + "function_call": null, + }); + if let Some(reasoning) = parts.reasoning { + message["reasoning_content"] = Value::from(reasoning); + } + json!({"finish_reason": parts.finish_reason, "index": index, "message": message}) + }) + .collect(); + completion["usage"] = usage; + Some(completion) +} + +/// An Anthropic message from its stream events. +fn anthropic_message(events: impl Iterator) -> Option { + let mut message: Option = None; + let mut blocks: BTreeMap = BTreeMap::new(); + let mut partial_json: BTreeMap = BTreeMap::new(); + for event in events { + let index = event["index"].as_u64().unwrap_or(0); + match event["type"].as_str() { + Some("message_start") => message = Some(event["message"].clone()), + Some("content_block_start") => { + blocks.insert(index, event["content_block"].clone()); + } + Some("content_block_delta") => { + let block = blocks.entry(index).or_insert_with(|| json!({})); + let delta = &event["delta"]; + let append = |block: &mut Value, field: &str, text: &str| { + let joined = format!("{}{text}", block[field].as_str().unwrap_or_default()); + block[field] = Value::from(joined); + }; + match delta["type"].as_str() { + Some("text_delta") => { + append(block, "text", delta["text"].as_str().unwrap_or_default()) + } + Some("thinking_delta") => append( + block, + "thinking", + delta["thinking"].as_str().unwrap_or_default(), + ), + Some("signature_delta") => block["signature"] = delta["signature"].clone(), + Some("input_json_delta") => partial_json + .entry(index) + .or_default() + .push_str(delta["partial_json"].as_str().unwrap_or_default()), + _ => {} + } + } + Some("message_delta") => { + let Some(message) = message.as_mut() else { + continue; + }; + if let Some(delta) = event["delta"].as_object() { + for (field, value) in delta { + message[field.as_str()] = value.clone(); + } + } + if let (Some(usage), Some(Value::Object(total))) = + (event["usage"].as_object(), message.get_mut("usage")) + { + for (field, value) in usage { + if !value.is_null() { + total.insert(field.clone(), value.clone()); + } + } + } + } + _ => {} + } + } + let mut message = message?; + for (index, json) in partial_json { + if let (Some(block), Ok(input)) = (blocks.get_mut(&index), serde_json::from_str(&json)) { + block["input"] = input; + } + } + message["content"] = Value::Array(blocks.into_values().collect()); + Some(message) +} + +/// Replaces each data URI in `text` that `matcher` accepts; `matcher` gets the text from just +/// after `data:` and returns how much of it the URI spans and its replacement. +fn replace_data_uris(text: &str, matcher: impl Fn(&str) -> Option<(usize, String)>) -> String { + let mut out = String::with_capacity(text.len()); + let mut rest = text; + while let Some(at) = rest.find("data:") { + let after = &rest[at + 5..]; + match matcher(after) { + Some((length, replacement)) => { + out.push_str(&rest[..at]); + out.push_str(&replacement); + rest = &after[length..]; + } + None => { + out.push_str(&rest[..at + 5]); + rest = after; + } + } + } + out.push_str(rest); + out +} + +/// The length of the leading run of `text` whose characters satisfy `accept`. +fn run(text: &str, accept: impl Fn(char) -> bool) -> usize { + text.find(|c: char| !accept(c)).unwrap_or(text.len()) +} + +fn is_base64(c: char) -> bool { + c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '=') +} + +/// The proxy's `truncate_base64_in_messages`: a data URI with more than 64 base64 characters +/// is logged as its size. +fn truncate_base64(value: Value) -> Value { + fn walk(value: Value, depth: usize) -> Value { + match value { + Value::String(text) => Value::String(truncate_text(&text)), + Value::Array(items) if depth <= MAX_TRUNCATION_DEPTH => Value::Array( + items + .into_iter() + .map(|item| match item { + Value::String(_) => walk(item, depth), + item => walk(item, depth + 1), + }) + .collect(), + ), + Value::Object(fields) if depth <= MAX_TRUNCATION_DEPTH => Value::Object( + fields + .into_iter() + .map(|(key, item)| match item { + Value::String(_) => (key, walk(item, depth)), + item => (key, walk(item, depth + 1)), + }) + .collect(), + ), + other => other, + } + } + walk(value, 0) +} + +fn truncate_text(text: &str) -> String { + replace_data_uris(text, |after| { + let mime = run(after, |c| !matches!(c, ';' | ',') && !c.is_whitespace()); + if mime == 0 || mime > 255 || !after[mime..].starts_with(";base64,") { + return None; + } + let payload_start = mime + ";base64,".len(); + let payload = run(&after[payload_start..], is_base64); + if payload == 0 { + return None; + } + let length = payload_start + payload; + if payload <= MAX_BASE64_LENGTH { + return Some((length, format!("data:{}", &after[..length]))); + } + Some(( + length, + format!( + "data:{};base64,[base64_data truncated: {}]", + &after[..mime], + base64_size(payload) + ), + )) + }) +} + +/// The proxy's `format_base64_size`. +fn base64_size(characters: usize) -> String { + let bytes = characters as f64 * 3.0 / 4.0; + match bytes { + bytes if bytes >= 1024.0 * 1024.0 => format!("{:.2}MB", bytes / (1024.0 * 1024.0)), + bytes if bytes >= 1024.0 => format!("{:.1}KB", bytes / 1024.0), + bytes => format!("{}B", bytes as u64), + } +} + +/// `s3_strip_base64_files`, as the proxy's `_strip_base64_from_messages_sync`: drops file and +/// non-text content parts and redacts inline base64 in what is left of each message. +pub(super) fn strip_base64_files(payload: &mut Map) { + let Some(Value::Array(messages)) = payload.get_mut("messages") else { + return; + }; + if messages.is_empty() { + return; + } + let stripped: Vec = std::mem::take(messages) + .into_iter() + .filter_map(|message| match message { + Value::Object(message) => Some(Value::Object(strip_message(message))), + _ => None, + }) + .collect(); + *messages = stripped; +} + +fn strip_message(mut message: Map) -> Map { + let content = match message.remove("content") { + Some(Value::Array(parts)) => Value::Array( + parts + .into_iter() + .filter(keep_part) + .map(|part| redact_base64(part, 0)) + .collect(), + ), + other => redact_base64(other.unwrap_or(Value::Null), 0), + }; + let mut stripped: Map = message + .into_iter() + .map(|(key, value)| (key, redact_base64(value, 0))) + .collect(); + stripped.insert("content".into(), content); + stripped +} + +fn keep_part(part: &Value) -> bool { + let Value::Object(part) = part else { + return true; + }; + !part.contains_key("file") + && part + .get("type") + .and_then(Value::as_str) + .is_none_or(|kind| kind == "text") +} + +fn redact_base64(value: Value, depth: usize) -> Value { + if depth > MAX_STRIP_DEPTH { + return Value::from("[MAX_DEPTH_REACHED]"); + } + match value { + Value::String(text) => Value::String(redact_text(&text)), + Value::Array(items) => Value::Array( + items + .into_iter() + .map(|item| redact_base64(item, depth + 1)) + .collect(), + ), + Value::Object(fields) => Value::Object( + fields + .into_iter() + .map(|(key, item)| (key, redact_base64(item, depth + 1))) + .collect(), + ), + other => other, + } +} + +/// Replaces `data:(application|image|audio|video)/;base64,` with +/// `[BASE64_REDACTED]`, the payload running over base64 characters and whitespace. +fn redact_text(text: &str) -> String { + replace_data_uris(text, |after| { + let kind = ["application/", "image/", "audio/", "video/"] + .into_iter() + .find(|kind| after.starts_with(kind))?; + let subtype = run(&after[kind.len()..], |c| { + c.is_ascii_alphanumeric() || matches!(c, '.' | '+' | '-') + }); + let rest = &after[kind.len() + subtype..]; + if subtype == 0 || !rest.starts_with(";base64,") { + return None; + } + let payload_start = kind.len() + subtype + ";base64,".len(); + let payload = run(&after[payload_start..], |c| { + is_base64(c) || c.is_whitespace() + }); + (payload > 0).then(|| (payload_start + payload, "[BASE64_REDACTED]".to_owned())) + }) +} + +/// `s3_log_prompts_only`: the payload without the response. +pub(super) fn prompts_only(payload: &mut Map) { + payload.insert("response".into(), Value::Null); +} + +#[cfg(test)] +pub(super) mod tests { + use std::{sync::Arc, time::Duration}; + + use axum::body::Bytes; + use http::HeaderMap; + use litellm_gateway_hooks::{ + Annotations, DeploymentFacts, KeyFacts, Outcome, RequestFacts, Usage, + }; + use rstest::rstest; + + use super::*; + + pub(crate) const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-x","choices":[{"index":0,"message":{"role":"assistant","content":"hi","tool_calls":[{"id":"t1","type":"function","function":{"name":"f","arguments":"{\"a\":1}"}}]},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":6,"total_tokens":11}}"#; + + pub(crate) fn record( + route: Route, + status: u16, + stream: bool, + request: Value, + response: Option<&str>, + ) -> Record { + let started_at = UNIX_EPOCH + Duration::from_secs(1_000); + let mut headers = HeaderMap::new(); + headers.insert("user-agent", "client/1.0".parse().unwrap()); + let key = KeyFacts { + hash: "hash-1".into(), + alias: Some("team-key".into()), + user_id: Some("user-1".into()), + team_id: Some("team-1".into()), + team_alias: Some("research".into()), + ..KeyFacts::default() + }; + let mut facts = RequestFacts::new( + "req-1".into(), + route, + Some(Arc::new(key)), + Some("end-1".into()), + "gpt".into(), + Arc::new(headers), + stream, + None, + Bytes::from(request.to_string()), + ); + facts.started_at = started_at; + Record { + request: Arc::new(facts), + outcome: Outcome { + deployment: Some(DeploymentFacts { + id: "gpt-1".into(), + model_group: "gpt".into(), + model: "openai/gpt-x".into(), + provider: None, + api_base: Some("https://api.example/v1/".into()), + }), + status, + usage: (status == 200).then_some(Usage { + input_tokens: 5, + output_tokens: 6, + ..Usage::default() + }), + cost: (status == 200).then_some(0.5), + time_to_first_byte: Some(Duration::from_millis(250)), + upstream_started_at: None, + ended_at: started_at + Duration::from_secs(2), + interrupted: false, + }, + annotations: Annotations::default(), + response_body: response.map(|body| Bytes::from(body.to_owned())), + } + } + + fn chat_request() -> Value { + json!({ + "model": "gpt", + "temperature": 0.2, + "metadata": {"tags": ["prod"]}, + "messages": [{"role": "user", "content": "hi"}], + }) + } + + #[test] + fn a_success_carries_the_proxys_fields() { + let record = record( + Route::ChatCompletions, + 200, + false, + chat_request(), + Some(COMPLETION), + ); + let payload = Value::Object(standard_logging_payload( + &record, + Redaction { + user_agent_tags: true, + ..Redaction::default() + }, + )); + assert_eq!(payload["id"], "chatcmpl-1"); + assert_eq!(payload["litellm_call_id"], "req-1"); + assert_eq!(payload["call_type"], "acompletion"); + assert_eq!(payload["status"], "success"); + assert_eq!(payload["custom_llm_provider"], "openai"); + assert_eq!(payload["model"], "gpt-x"); + assert_eq!(payload["model_group"], "gpt"); + assert_eq!(payload["model_id"], "gpt-1"); + assert_eq!(payload["api_base"], "https://api.example/v1"); + assert_eq!(payload["response_cost"], 0.5); + assert_eq!(payload["prompt_tokens"], 5); + assert_eq!(payload["completion_tokens"], 6); + assert_eq!(payload["total_tokens"], 11); + assert_eq!(payload["startTime"], 1000.0); + assert_eq!(payload["endTime"], 1002.0); + assert_eq!(payload["completionStartTime"], 1002.0); + assert_eq!(payload["response_time"], 2.0); + assert_eq!(payload["stream"], Value::Null); + assert_eq!(payload["end_user"], "end-1"); + assert_eq!( + payload["messages"], + json!([{"role": "user", "content": "hi"}]) + ); + assert_eq!( + payload["response"], + serde_json::from_str::(COMPLETION).unwrap() + ); + assert_eq!(payload["model_parameters"], json!({"temperature": 0.2})); + assert_eq!( + payload["request_tags"], + json!(["prod", "User-Agent: client", "User-Agent: client/1.0"]) + ); + let metadata = &payload["metadata"]; + assert_eq!(metadata.as_object().unwrap().len(), METADATA_KEYS.len()); + assert_eq!(metadata["user_api_key_hash"], "hash-1"); + assert_eq!(metadata["user_api_key_alias"], "team-key"); + assert_eq!(metadata["user_api_key_team_alias"], "research"); + assert_eq!(metadata["user_api_key_team_id"], "team-1"); + assert_eq!(metadata["user_api_key_user_id"], "user-1"); + assert_eq!(metadata["user_api_key_end_user_id"], "end-1"); + assert_eq!(metadata["requester_metadata"], json!({"tags": ["prod"]})); + assert_eq!(metadata["usage_object"]["total_tokens"], 11); + assert_eq!(payload["error_information"]["error_code"], ""); + assert_eq!(payload["error_str"], Value::Null); + } + + #[test] + fn a_failure_carries_the_error_and_an_empty_response() { + let error = r#"{"error":{"message":"slow down","type":"rate_limit_error"}}"#; + let record = record( + Route::ChatCompletions, + 429, + false, + chat_request(), + Some(error), + ); + let payload = Value::Object(standard_logging_payload(&record, Redaction::default())); + assert_eq!(payload["status"], "failure"); + assert_eq!(payload["id"], "req-1"); + assert_eq!(payload["response"], json!({})); + assert_eq!(payload["response_cost"], 0.0); + assert_eq!(payload["total_tokens"], 0); + assert_eq!(payload["error_str"], "slow down"); + assert_eq!( + payload["error_information"], + json!({ + "error_code": "429", + "error_class": "RateLimitError", + "llm_provider": "openai", + "traceback": "", + "error_message": "slow down", + }) + ); + assert_eq!(payload["metadata"]["usage_object"], json!({})); + assert_eq!(payload["request_tags"], json!(["prod"])); + } + + #[test] + fn a_chat_stream_is_put_back_together() { + let stream = concat!( + "data: {\"id\":\"chatcmpl-9\",\"created\":7,\"model\":\"gpt-x\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"he\"},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-9\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"llo\",\"tool_calls\":[{\"index\":0,\"id\":\"t1\",\"function\":{\"name\":\"f\",\"arguments\":\"{\\\"a\\\"\"}}]},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-9\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\":1}\"}}]},\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n", + "data: [DONE]\n\n", + ); + let record = record( + Route::ChatCompletions, + 200, + true, + chat_request(), + Some(stream), + ); + let payload = Value::Object(standard_logging_payload(&record, Redaction::default())); + assert_eq!(payload["id"], "chatcmpl-9"); + assert_eq!(payload["stream"], true); + assert_eq!(payload["cache_hit"], false); + assert_eq!(payload["completionStartTime"], 1000.25); + assert_eq!(payload["response_time"], 0.25); + let response = &payload["response"]; + assert_eq!(response["object"], "chat.completion"); + assert_eq!(response["created"], 7); + assert_eq!(response["usage"]["total_tokens"], 3); + assert_eq!( + response["choices"], + json!([{ + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": "hello", + "role": "assistant", + "tool_calls": [{"id": "t1", "type": "function", "function": {"name": "f", "arguments": "{\"a\":1}"}}], + "function_call": null, + }, + }]) + ); + } + + #[test] + fn a_messages_stream_is_put_back_together_with_its_system_prompt() { + let stream = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-x\",\"stop_reason\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":1}}}\n\n", + "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n", + "data: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"tool_use\",\"id\":\"t\",\"name\":\"f\",\"input\":{}}}\n\n", + "data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"a\\\":\"}}\n\n", + "data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"1}\"}}\n\n", + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":42}}\n\n", + ); + let request = json!({"model": "claude", "system": "be brief", "messages": [{"role": "user", "content": "hi"}]}); + let record = record(Route::Messages, 200, true, request, Some(stream)); + let payload = Value::Object(standard_logging_payload(&record, Redaction::default())); + assert_eq!(payload["id"], "msg_1"); + assert_eq!(payload["call_type"], "anthropic_messages"); + assert_eq!( + payload["messages"], + json!([{"role": "system", "content": "be brief"}, {"role": "user", "content": "hi"}]) + ); + let response = &payload["response"]; + assert_eq!(response["stop_reason"], "tool_use"); + assert_eq!( + response["usage"], + json!({"input_tokens": 10, "output_tokens": 42}) + ); + assert_eq!( + response["content"], + json!([{"type": "text", "text": "hi"}, {"type": "tool_use", "id": "t", "name": "f", "input": {"a": 1}}]) + ); + } + + #[test] + fn turning_off_message_logging_redacts_messages_and_response() { + let record = record( + Route::ChatCompletions, + 200, + false, + chat_request(), + Some(COMPLETION), + ); + let redaction = Redaction { + messages: true, + ..Redaction::default() + }; + let payload = Value::Object(standard_logging_payload(&record, redaction)); + assert_eq!( + payload["messages"], + json!([{"role": "user", "content": REDACTED}]) + ); + let message = &payload["response"]["choices"][0]["message"]; + assert_eq!(message["content"], REDACTED); + assert_eq!(message["tool_calls"][0]["function"]["arguments"], REDACTED); + assert_eq!(message["tool_calls"][0]["function"]["name"], "f"); + assert_eq!(payload["response"]["usage"]["total_tokens"], 11); + } + + #[test] + fn a_messages_response_is_redacted_to_text() { + let message = + r#"{"id":"msg_1","type":"message","content":[{"type":"text","text":"secret"}]}"#; + let record = record( + Route::Messages, + 200, + false, + json!({"messages": []}), + Some(message), + ); + let redaction = Redaction { + messages: true, + ..Redaction::default() + }; + let payload = Value::Object(standard_logging_payload(&record, redaction)); + assert_eq!(payload["response"], json!({"text": REDACTED})); + assert_eq!(payload["id"], "msg_1"); + } + + #[test] + fn the_client_can_opt_out_of_redaction() { + let mut record = record( + Route::ChatCompletions, + 200, + false, + chat_request(), + Some(COMPLETION), + ); + let mut headers = HeaderMap::new(); + headers.insert("litellm-disable-message-redaction", "true".parse().unwrap()); + let request = Arc::get_mut(&mut record.request).unwrap(); + request.client_headers = Arc::new(headers); + let redaction = Redaction { + messages: true, + ..Redaction::default() + }; + let payload = Value::Object(standard_logging_payload(&record, redaction)); + assert_eq!( + payload["messages"], + json!([{"role": "user", "content": "hi"}]) + ); + } + + #[test] + fn redacting_key_info_leaves_the_fields_null() { + let record = record( + Route::ChatCompletions, + 200, + false, + chat_request(), + Some(COMPLETION), + ); + let redaction = Redaction { + user_api_key_info: true, + ..Redaction::default() + }; + let payload = Value::Object(standard_logging_payload(&record, redaction)); + assert_eq!(payload["metadata"]["user_api_key_team_alias"], Value::Null); + assert_eq!(payload["metadata"]["user_api_key_hash"], Value::Null); + } + + #[rstest] + #[case::short("data:image/png;base64,AAAA", "data:image/png;base64,AAAA")] + #[case::long( + &format!("see data:image/png;base64,{} end", "A".repeat(4000)), + "see data:image/png;base64,[base64_data truncated: 2.9KB] end" + )] + #[case::bytes(&format!("data:audio/wav;base64,{}", "A".repeat(100)), "data:audio/wav;base64,[base64_data truncated: 75B]")] + #[case::not_base64("data:text/plain,hello", "data:text/plain,hello")] + fn long_data_uris_in_messages_are_logged_as_their_size( + #[case] text: &str, + #[case] expected: &str, + ) { + assert_eq!( + truncate_base64(json!([{"content": [{"image_url": {"url": text}}]}])), + json!([{"content": [{"image_url": {"url": expected}}]}]) + ); + } + + #[test] + fn stripping_base64_drops_file_parts_and_redacts_inline_data() { + let mut payload = Map::new(); + payload.insert( + "messages".into(), + json!([ + {"role": "user", "content": [ + {"type": "text", "text": "look data:image/png;base64,QUJD"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,QUJD"}}, + {"file": {"file_data": "x"}}, + ]}, + {"role": "assistant", "content": "data:application/pdf;base64,QUJD\nQUJD"}, + {"role": "tool"}, + "not a message", + ]), + ); + strip_base64_files(&mut payload); + assert_eq!( + payload["messages"], + json!([ + {"role": "user", "content": [{"type": "text", "text": "look [BASE64_REDACTED]"}]}, + {"role": "assistant", "content": "[BASE64_REDACTED]"}, + {"role": "tool", "content": null}, + ]) + ); + } + + #[test] + fn prompts_only_drops_the_response() { + let record = record( + Route::ChatCompletions, + 200, + false, + chat_request(), + Some(COMPLETION), + ); + let mut payload = standard_logging_payload(&record, Redaction::default()); + prompts_only(&mut payload); + assert_eq!(payload["response"], Value::Null); + assert_eq!( + payload["messages"], + json!([{"role": "user", "content": "hi"}]) + ); + assert_eq!(payload["id"], "chatcmpl-1"); + } + + #[test] + fn without_a_response_body_the_id_is_the_request_id() { + let record = record(Route::ChatCompletions, 200, false, chat_request(), None); + let payload = standard_logging_payload(&record, Redaction::default()); + assert_eq!(payload["id"], "req-1"); + assert_eq!(payload["response"], json!({})); + } +} diff --git a/litellm-rust/crates/gateway/src/s3_log/upload.rs b/litellm-rust/crates/gateway/src/s3_log/upload.rs new file mode 100644 index 00000000000..e78f0d89196 --- /dev/null +++ b/litellm-rust/crates/gateway/src/s3_log/upload.rs @@ -0,0 +1,278 @@ +//! One signed PUT per object, retried as the proxy's `_put_with_retries` retries it: three +//! attempts, a transport error or a 403, 429, 500 or 503 retried after a jittered backoff, and +//! a 400 or 403 that names an object-specific error code dropped instead of kept for later. + +use std::{ + collections::BTreeMap, + sync::Arc, + time::{Duration, Instant, SystemTime}, +}; + +use base64::Engine; +use litellm_auth_aws::{AwsAuthService, Credentials, sign_s3}; +use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool, Verify}; +use md5::Md5; +use sha2::{Digest, Sha256}; + +use super::{Element, config::Settings, key::object_url}; + +const ATTEMPTS: u32 = 3; +const RETRYABLE: [u16; 4] = [403, 429, 500, 503]; +/// Codes S3 answers that mean this object can never be stored, so retrying it is pointless. +const TERMINAL_CODES: [&str; 7] = [ + "EntityTooLarge", + "InvalidArgument", + "MalformedXML", + "InvalidDigest", + "KeyTooLongError", + "BadDigest", + "InvalidRequest", +]; +/// Credentials without an expiry are resolved again after this long. +const CREDENTIALS_TTL: Duration = Duration::from_secs(600); +/// Credentials are refreshed this long before they expire. +const CREDENTIALS_MARGIN: Duration = Duration::from_secs(300); + +/// How one object's upload ended. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum Upload { + Delivered, + /// Failed for now; it goes back on the queue. + Retry, + /// S3 rejected the object itself. + Dropped, +} + +/// What one attempt got back. +enum Attempt { + Response { status: u16, body: String }, + Transport(String), +} + +impl Attempt { + fn is_terminal(&self) -> bool { + matches!(self, Self::Response { status, body } if is_terminal(*status, body)) + } +} + +/// A 400 or 403 whose `` S3 reserves for a rejection of this object. +fn is_terminal(status: u16, body: &str) -> bool { + matches!(status, 400 | 403) + && error_code(body).is_some_and(|code| TERMINAL_CODES.contains(&code)) +} + +fn error_code(body: &str) -> Option<&str> { + let start = body.find("")? + "".len(); + let length = body[start..].find("")?; + Some(&body[start..start + length]).filter(|code| !code.is_empty() && !code.contains('<')) +} + +/// `(1 << attempt) * (0.5 + random())` seconds. +fn backoff(attempt: u32) -> Duration { + Duration::from_secs_f64(f64::from(1u32 << attempt) * (0.5 + rand::random::())) +} + +pub(super) struct Uploader { + settings: Arc, + pool: Arc, + http: HttpClientConfig, + auth: AwsAuthService, + credentials: tokio::sync::Mutex>, +} + +impl Uploader { + pub fn new( + settings: Arc, + pool: Arc, + http: &HttpClientConfig, + ) -> Self { + let verify = match &settings.verify { + None => http.verify.clone(), + Some(litellm_http::SslVerify::Enabled) => Verify::BuiltInRoots, + Some(litellm_http::SslVerify::Disabled) => Verify::Disabled, + Some(litellm_http::SslVerify::CaBundle(path)) => Verify::CaBundle(path.clone()), + }; + Self { + settings, + pool, + http: HttpClientConfig { + verify, + ..http.clone() + }, + auth: AwsAuthService::default(), + credentials: tokio::sync::Mutex::new(None), + } + } + + pub async fn upload(&self, element: &Element) -> Upload { + let url = object_url(&self.settings, &element.key); + let headers = self.headers(element); + let mut attempt = 0; + loop { + let result = match self.put(&url, element, &headers).await { + Ok(result) => result, + Err(error) => { + tracing::error!(%error, key = %element.key, "s3 logging: cannot sign the upload"); + return Upload::Retry; + } + }; + let retry = match &result { + Attempt::Transport(_) => true, + Attempt::Response { status, .. } => { + RETRYABLE.contains(status) + && !(self.settings.drop_on_terminal_error && result.is_terminal()) + } + }; + if retry && attempt + 1 < ATTEMPTS { + let wait = backoff(attempt); + let reason = match &result { + Attempt::Transport(error) => error.clone(), + Attempt::Response { status, .. } => status.to_string(), + }; + tracing::debug!(attempt = attempt + 1, %reason, ?wait, key = %element.key, "s3 logging: upload failed, retrying"); + tokio::time::sleep(wait).await; + attempt += 1; + continue; + } + return match result { + Attempt::Response { status, .. } if (200..300).contains(&status) => { + Upload::Delivered + } + Attempt::Response { status, body } => { + tracing::error!(status, body = %body.chars().take(512).collect::(), key = %element.key, "s3 logging: upload failed"); + if self.settings.drop_on_terminal_error && is_terminal(status, &body) { + tracing::warn!(status, key = %element.key, "s3 logging: dropping the object, S3 rejected it"); + return Upload::Dropped; + } + Upload::Retry + } + Attempt::Transport(error) => { + tracing::error!(%error, key = %element.key, "s3 logging: upload failed"); + Upload::Retry + } + }; + } + } + + /// The proxy's PUT headers, before signing. + fn headers(&self, element: &Element) -> BTreeMap { + let body = &element.body; + let sha256: String = Sha256::digest(body) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect(); + let md5 = + base64::engine::general_purpose::STANDARD.encode(::digest(body)); + let mut headers = BTreeMap::from([ + ("Content-Type".to_owned(), element.content_type.to_owned()), + ("Content-MD5".to_owned(), md5), + ("x-amz-content-sha256".to_owned(), sha256), + ("Content-Language".to_owned(), "en".to_owned()), + ( + "Content-Disposition".to_owned(), + format!("inline; filename=\"{}\"", element.download_name), + ), + ( + "Cache-Control".to_owned(), + "private, immutable, max-age=31536000, s-maxage=0".to_owned(), + ), + ]); + if let Some(algorithm) = &self.settings.server_side_encryption { + headers.insert("x-amz-server-side-encryption".into(), algorithm.clone()); + } + if let Some(key_id) = &self.settings.sse_kms_key_id { + headers.insert( + "x-amz-server-side-encryption-aws-kms-key-id".into(), + key_id.clone(), + ); + } + headers + } + + async fn put( + &self, + url: &str, + element: &Element, + headers: &BTreeMap, + ) -> Result> { + let credentials = self.credentials().await?; + let signature = sign_s3( + "PUT", + url, + &element.body, + headers, + &self.settings.region, + &credentials, + SystemTime::now(), + )?; + let client = self.pool.client(&self.http, ClientVariant::Provider)?; + let mut request = client + .put(url) + .timeout(self.settings.attempt_timeout) + .body(element.body.clone()); + for (name, value) in headers.iter().chain(&signature) { + request = request.header(name, value); + } + Ok(match request.send().await { + Ok(response) => Attempt::Response { + status: response.status().as_u16(), + body: response.text().await.unwrap_or_default(), + }, + Err(error) => Attempt::Transport(error.to_string()), + }) + } + + async fn credentials(&self) -> Result { + let mut cached = self.credentials.lock().await; + if let Some((credentials, until)) = cached.as_ref() + && Instant::now() < *until + { + return Ok(credentials.clone()); + } + let env = |name: &str| std::env::var(name).ok(); + let credentials = self + .auth + .resolve_credentials(self.settings.auth.clone(), &env) + .await?; + let until = credentials + .expiry() + .and_then(|expiry| expiry.duration_since(SystemTime::now()).ok()) + .map_or(CREDENTIALS_TTL, |left| { + left.saturating_sub(CREDENTIALS_MARGIN) + }); + *cached = Some((credentials.clone(), Instant::now() + until)); + Ok(credentials) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::too_large(400, "EntityTooLarge", true)] + #[case::bad_digest(403, "BadDigest", true)] + #[case::access_denied(403, "AccessDenied", false)] + #[case::slow_down(503, "InvalidRequest", false)] + #[case::not_xml(400, "bad request", false)] + #[case::empty_code(400, "", false)] + fn only_object_rejections_are_terminal( + #[case] status: u16, + #[case] body: &str, + #[case] expected: bool, + ) { + assert_eq!(is_terminal(status, body), expected); + } + + #[rstest] + #[case(0, 0.5, 1.5)] + #[case(1, 1.0, 3.0)] + fn backoff_doubles_with_jitter(#[case] attempt: u32, #[case] low: f64, #[case] high: f64) { + for _ in 0..50 { + let wait = backoff(attempt).as_secs_f64(); + assert!((low..high).contains(&wait), "{wait}"); + } + } +} diff --git a/litellm-rust/crates/gateway/src/spend.rs b/litellm-rust/crates/gateway/src/spend.rs index 56491896179..92e949ed71e 100644 --- a/litellm-rust/crates/gateway/src/spend.rs +++ b/litellm-rust/crates/gateway/src/spend.rs @@ -597,7 +597,7 @@ fn metadata(record: &Record) -> Value { /// `spend_logs_metadata` as the proxy merges it: the team's, then the key's, then what hooks /// attached for the request, later keys winning. -fn spend_logs_metadata(record: &Record) -> Value { +pub(crate) fn spend_logs_metadata(record: &Record) -> Value { let mut merged = Map::new(); if let Some(key) = record.key() { for metadata in [&key.team_metadata, &key.metadata] { @@ -617,7 +617,7 @@ fn spend_logs_metadata(record: &Record) -> Value { /// The proxy's usage object for a response, and its `additional_usage_values`: everything but /// the three totals. -fn usage_values(usage: &Usage) -> (Value, Value) { +pub(crate) fn usage_values(usage: &Usage) -> (Value, Value) { let mut additional = Map::new(); additional.insert( "prompt_tokens_details".into(), diff --git a/litellm-rust/crates/gateway/tests/s3_log.rs b/litellm-rust/crates/gateway/tests/s3_log.rs new file mode 100644 index 00000000000..629c217e75b --- /dev/null +++ b/litellm-rust/crates/gateway/tests/s3_log.rs @@ -0,0 +1,313 @@ +//! The `s3_v2` observer against a mock S3: a request ends up as a signed PUT under the proxy's +//! key layout, a 503 is retried, and an object S3 rejects is dropped instead of kept. + +use std::{ + collections::{BTreeMap, HashMap}, + sync::Arc, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use base64::Engine; +use litellm_auth_aws::{Credentials, sign_s3}; +use litellm_config::Config; +use litellm_gateway::s3_log::S3Log; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture, +}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; +use tokio::net::TcpListener; +use tower::ServiceExt; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{method, path_regex}, +}; + +const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-x","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":6,"total_tokens":11}}"#; +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":3,"output_tokens":4}}"#; +const OBJECT_KEY: &str = + r"^/logs/audit/research/\d{4}-\d{2}-\d{2}/time-\d{2}-\d{2}-\d{2}-\d{6}_chatcmpl-1\.json$"; + +/// Answers Anthropic and OpenAI calls. +async fn upstream() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(|request: Request| async move { + let body = match request.uri().path().ends_with("/chat/completions") { + true => COMPLETION, + false => MESSAGE, + }; + ([(CONTENT_TYPE, "application/json")], body).into_response() + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + format!("http://{address}") +} + +struct Keys(HashMap); + +impl KeyLookup for Keys { + fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> { + Box::pin(async move { Ok(self.0.get(hash).cloned()) }) + } +} + +fn keys() -> Arc { + let status = KeyStatus::Active { + expires_at: None, + models: KeyModels::default(), + allowed_routes: AllowedRoutes::default(), + team_id: Some("team-1".into()), + profile: Arc::new(KeyProfile { + alias: Some("team-key".into()), + user_id: Some("user-1".into()), + team_alias: Some("research".into()), + ..KeyProfile::default() + }), + }; + let hash = KeyHash::from_token(&litellm_auth_types::SecretValue::new("sk-team")); + Arc::new(Keys(HashMap::from([(hash, status)]))) +} + +struct Harness { + app: Router, + inference: Arc, + s3_log: Arc, +} + +async fn harness(s3: &MockServer) -> Harness { + let upstream = upstream().await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: gpt + litellm_params: + model: openai/gpt-x + api_base: {upstream}/v1 + api_key: sk-upstream + model_info: + id: gpt-1 + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {upstream} + api_key: sk-upstream +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master +litellm_settings: + callbacks: [s3_v2] + s3_callback_params: + s3_bucket_name: logs + s3_path: audit + s3_region_name: us-east-1 + s3_endpoint_url: {endpoint} + s3_aws_access_key_id: AKIDTEST + s3_aws_secret_access_key: secret + s3_use_team_prefix: true + s3_server_side_encryption: AES256 +", + endpoint = s3.uri(), + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let s3_log = S3Log::from_config(&config, &inference) + .await + .unwrap() + .unwrap(); + litellm_gateway::hooks::install(&inference, &config, Vec::new(), vec![s3_log.clone()]) + .await + .unwrap(); + let app = + litellm_gateway::router_with_keys(inference.clone(), &config, None, None, Some(keys())); + Harness { + app, + inference, + s3_log, + } +} + +async fn send(app: &Router, path: &str, body: Value) { + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer sk-team") + .header("user-agent", "client/1.0") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + assert_eq!(status, 200, "{body:?}"); +} + +/// `YYYYMMDDTHHMMSSZ` as a time. +fn amz_date(text: &str) -> SystemTime { + let number = |range: std::ops::Range| text[range].parse::().unwrap(); + let (year, month, day) = (number(0..4), number(4..6), number(6..8)); + let year = if month <= 2 { year - 1 } else { year }; + let era = year.div_euclid(400); + let year_of_era = year - era * 400; + let day_of_year = (153 * (month + if month > 2 { -3 } else { 9 }) + 2) / 5 + day - 1; + let day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year; + let days = era * 146_097 + day_of_era - 719_468; + let seconds = days * 86_400 + number(9..11) * 3600 + number(11..13) * 60 + number(13..15); + UNIX_EPOCH + Duration::from_secs(seconds as u64) +} + +#[tokio::test] +async fn a_request_is_written_as_a_signed_put_and_a_503_is_retried() { + let s3 = MockServer::start().await; + Mock::given(method("PUT")) + .and(path_regex(OBJECT_KEY)) + .respond_with(ResponseTemplate::new(503)) + .up_to_n_times(1) + .with_priority(1) + .mount(&s3) + .await; + Mock::given(method("PUT")) + .and(path_regex(OBJECT_KEY)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&s3) + .await; + let harness = harness(&s3).await; + send( + &harness.app, + "/v1/chat/completions", + json!({"model": "gpt", "temperature": 0.5, "messages": [{"role": "user", "content": "hello"}]}), + ) + .await; + + litellm_gateway::hooks::close(&harness.inference).await; + + let requests = s3.received_requests().await.unwrap(); + assert_eq!(requests.len(), 2, "the 503 is retried once"); + assert_eq!(requests[0].url, requests[1].url); + let put = &requests[1]; + let header = |name: &str| { + put.headers + .get(name) + .map(|value| value.to_str().unwrap().to_owned()) + }; + let payload: Value = serde_json::from_slice(&put.body).unwrap(); + assert_eq!(payload["id"], "chatcmpl-1"); + assert_eq!(payload["status"], "success"); + assert_eq!(payload["model_group"], "gpt"); + assert_eq!(payload["model_id"], "gpt-1"); + assert_eq!( + payload["messages"], + json!([{"role": "user", "content": "hello"}]) + ); + assert_eq!( + payload["response"]["choices"][0]["message"]["content"], + "hi" + ); + assert_eq!(payload["metadata"]["user_api_key_team_alias"], "research"); + assert_eq!(payload["total_tokens"], 11); + + assert_eq!(header("content-type").as_deref(), Some("application/json")); + assert_eq!(header("content-language").as_deref(), Some("en")); + assert_eq!( + header("cache-control").as_deref(), + Some("private, immutable, max-age=31536000, s-maxage=0") + ); + assert_eq!( + header("x-amz-server-side-encryption").as_deref(), + Some("AES256") + ); + let disposition = header("content-disposition").unwrap(); + assert!( + disposition.starts_with("inline; filename=\"time-") + && disposition.ends_with("_chatcmpl-1.json\""), + "{disposition}" + ); + let sha256: String = Sha256::digest(&put.body) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect(); + assert_eq!(header("x-amz-content-sha256"), Some(sha256)); + let md5 = base64::engine::general_purpose::STANDARD + .encode(::digest(&put.body)); + assert_eq!(header("content-md5"), Some(md5)); + + let signed: BTreeMap = [ + "content-type", + "content-md5", + "x-amz-content-sha256", + "content-language", + "content-disposition", + "cache-control", + "x-amz-server-side-encryption", + ] + .into_iter() + .map(|name| (name.to_owned(), header(name).unwrap())) + .collect(); + let expected = sign_s3( + "PUT", + &format!("{}{}", s3.uri(), put.url.path()), + &put.body, + &signed, + "us-east-1", + &Credentials::new("AKIDTEST", "secret", None, None, "test"), + amz_date(&header("x-amz-date").unwrap()), + ) + .unwrap(); + assert_eq!(header("x-amz-date"), Some(expected["X-Amz-Date"].clone())); + assert_eq!( + header("authorization"), + Some(expected["Authorization"].clone()) + ); + assert!( + expected["Authorization"].contains("SignedHeaders=cache-control;content-disposition;content-language;content-md5;content-type;host;x-amz-content-sha256;x-amz-date;x-amz-server-side-encryption"), + "{}", + expected["Authorization"] + ); +} + +#[tokio::test] +async fn an_object_s3_rejects_is_dropped_while_the_rest_are_written() { + let s3 = MockServer::start().await; + Mock::given(method("PUT")) + .and(path_regex(r"_msg_1\.json$")) + .respond_with( + ResponseTemplate::new(400) + .set_body_string("EntityTooLarge"), + ) + .expect(1) + .mount(&s3) + .await; + Mock::given(method("PUT")) + .and(path_regex(OBJECT_KEY)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&s3) + .await; + let harness = harness(&s3).await; + send( + &harness.app, + "/v1/chat/completions", + json!({"model": "gpt", "messages": [{"role": "user", "content": "hello"}]}), + ) + .await; + send( + &harness.app, + "/v1/messages", + json!({"model": "claude", "max_tokens": 8, "messages": [{"role": "user", "content": "hello"}]}), + ) + .await; + + litellm_gateway::hooks::close(&harness.inference).await; + harness.s3_log.flush().await; + + assert_eq!(s3.received_requests().await.unwrap().len(), 2); +} diff --git a/litellm-rust/docs/customer-proxy/s3-log.svg b/litellm-rust/docs/customer-proxy/s3-log.svg new file mode 100644 index 00000000000..1618e7d8044 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/s3-log.svg @@ -0,0 +1,61 @@ + + + + + + s3_v2 observer: each request's logging payload as an S3 object, off the request path + + + At startup + on when callbacks lists s3_v2 (or s3_log); + success_callback / failure_callback + log only successes / only failures + s3_callback_params, os.environ/ resolved; + bad numbers and flags warn, use defaults + credentials: keys, profile, role, web + identity or the default chain + no s3_bucket_name stops startup + + + Request finishes (observe) + StandardLoggingPayload: key, team, + deployment, usage, cost, timings, + messages, response (streams rebuilt) + turn_off_message_logging, strip_base64, + log_prompts_only applied here + queued; past s3_max_queue_size the + oldest are dropped with a warning + nothing awaits S3 here + + + + Every 10 s, and at shutdown + the whole queue, s3_max_concurrent_uploads + PUTs at a time (16), SigV4 signed + s3_batch_file_upload: one .jsonl per + folder instead of one .json per request + failed for now: back at the front, sent + after newer ones on the next flush + DEFAULT_S3_FLUSH_INTERVAL_SECONDS + + + + Object key + {s3_path}/[team_alias/][key_alias/] + YYYY-MM-DD[/HH]/time-HH-MM-SS-ffffff_{id}.json + prefixes with s3_use_team_prefix / s3_use_key_prefix + /HH with s3_partition_granularity: hour + id: the response's id, else the request's call id + past S3's 1024 / 255 byte limits: head + sha256 + path style, or s3_use_virtual_hosted_style + + + Upload and retries + 3 attempts; transport error, 403, 429, 500, 503 + retried after 0.5-1.5 s, then 1-3 s + 400 / 403 with an object error code (EntityTooLarge, + BadDigest, ...) dropped (s3_drop_on_terminal_error) + kept until s3_max_retry_age_seconds (3600) + nothing delivered: all kept, the bucket is down + Content-MD5, SSE and Content-Disposition as the proxy +