diff --git a/litellm-rust/crates/gateway-hooks/src/hook.rs b/litellm-rust/crates/gateway-hooks/src/hook.rs index 9246f77d1a2..7d92d15769c 100644 --- a/litellm-rust/crates/gateway-hooks/src/hook.rs +++ b/litellm-rust/crates/gateway-hooks/src/hook.rs @@ -146,7 +146,7 @@ pub struct RequestEdit { allowed: Targets, writer: &'static str, header_ops: Vec, - body_fields: Vec<(String, Value)>, + body_fields: Vec<(Vec, Value)>, response_headers: Vec<(HeaderName, HeaderValue)>, annotations: Annotations, } @@ -183,8 +183,15 @@ impl RequestEdit { /// Sets a top-level field of the upstream body, replacing what the client sent. pub fn set_upstream_body_field(&mut self, field: impl Into, value: Value) { - if self.may(Targets::UPSTREAM_BODY) { - self.body_fields.push((field.into(), value)); + self.set_upstream_body_path(vec![field.into()], value); + } + + /// Sets a field inside a top-level object of the upstream body, such as + /// `["metadata", "user_id"]`, keeping the object's other fields. A missing or non-object + /// step becomes an object. + pub fn set_upstream_body_path(&mut self, path: Vec, value: Value) { + if !path.is_empty() && self.may(Targets::UPSTREAM_BODY) { + self.body_fields.push((path, value)); } } @@ -202,7 +209,11 @@ impl RequestEdit { /// Whether any change reaches the provider. pub fn changes_upstream(&self) -> bool { - !self.header_ops.is_empty() || !self.body_fields.is_empty() + self.changes_upstream_headers() || !self.body_fields.is_empty() + } + + pub fn changes_upstream_headers(&self) -> bool { + !self.header_ops.is_empty() } /// Applies the header changes in order to the headers about to be sent. @@ -221,8 +232,24 @@ impl RequestEdit { /// Applies the body changes in order to the body about to be sent. pub fn apply_upstream_body(&self, body: &mut Map) { - for (field, value) in &self.body_fields { - body.insert(field.clone(), value.clone()); + for (path, value) in &self.body_fields { + let Some((last, parents)) = path.split_last() else { + continue; + }; + let mut object = &mut *body; + for step in parents { + let slot = object + .entry(step.clone()) + .or_insert_with(|| Value::Object(Map::new())); + if !slot.is_object() { + *slot = Value::Object(Map::new()); + } + let Value::Object(inner) = slot else { + unreachable!("just made an object") + }; + object = inner; + } + object.insert(last.clone(), value.clone()); } } @@ -328,6 +355,25 @@ mod tests { ); } + #[rstest] + fn nested_body_fields_keep_their_siblings() { + let mut edit = RequestEdit::default(); + edit.scope("hook", Targets::UPSTREAM_BODY); + edit.set_upstream_body_path(vec!["metadata".into(), "user_id".into()], Value::from("u")); + edit.set_upstream_body_path(vec!["extra".into(), "a".into()], Value::from(1)); + + let serde_json::Value::Object(mut body) = + serde_json::json!({"metadata": {"tag": "t"}, "extra": "not an object"}) + else { + unreachable!() + }; + edit.apply_upstream_body(&mut body); + assert_eq!( + Value::Object(body), + serde_json::json!({"metadata": {"tag": "t", "user_id": "u"}, "extra": {"a": 1}}) + ); + } + #[rstest] fn undeclared_writes_are_dropped() { let mut edit = RequestEdit::default(); diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 118eb2cacad..e5fac1cafcc 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -93,6 +93,18 @@ async fn handle( None => route, }; + // Chat forwards no client headers of its own; hooks may set some. + let extra_headers = tracking + .as_ref() + .and_then(|tracking| tracking.edit_headers(Vec::new())) + .filter(|headers| !headers.is_empty()) + .map(|headers| { + headers + .into_iter() + .map(|(name, value)| (name, Value::from(value))) + .collect::>() + }); + let extra_headers = &extra_headers; let cache_headers = crate::caching::CacheHeaders::default(); let mut body = Some(body); let routed = gateway @@ -120,7 +132,7 @@ async fn handle( &deployment.model, deployment.custom_llm_provider.as_deref(), ) { - return openai(gateway, deployment, body).await; + return openai(gateway, deployment, body, extra_headers.clone()).await; } let (messages, optional_params) = split(body); let call = ChatCompletionsCall { @@ -130,7 +142,7 @@ async fn handle( api_key: deployment.api_key.clone(), api_base: deployment.api_base.clone(), custom_llm_provider: deployment.custom_llm_provider.clone(), - extra_headers: None, + extra_headers: extra_headers.clone(), timeout: deployment.timeout, }; Ok(litellm_host_http::serve_unary( @@ -166,6 +178,7 @@ async fn openai( gateway: &Gateway, deployment: &Deployment, body: Map, + extra_headers: Option>, ) -> Result { let machine = gateway.chat_completions.clone().openai_machine( OpenAiChatCall { @@ -174,7 +187,7 @@ async fn openai( api_key: deployment.api_key.clone(), api_base: deployment.api_base.clone(), custom_llm_provider: deployment.custom_llm_provider.clone(), - extra_headers: None, + extra_headers, timeout: deployment.timeout, }, None, diff --git a/litellm-rust/crates/gateway-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs index 07833d8cc93..4893a91e835 100644 --- a/litellm-rust/crates/gateway-inference/src/hooks.rs +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -154,6 +154,20 @@ impl Tracking { self.admission.edit().apply_upstream_body(body); } + /// The headers the gateway would forward, after the hooks' changes; `None` when no hook + /// changes upstream headers. + pub(crate) fn edit_headers( + &self, + forwarded: Vec<(String, String)>, + ) -> Option> { + let edit = self.admission.edit(); + edit.changes_upstream_headers().then(|| { + let mut forwarded = forwarded; + edit.apply_upstream_headers(&mut forwarded); + forwarded + }) + } + /// Notes the deployment an attempt went to; the last one is the one that served. pub(crate) fn attempt(&self, name: &str, attempt: &Attempt<'_>) { if !self.pipeline.reads(Facts::DEPLOYMENT) { diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index 03f1f57b905..800cabe7659 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -82,6 +82,7 @@ async fn handle( None => route, }; + let upstream_headers = &upstream_headers(headers, tracking.as_ref()); let cache_headers = crate::caching::CacheHeaders::default(); let mut body = Some(body); let routed = gateway @@ -103,7 +104,7 @@ async fn handle( async move { crate::observe::record(name, attempt.deployment); request::authorize(identity, name, attempt.deployment).await?; - let call = project(attempt.deployment, body?, headers)?; + let call = project(attempt.deployment, body?, upstream_headers)?; let machine = route.machine(call, cache_options.policy); let stream = Sse::::new(Json, |error| { Bytes::from(Error::from(error).sse_frame()) @@ -123,10 +124,61 @@ async fn handle( } } +/// Client headers the provider gets: the API headers scoped to the providers that take them, +/// and any header a hook set. +#[derive(Default)] +struct UpstreamHeaders { + scoped: Option, + extra: Option>, +} + +fn upstream_headers( + headers: &HeaderMap, + tracking: Option<&crate::hooks::Tracking>, +) -> UpstreamHeaders { + let client = |name: &str| Some(headers.get(name)?.to_str().ok()?.to_owned()); + let Some(tracking) = tracking else { + return UpstreamHeaders { + scoped: scoped_headers(client), + extra: None, + }; + }; + let mut names: Vec<&str> = header_scopes() + .iter() + .flat_map(|(_, names)| names.iter().copied()) + .collect(); + names.sort_unstable(); + names.dedup(); + let forwarded = names + .iter() + .filter_map(|name| Some(((*name).to_owned(), client(name)?))) + .collect(); + let Some(sent) = tracking.edit_headers(forwarded) else { + return UpstreamHeaders { + scoped: scoped_headers(client), + extra: None, + }; + }; + let value = |name: &str| { + sent.iter() + .find(|(sent, _)| sent.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.clone()) + }; + let extra: Map = sent + .iter() + .filter(|(name, _)| !names.iter().any(|scoped| scoped.eq_ignore_ascii_case(name))) + .map(|(name, value)| (name.clone(), Value::from(value.as_str()))) + .collect(); + UpstreamHeaders { + scoped: scoped_headers(value), + extra: (!extra.is_empty()).then_some(extra), + } +} + fn project( deployment: &Deployment, body: Map, - headers: &HeaderMap, + headers: &UpstreamHeaders, ) -> Result { let body = body .into_iter() @@ -141,28 +193,31 @@ fn project( api_base: deployment.api_base.clone(), custom_llm_provider: deployment.custom_llm_provider.clone(), litellm_params: deployment.litellm_params.clone(), - extra_headers: None, - provider_specific_header: anthropic_api_headers(headers), + extra_headers: headers.extra.clone(), + provider_specific_header: headers.scoped.clone(), timeout: deployment.timeout, shaping: deployment.shaping.clone(), }) } -fn anthropic_api_headers(headers: &HeaderMap) -> Option { - let scopes = [ +fn header_scopes() -> [(&'static str, &'static [&'static str]); 2] { + [ ( ANTHROPIC_API_HEADER_PROVIDERS, ANTHROPIC_API_HEADERS.as_slice(), ), ("openrouter", OPENROUTER_API_HEADERS.as_slice()), - ]; - let forwarded: Vec = scopes + ] +} + +fn scoped_headers(get: impl Fn(&str) -> Option) -> Option { + let forwarded: Vec = header_scopes() .into_iter() .filter_map(|(provider, names)| { let extra_headers: Map = names .iter() .filter_map(|name| { - let value = headers.get(*name)?.to_str().ok()?; + let value = get(name)?; Some(((*name).to_owned(), Value::from(value))) }) .collect(); diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index a2097417ce5..fbab39ef426 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -20,6 +20,8 @@ pub enum Error { DatabaseUrl(String), #[error("reading a gateway setting from the secret source failed")] Secret(#[from] litellm_secrets::Error), + #[error("litellm_settings.hooks: {0}")] + Hooks(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 57f05ec9322..c97c1d6515b 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -10,6 +10,7 @@ mod key_status; mod keys; pub mod metrics; pub mod rate_limits; +pub mod request_hooks; 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 3e648f5a758..b8a35eddbcd 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -128,6 +128,7 @@ async fn main() -> Result<(), Box> { ), )); } + hooks.extend(litellm_gateway::request_hooks::from_config(&config)?); let mut observers: Vec> = Vec::new(); let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config)); let spend_tracking = spend.is_some(); diff --git a/litellm-rust/crates/gateway/src/request_hooks.rs b/litellm-rust/crates/gateway/src/request_hooks.rs new file mode 100644 index 00000000000..348a6852504 --- /dev/null +++ b/litellm-rust/crates/gateway/src/request_hooks.rs @@ -0,0 +1,556 @@ +//! Built-in hooks a config turns on and configures under `litellm_settings.hooks`, run in the +//! order listed: +//! +//! ```yaml +//! litellm_settings: +//! hooks: +//! - forward_headers: +//! headers: ["x-request-origin", "x-trace-*"] +//! - spend_metadata: +//! from_headers: +//! x-litellm-spend-logs-metadata: json +//! x-cost-center: cost_center +//! from_key_metadata: [cost_center] +//! - strip_headers: +//! headers: ["x-litellm-spend-logs-metadata", "x-cost-center"] +//! - provider_metadata: +//! body: +//! metadata.user_id: "{end_user|key.user_id}" +//! headers: +//! x-caller-team: "{key.team_alias}" +//! ``` + +use std::sync::Arc; + +use futures_util::future::BoxFuture; +use http::{HeaderName, HeaderValue}; +use litellm_config::Config; +use litellm_gateway_hooks::{ + Facts, HeaderPattern, Hook, Refusal, RequestEdit, RequestFacts, Route, Settle, Targets, +}; +use serde::Deserialize; +use serde_json::{Map, Value}; + +use crate::{Error, spend::SPEND_LOGS_METADATA}; + +/// Headers a hook never forwards: credentials, and those the HTTP client or a request signer +/// sets itself. +const PROTECTED: [&str; 13] = [ + "authorization", + "proxy-authorization", + "x-api-key", + "api-key", + "x-litellm-api-key", + "x-goog-api-key", + "cookie", + "host", + "content-length", + "content-type", + "transfer-encoding", + "connection", + "date", +]; + +fn protected(name: &str) -> bool { + PROTECTED + .iter() + .any(|protected| protected.eq_ignore_ascii_case(name)) + || name + .get(..6) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case("x-amz-")) +} + +/// The hooks `litellm_settings.hooks` lists, in its order. +pub fn from_config(config: &Config) -> Result>, Error> { + let Some(entries) = config.litellm_settings.additional_fields.get("hooks") else { + return Ok(Vec::new()); + }; + let litellm_config::Value::Sequence(entries) = entries else { + return Err(Error::Hooks("must be a list".into())); + }; + entries.iter().map(hook).collect() +} + +fn hook(entry: &litellm_config::Value) -> Result, Error> { + let (name, settings) = match entry { + litellm_config::Value::Mapping(mapping) if mapping.len() == 1 => { + let (name, settings) = mapping.iter().next().expect("one entry"); + (name.as_str().unwrap_or_default(), settings.clone()) + } + litellm_config::Value::String(name) => (name.as_str(), litellm_config::Value::Null), + _ => { + return Err(Error::Hooks( + "each entry is a hook name with its settings".into(), + )); + } + }; + let built: Result, String> = match name { + "forward_headers" => parse(settings) + .and_then(patterns) + .map(|patterns| Arc::new(ForwardHeaders { patterns }) as Arc), + "strip_headers" => parse(settings) + .and_then(patterns) + .map(|patterns| Arc::new(StripHeaders { patterns }) as Arc), + "spend_metadata" => parse(settings) + .and_then(SpendMetadata::new) + .map(|hook| Arc::new(hook) as Arc), + "provider_metadata" => parse(settings) + .and_then(ProviderMetadata::new) + .map(|hook| Arc::new(hook) as Arc), + other => return Err(Error::Hooks(format!("unknown hook {other}"))), + }; + built.map_err(|error| Error::Hooks(format!("{name}: {error}"))) +} + +fn parse(settings: litellm_config::Value) -> Result { + T::deserialize(settings).map_err(|error| error.to_string()) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct HeaderList { + headers: Vec, +} + +fn patterns(list: HeaderList) -> Result, String> { + list.headers + .iter() + .map(|pattern| { + HeaderPattern::parse(pattern).map_err(|_| format!("{pattern} is not a header name")) + }) + .collect() +} + +/// `forward_headers`: copies the client headers it names onto the provider request. +struct ForwardHeaders { + patterns: Vec, +} + +impl Hook for ForwardHeaders { + fn name(&self) -> &'static str { + "forward_headers" + } + + fn reads(&self) -> Facts { + Facts::CLIENT_HEADERS + } + + fn writes(&self) -> Targets { + Targets::UPSTREAM_HEADERS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + for (name, value) in request.client_headers.iter() { + if !protected(name.as_str()) + && self + .patterns + .iter() + .any(|pattern| pattern.matches(name.as_str())) + { + edit.set_upstream_header(name.clone(), value.clone()); + } + } + Box::pin(async { Ok(None) }) + } +} + +/// `strip_headers`: removes the headers it names from the provider request, whether a hook set +/// them or the gateway forwards them by default. +struct StripHeaders { + patterns: Vec, +} + +impl Hook for StripHeaders { + fn name(&self) -> &'static str { + "strip_headers" + } + + fn reads(&self) -> Facts { + Facts::NONE + } + + fn writes(&self) -> Targets { + Targets::UPSTREAM_HEADERS + } + + fn run<'a>( + &'a self, + _request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + for pattern in &self.patterns { + edit.remove_upstream_headers(pattern.clone()); + } + Box::pin(async { Ok(None) }) + } +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct SpendMetadataSettings { + /// Header name to `json`, merged as an object, or to the key it is stored under. + #[serde(default)] + from_headers: litellm_config::Value, + /// Fields of the key's metadata copied under the same names. + #[serde(default)] + from_key_metadata: Vec, +} + +enum HeaderRule { + Json, + Field(String), +} + +/// `spend_metadata`: adds caller-supplied values to the request's `spend_logs_metadata` +/// without changing the request. +struct SpendMetadata { + headers: Vec<(HeaderName, HeaderRule)>, + key_fields: Vec, +} + +impl SpendMetadata { + fn new(settings: SpendMetadataSettings) -> Result { + let headers = match settings.from_headers { + litellm_config::Value::Null => Vec::new(), + litellm_config::Value::Mapping(mapping) => mapping + .iter() + .map(|(name, rule)| { + let (Some(name), Some(rule)) = (name.as_str(), rule.as_str()) else { + return Err("from_headers maps header names to json or a key".to_owned()); + }; + let name = HeaderName::from_bytes(name.as_bytes()) + .map_err(|_| format!("{name} is not a header name"))?; + let rule = match rule { + "json" => HeaderRule::Json, + field => HeaderRule::Field(field.to_owned()), + }; + Ok((name, rule)) + }) + .collect::>()?, + _ => return Err("from_headers maps header names to json or a key".into()), + }; + Ok(Self { + headers, + key_fields: settings.from_key_metadata, + }) + } +} + +impl Hook for SpendMetadata { + fn name(&self) -> &'static str { + "spend_metadata" + } + + fn reads(&self) -> Facts { + Facts::CLIENT_HEADERS | Facts::KEY + } + + fn writes(&self) -> Targets { + Targets::ANNOTATIONS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + let mut metadata = Map::new(); + for (name, rule) in &self.headers { + let Some(value) = request + .client_headers + .get(name) + .and_then(|value| value.to_str().ok()) + else { + continue; + }; + match rule { + HeaderRule::Json => match serde_json::from_str::(value) { + Ok(Value::Object(fields)) => metadata.extend(fields), + _ => tracing::warn!( + header = name.as_str(), + "spend metadata header is not a JSON object; ignored" + ), + }, + HeaderRule::Field(field) => { + metadata.insert(field.clone(), Value::from(value)); + } + } + } + // The key's own values win over what the caller sent. + if let Some(key) = request.key.as_deref() { + for field in &self.key_fields { + if let Some(value) = key.metadata.get(field) { + metadata.insert(field.clone(), value.clone()); + } + } + } + if !metadata.is_empty() { + edit.annotate(SPEND_LOGS_METADATA, Value::Object(metadata)); + } + Box::pin(async { Ok(None) }) + } +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct ProviderMetadataSettings { + /// Dotted body paths, such as `metadata.user_id`, to templates. + #[serde(default)] + body: std::collections::BTreeMap, + #[serde(default)] + headers: std::collections::BTreeMap, +} + +/// A fact a template can name. +#[derive(Clone, Debug, PartialEq)] +enum Fact { + EndUser, + Model, + RequestId, + KeyAlias, + KeyUserId, + KeyTeamId, + KeyTeamAlias, + KeyMetadata(String), + TeamMetadata(String), + Header(HeaderName), +} + +impl Fact { + fn parse(name: &str) -> Result { + Ok(match name { + "end_user" => Self::EndUser, + "model" => Self::Model, + "request_id" => Self::RequestId, + "key.alias" => Self::KeyAlias, + "key.user_id" => Self::KeyUserId, + "key.team_id" => Self::KeyTeamId, + "key.team_alias" => Self::KeyTeamAlias, + other => { + if let Some(field) = other.strip_prefix("key.metadata.") { + Self::KeyMetadata(field.to_owned()) + } else if let Some(field) = other.strip_prefix("team.metadata.") { + Self::TeamMetadata(field.to_owned()) + } else if let Some(header) = other.strip_prefix("header.") { + Self::Header( + HeaderName::from_bytes(header.as_bytes()) + .map_err(|_| format!("{header} is not a header name"))?, + ) + } else { + return Err(format!("unknown fact {other}")); + } + } + }) + } + + fn value(&self, request: &RequestFacts) -> Option { + let key = request.key.as_deref(); + let scalar = |value: &Value| match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + Value::Bool(flag) => Some(flag.to_string()), + _ => None, + }; + match self { + Self::EndUser => request.end_user.as_deref().map(str::to_owned), + Self::Model => Some(request.requested_model.to_string()), + Self::RequestId => Some(request.request_id.to_string()), + Self::KeyAlias => key?.alias.clone(), + Self::KeyUserId => key?.user_id.clone(), + Self::KeyTeamId => key?.team_id.clone(), + Self::KeyTeamAlias => key?.team_alias.clone(), + Self::KeyMetadata(field) => key?.metadata.get(field).and_then(scalar), + Self::TeamMetadata(field) => key?.team_metadata.get(field).and_then(scalar), + Self::Header(name) => request + .client_headers + .get(name) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned), + } + .filter(|value| !value.is_empty()) + } +} + +#[derive(Debug, PartialEq)] +enum Part { + Text(String), + /// The first of these facts that is set. + Facts(Vec), +} + +/// Text with `{fact|fallback}` placeholders, parsed once. A placeholder none of whose facts are +/// set leaves the whole value unset. +#[derive(Debug, PartialEq)] +struct Template(Vec); + +impl Template { + fn parse(source: &str) -> Result { + let mut parts = Vec::new(); + let mut rest = source; + while let Some(open) = rest.find('{') { + if open > 0 { + parts.push(Part::Text(rest[..open].to_owned())); + } + let close = rest[open..] + .find('}') + .ok_or_else(|| format!("unclosed placeholder in {source}"))?; + let facts = rest[open + 1..open + close] + .split('|') + .map(|fact| Fact::parse(fact.trim())) + .collect::>()?; + parts.push(Part::Facts(facts)); + rest = &rest[open + close + 1..]; + } + if !rest.is_empty() { + parts.push(Part::Text(rest.to_owned())); + } + Ok(Self(parts)) + } + + fn reads(&self) -> Facts { + let mut reads = Facts::NONE; + for part in &self.0 { + let Part::Facts(facts) = part else { continue }; + for fact in facts { + reads |= match fact { + Fact::EndUser => Facts::END_USER, + Fact::Header(_) => Facts::CLIENT_HEADERS, + Fact::Model | Fact::RequestId => Facts::NONE, + _ => Facts::KEY, + }; + } + } + reads + } + + fn render(&self, request: &RequestFacts) -> Option { + let mut rendered = String::new(); + for part in &self.0 { + match part { + Part::Text(text) => rendered.push_str(text), + Part::Facts(facts) => { + rendered.push_str(&facts.iter().find_map(|fact| fact.value(request))?) + } + } + } + Some(rendered) + } +} + +/// `provider_metadata`: sets body fields and headers on the provider request from facts about +/// the caller. Body fields are set on `/v1/messages` requests only. +struct ProviderMetadata { + body: Vec<(Vec, Template)>, + headers: Vec<(HeaderName, Template)>, + reads: Facts, +} + +impl ProviderMetadata { + fn new(settings: ProviderMetadataSettings) -> Result { + let body = settings + .body + .iter() + .map(|(path, template)| { + let path: Vec = path.split('.').map(str::to_owned).collect(); + if path.iter().any(String::is_empty) { + return Err(format!("{} is not a body path", path.join("."))); + } + Ok((path, Template::parse(template)?)) + }) + .collect::, String>>()?; + let headers = settings + .headers + .iter() + .map(|(name, template)| { + let header = HeaderName::from_bytes(name.as_bytes()) + .map_err(|_| format!("{name} is not a header name"))?; + if protected(header.as_str()) { + return Err(format!("{name} cannot be set by a hook")); + } + Ok((header, Template::parse(template)?)) + }) + .collect::, String>>()?; + let reads = body + .iter() + .map(|(_, template)| template) + .chain(headers.iter().map(|(_, template)| template)) + .fold(Facts::NONE, |reads, template| reads | template.reads()); + Ok(Self { + body, + headers, + reads, + }) + } +} + +impl Hook for ProviderMetadata { + fn name(&self) -> &'static str { + "provider_metadata" + } + + fn reads(&self) -> Facts { + self.reads + } + + fn writes(&self) -> Targets { + Targets::UPSTREAM_BODY | Targets::UPSTREAM_HEADERS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + // Body paths name Messages API fields; Chat Completions takes the headers only. + let body = match request.route { + Route::Messages => self.body.as_slice(), + Route::ChatCompletions => &[], + }; + for (path, template) in body { + if let Some(value) = template.render(request) { + edit.set_upstream_body_path(path.clone(), Value::from(value)); + } + } + for (name, template) in &self.headers { + if let Some(value) = template + .render(request) + .and_then(|value| HeaderValue::from_str(&value).ok()) + { + edit.set_upstream_header(name.clone(), value); + } + } + Box::pin(async { Ok(None) }) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + fn templates_name_facts_with_fallbacks() { + assert_eq!( + Template::parse("team-{key.team_alias|key.team_id}").unwrap(), + Template(vec![ + Part::Text("team-".into()), + Part::Facts(vec![Fact::KeyTeamAlias, Fact::KeyTeamId]), + ]) + ); + assert!(Template::parse("{nope}").is_err()); + assert!(Template::parse("{end_user").is_err()); + } + + #[rstest] + #[case::credential("Authorization", true)] + #[case::signed("X-Amz-Date", true)] + #[case::trace("x-trace-id", false)] + fn credentials_are_never_forwarded(#[case] name: &str, #[case] expected: bool) { + assert_eq!(protected(name), expected); + } +} diff --git a/litellm-rust/crates/gateway/tests/request_hooks.rs b/litellm-rust/crates/gateway/tests/request_hooks.rs new file mode 100644 index 00000000000..af8a74b3079 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/request_hooks.rs @@ -0,0 +1,339 @@ +//! The configurable request hooks on real requests: headers forwarded and stripped, provider +//! metadata set from the caller, and spend metadata handed to observers. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; + +use axum::{ + Router, + body::{Body, Bytes, to_bytes}, + http::{HeaderMap, Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; + +use futures_util::future::BoxFuture; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture, +}; +use litellm_gateway_hooks::{ + Facts, KeyFacts, Observer, Observers, Pipeline, Record, RequestFacts, Route, +}; +use rstest::rstest; +use serde_json::{Map, Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#; + +type Seen = Arc>>; + +async fn upstream() -> (String, Seen) { + let seen = Seen::default(); + let recorded = seen.clone(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(move |request: Request| { + let recorded = recorded.clone(); + async move { + let headers = request.headers().clone(); + let body = to_bytes(request.into_body(), 1 << 20).await.unwrap(); + recorded.lock().unwrap().push(( + headers, + serde_json::from_slice(&body).unwrap_or(Value::Null), + )); + ([(CONTENT_TYPE, "application/json")], MESSAGE).into_response() + } + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + (format!("http://{address}"), seen) +} + +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()) }) + } +} + +const HOOKS: &str = r#" + hooks: + - forward_headers: + headers: ["x-trace-*", "x-cost-center", "authorization"] + - strip_headers: + headers: ["x-cost-center", "anthropic-beta"] + - provider_metadata: + body: + metadata.user_id: "{end_user|key.user_id}" + headers: + x-caller-team: "team-{key.team_alias}" +"#; + +async fn gateway(hooks: &str) -> (Router, Seen) { + let (address, seen) = upstream().await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {address} + api_key: sk-upstream +general_settings: + master_key: sk-master +litellm_settings:{hooks}" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let hooks = litellm_gateway::request_hooks::from_config(&config).unwrap(); + litellm_gateway::hooks::install(&inference, &config, hooks, Vec::new()) + .await + .unwrap(); + let key = KeyStatus::Active { + expires_at: None, + models: KeyModels::default(), + allowed_routes: AllowedRoutes::default(), + team_id: Some("team-1".into()), + profile: Arc::new(KeyProfile { + user_id: Some("user-1".into()), + team_alias: Some("research".into()), + ..KeyProfile::default() + }), + }; + let keys: Arc = Arc::new(Keys(HashMap::from([( + KeyHash::from_token(&SecretValue::new("sk-key")), + key, + )]))); + let app = litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys)); + (app, seen) +} + +async fn send(app: &Router, path: &str, body: Value) -> u16 { + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer sk-key") + .header("x-trace-id", "trace-1") + .header("x-cost-center", "cc-9") + .header("anthropic-beta", "some-beta") + .header("x-other", "kept-out") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + response.status().as_u16() +} + +fn header<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> { + headers.get(name).map(|value| value.to_str().unwrap()) +} + +#[rstest] +#[tokio::test] +async fn messages_carry_the_configured_headers_and_metadata() { + let (app, seen) = gateway(HOOKS).await; + let body = json!({ + "model": "claude", + "max_tokens": 8, + "metadata": {"tag": "t"}, + "messages": [{"role": "user", "content": "hi"}], + }); + assert_eq!(send(&app, "/v1/messages", body).await, 200); + let (headers, body) = seen.lock().unwrap().pop().unwrap(); + assert_eq!(header(&headers, "x-trace-id"), Some("trace-1")); + assert_eq!(header(&headers, "x-caller-team"), Some("team-research")); + // Stripped after forwarding, and the gateway's own pass-through is stripped too. + assert_eq!(header(&headers, "x-cost-center"), None); + assert_eq!(header(&headers, "anthropic-beta"), None); + assert_eq!(header(&headers, "x-other"), None); + // Credentials are never forwarded, even when listed. + assert_ne!(header(&headers, "authorization"), Some("Bearer sk-key")); + assert_eq!(header(&headers, "x-api-key"), Some("sk-upstream")); + // The provider keeps only the `metadata` fields it knows. + assert_eq!(body["metadata"]["user_id"], "user-1"); +} + +#[rstest] +#[tokio::test] +async fn the_end_user_comes_before_the_key_user() { + let (app, seen) = gateway(HOOKS).await; + let body = json!({ + "model": "claude", + "max_tokens": 8, + "litellm_metadata": {"user": "end-7"}, + "messages": [{"role": "user", "content": "hi"}], + }); + assert_eq!(send(&app, "/v1/messages", body).await, 200); + let (_, body) = seen.lock().unwrap().pop().unwrap(); + assert_eq!(body["metadata"]["user_id"], "end-7"); +} + +#[rstest] +#[tokio::test] +async fn chat_completions_take_the_headers_only() { + let (app, seen) = gateway(HOOKS).await; + let body = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + assert_eq!(send(&app, "/v1/chat/completions", body).await, 200); + let (headers, body) = seen.lock().unwrap().pop().unwrap(); + assert_eq!(header(&headers, "x-trace-id"), Some("trace-1")); + assert_eq!(header(&headers, "x-caller-team"), Some("team-research")); + assert_eq!(header(&headers, "x-cost-center"), None); + assert_eq!(body.get("metadata"), None); +} + +#[rstest] +#[tokio::test] +async fn without_hooks_the_default_pass_through_is_unchanged() { + let (app, seen) = gateway("\n drop_params: true\n").await; + let body = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + assert_eq!(send(&app, "/v1/messages", body).await, 200); + let (headers, _) = seen.lock().unwrap().pop().unwrap(); + assert_eq!(header(&headers, "anthropic-beta"), Some("some-beta")); + assert_eq!(header(&headers, "x-trace-id"), None); +} + +struct Annotations(Mutex>>); + +impl Observer for Annotations { + fn name(&self) -> &'static str { + "annotations" + } + + fn reads(&self) -> Facts { + Facts::ANNOTATIONS + } + + fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> { + Box::pin(async move { + self.0.lock().unwrap().push( + record + .annotations + .get(litellm_gateway::spend::SPEND_LOGS_METADATA) + .cloned(), + ); + }) + } +} + +#[rstest] +#[tokio::test] +async fn spend_metadata_reaches_observers() { + let config = Config::from_yaml( + r#" +model_list: [] +litellm_settings: + hooks: + - spend_metadata: + from_headers: + x-litellm-spend-logs-metadata: json + x-cost-center: cost_center + from_key_metadata: [owner] +"#, + ) + .unwrap(); + let observed = Arc::new(Annotations(Mutex::default())); + let pipeline = Pipeline::new( + litellm_gateway::request_hooks::from_config(&config).unwrap(), + Observers::start(vec![observed.clone()], 8), + Arc::default(), + ); + let run = async |headers: &[(&'static str, &'static str)]| { + let mut map = HeaderMap::new(); + for (name, value) in headers { + map.insert(*name, value.parse().unwrap()); + } + let Value::Object(metadata) = json!({"owner": "platform"}) else { + unreachable!() + }; + let key = KeyFacts { + metadata: Arc::new(metadata), + ..KeyFacts::default() + }; + let request = RequestFacts::new( + "r".into(), + Route::Messages, + Some(Arc::new(key)), + None, + "claude".into(), + Arc::new(map), + false, + None, + Bytes::from_static(b"{}"), + ); + let admission = pipeline.admit(Arc::new(request)).await.unwrap(); + pipeline.finish( + admission, + litellm_gateway_hooks::Outcome { + deployment: None, + status: 200, + usage: None, + cost: None, + time_to_first_byte: None, + ended_at: std::time::SystemTime::now(), + interrupted: false, + }, + None, + ); + }; + run(&[ + ( + "x-litellm-spend-logs-metadata", + r#"{"project": "p1", "owner": "caller"}"#, + ), + ("x-cost-center", "cc-9"), + ]) + .await; + run(&[("x-litellm-spend-logs-metadata", "not json")]).await; + pipeline.close().await; + let mut expected = Map::new(); + expected.insert("project".into(), json!("p1")); + expected.insert("owner".into(), json!("platform")); + expected.insert("cost_center".into(), json!("cc-9")); + assert_eq!( + *observed.0.lock().unwrap(), + [ + Some(Value::Object(expected)), + Some(json!({"owner": "platform"})) + ] + ); +} + +#[rstest] +#[case::unknown_hook(" - rewrite_everything: {}\n", "unknown hook rewrite_everything")] +#[case::bad_template( + " - provider_metadata:\n headers:\n x-a: \"{key.secret}\"\n", + "provider_metadata: unknown fact key.secret" +)] +#[case::credential_header( + " - provider_metadata:\n headers:\n authorization: \"{key.user_id}\"\n", + "provider_metadata: authorization cannot be set by a hook" +)] +#[case::unknown_setting( + " - strip_headers:\n names: [x-a]\n", + "strip_headers: unknown field `names`" +)] +fn misconfigured_hooks_stop_the_gateway(#[case] hooks: &str, #[case] message: &str) { + let config = Config::from_yaml(&format!( + "model_list: []\nlitellm_settings:\n hooks:\n{hooks}" + )) + .unwrap(); + let error = litellm_gateway::request_hooks::from_config(&config) + .err() + .unwrap() + .to_string(); + assert!(error.contains(message), "{error}"); +} diff --git a/litellm-rust/docs/customer-proxy/request-hooks.svg b/litellm-rust/docs/customer-proxy/request-hooks.svg new file mode 100644 index 00000000000..74e96af77e9 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/request-hooks.svg @@ -0,0 +1,56 @@ + + + + + + Configurable request hooks: litellm_settings.hooks, run top to bottom + + + Client request + x-trace-id: trace-1 + x-cost-center: cc-9 + x-litellm-spend-logs- + metadata: {...} + anthropic-beta: ... + authorization: Bearer sk-... + key: user, team alias, + metadata; end user + + + + forward_headers copies listed client headers (exact or prefix-*) + + spend_metadata headers (json or one key) + key metadata + + strip_headers removes forwarded and pass-through headers + + provider_metadata "{end_user|key.user_id}" templates + + + + + + Provider request + x-trace-id: trace-1 + x-caller-team: research + metadata.user_id: user-1 + (messages only) + no x-cost-center, no + anthropic-beta, never + client credentials + + + + Spend log + the spend_metadata annotation lands in + LiteLLM_SpendLogs.metadata.spend_logs_metadata + + + + Rules + Each hook declares what it reads and writes; a later hook in the list wins on the same header or field. + Header edits reach the provider on both endpoints and every provider, Bedrock's signed requests included. + Credentials, signer and framing headers (authorization, x-api-key, cookie, x-amz-*, ...) are never forwarded or set. + An unknown hook, setting or template fact stops the gateway at startup; a bad JSON header is dropped with a warning. + Templates name facts only: end_user, model, request_id, key.alias|user_id|team_id|team_alias, key.metadata.*, team.metadata.*, header.* +