diff --git a/litellm-rust/crates/gateway-hooks/src/facts.rs b/litellm-rust/crates/gateway-hooks/src/facts.rs index 2b8d8dc9e11..b0ed1c58fe2 100644 --- a/litellm-rust/crates/gateway-hooks/src/facts.rs +++ b/litellm-rust/crates/gateway-hooks/src/facts.rs @@ -231,6 +231,13 @@ impl Annotations { pub fn is_empty(&self) -> bool { self.0.is_empty() } + + /// Sets each of `later`'s values in turn, as if they had been written after this one's. + pub(crate) fn extend(&mut self, later: Annotations) { + for (key, value) in later.0 { + self.set(key, value); + } + } } /// How a request ended. diff --git a/litellm-rust/crates/gateway-hooks/src/hook.rs b/litellm-rust/crates/gateway-hooks/src/hook.rs index 7d92d15769c..15457dfb7d2 100644 --- a/litellm-rust/crates/gateway-hooks/src/hook.rs +++ b/litellm-rust/crates/gateway-hooks/src/hook.rs @@ -152,9 +152,21 @@ pub struct RequestEdit { } impl RequestEdit { - pub(crate) fn scope(&mut self, writer: &'static str, allowed: Targets) { - self.writer = writer; - self.allowed = allowed; + /// An empty edit for one hook, which may only make the changes `allowed` names. + pub(crate) fn scoped(writer: &'static str, allowed: Targets) -> Self { + Self { + allowed, + writer, + ..Self::default() + } + } + + /// Appends `later`'s changes after this edit's, as if its hook had run after this one's. + pub(crate) fn extend(&mut self, later: RequestEdit) { + self.header_ops.extend(later.header_ops); + self.body_fields.extend(later.body_fields); + self.response_headers.extend(later.response_headers); + self.annotations.extend(later.annotations); } fn may(&self, target: Targets) -> bool { @@ -314,18 +326,18 @@ mod tests { #[rstest] fn header_changes_apply_in_order() { - let mut edit = RequestEdit::default(); - edit.scope("first", Targets::UPSTREAM_HEADERS); + let mut edit = RequestEdit::scoped("first", Targets::UPSTREAM_HEADERS); edit.set_upstream_header( HeaderName::from_static("x-a"), HeaderValue::from_static("1"), ); - edit.scope("second", Targets::UPSTREAM_HEADERS); - edit.remove_upstream_headers(HeaderPattern::parse("x-*").unwrap()); - edit.set_upstream_header( + let mut second = RequestEdit::scoped("second", Targets::UPSTREAM_HEADERS); + second.remove_upstream_headers(HeaderPattern::parse("x-*").unwrap()); + second.set_upstream_header( HeaderName::from_static("x-b"), HeaderValue::from_static("2"), ); + edit.extend(second); let mut sent = headers(&[("X-A", "client"), ("content-type", "application/json")]); edit.apply_upstream_headers(&mut sent); @@ -337,11 +349,11 @@ mod tests { #[rstest] fn later_hooks_win_on_the_same_body_field() { - let mut edit = RequestEdit::default(); - edit.scope("first", Targets::UPSTREAM_BODY); + let mut edit = RequestEdit::scoped("first", Targets::UPSTREAM_BODY); edit.set_upstream_body_field("metadata", Value::from("a")); - edit.scope("second", Targets::UPSTREAM_BODY); - edit.set_upstream_body_field("metadata", Value::from("b")); + let mut second = RequestEdit::scoped("second", Targets::UPSTREAM_BODY); + second.set_upstream_body_field("metadata", Value::from("b")); + edit.extend(second); let serde_json::Value::Object(mut body) = serde_json::json!({"model": "m", "metadata": "client"}) @@ -357,8 +369,7 @@ mod tests { #[rstest] fn nested_body_fields_keep_their_siblings() { - let mut edit = RequestEdit::default(); - edit.scope("hook", Targets::UPSTREAM_BODY); + let mut edit = RequestEdit::scoped("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)); @@ -376,8 +387,7 @@ mod tests { #[rstest] fn undeclared_writes_are_dropped() { - let mut edit = RequestEdit::default(); - edit.scope("reader", Targets::ANNOTATIONS); + let mut edit = RequestEdit::scoped("reader", Targets::ANNOTATIONS); edit.set_upstream_body_field("metadata", Value::Null); edit.add_response_header( HeaderName::from_static("x-a"), diff --git a/litellm-rust/crates/gateway-hooks/src/pipeline.rs b/litellm-rust/crates/gateway-hooks/src/pipeline.rs index 64afeeb0838..addbf7332f8 100644 --- a/litellm-rust/crates/gateway-hooks/src/pipeline.rs +++ b/litellm-rust/crates/gateway-hooks/src/pipeline.rs @@ -3,6 +3,7 @@ use std::{sync::Arc, time::SystemTime}; use bytes::Bytes; +use futures_util::future::join_all; use crate::{ facts::{Facts, Outcome, Record, RequestFacts}, @@ -19,7 +20,7 @@ pub struct Pipeline { } impl Pipeline { - /// `hooks` run in the order given, which is the config's. + /// `hooks` are given in the config's order, which decides whose changes win. pub fn new(hooks: Vec>, observers: Observers, prices: Arc) -> Self { let mut reads = observers.reads(); for hook in &hooks { @@ -60,8 +61,9 @@ impl Pipeline { &self.prices } - /// Runs every hook in order. When one refuses, the hooks before it are told the request - /// ended with that refusal, and so are the observers. + /// Runs every hook at once. Their changes apply in config order, and when several refuse, + /// the first in config order answers. On a refusal the hooks that let the request through + /// are told it ended with that refusal, and so are the observers. pub async fn admit(&self, request: Arc) -> Result { let admitted = self.run_hooks(request).await; // Refused or not, the request's facts outlive this call (in observer records, or for @@ -74,38 +76,57 @@ impl Pipeline { admitted.map_err(|(_, refusal)| refusal) } + /// Hooks read only the request and write only their own edit, so none waits on another; + /// admission costs the slowest hook's round trips rather than the sum of them. async fn run_hooks( &self, request: Arc, ) -> Result, Refusal)> { + let mut edits: Vec = self + .hooks + .iter() + .map(|hook| RequestEdit::scoped(hook.name(), hook.writes())) + .collect(); + let results = join_all( + self.hooks + .iter() + .zip(edits.iter_mut()) + .map(|(hook, edit)| hook.run(&request, edit)), + ) + .await; let mut admission = Admission { request, edit: RequestEdit::default(), settles: Vec::new(), }; - for hook in &self.hooks { - admission.edit.scope(hook.name(), hook.writes()); - match hook.run(&admission.request, &mut admission.edit).await { - Ok(Some(settle)) => admission.settles.push(settle), - Ok(None) => {} + let mut refused = None; + for (result, edit) in results.into_iter().zip(edits) { + match result { + Ok(settle) => { + admission.settles.extend(settle); + admission.edit.extend(edit); + } Err(refusal) => { - let outcome = Outcome { - deployment: None, - status: refusal.status, - usage: None, - cost: None, - time_to_first_byte: None, - upstream_started_at: None, - ended_at: SystemTime::now(), - interrupted: false, - }; - let request = Arc::clone(&admission.request); - self.finish(admission, outcome, None); - return Err((request, refusal)); + refused.get_or_insert(refusal); } } } - Ok(admission) + let Some(refusal) = refused else { + return Ok(admission); + }; + let outcome = Outcome { + deployment: None, + status: refusal.status, + usage: None, + cost: None, + time_to_first_byte: None, + upstream_started_at: None, + ended_at: SystemTime::now(), + interrupted: false, + }; + let request = Arc::clone(&admission.request); + self.finish(admission, outcome, None); + Err((request, refusal)) } /// Records how the request ended: prices its usage, lets hooks settle, and hands the @@ -288,7 +309,7 @@ mod tests { } #[tokio::test] - async fn hooks_run_in_order_and_settle_with_the_cost() { + async fn later_hooks_win_and_settle_with_the_cost() { let log = Log::default(); let costs = Arc::new(Costs(Mutex::default())); let mut prices = Prices::default(); @@ -370,7 +391,7 @@ mod tests { } #[tokio::test] - async fn a_refusal_settles_the_hooks_before_it() { + async fn a_refusal_settles_the_hooks_that_let_the_request_through() { let log = Log::default(); let pipeline = Pipeline::new( vec![ @@ -383,6 +404,94 @@ mod tests { ); let refusal = pipeline.admit(request()).await.err().unwrap(); assert_eq!(refusal.status, 429); - assert_eq!(*log.lock().unwrap(), ["run a", "run limit", "settle a 429"]); + assert_eq!( + *log.lock().unwrap(), + [ + "run a", + "run limit", + "run c", + "settle a 429", + "settle c 429" + ] + ); + } + + struct Refuses(u16); + + impl Hook for Refuses { + fn name(&self) -> &'static str { + "refuses" + } + + fn reads(&self) -> Facts { + Facts::KEY + } + + fn writes(&self) -> Targets { + Targets::ADMISSION + } + + fn run<'a>( + &'a self, + _request: &'a Arc, + _edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { Err(Refusal::new(self.0, "refused", "no")) }) + } + } + + #[tokio::test] + async fn the_first_refusal_in_config_order_answers() { + let pipeline = Pipeline::new( + vec![Arc::new(Refuses(400)), Arc::new(Refuses(429))], + Observers::start(Vec::new(), 8), + Arc::default(), + ); + assert_eq!(pipeline.admit(request()).await.err().unwrap().status, 400); + } + + /// Each hook waits until every hook has started, which only finishes when they run at once. + struct Meets(Arc); + + impl Hook for Meets { + fn name(&self) -> &'static str { + "meets" + } + + fn reads(&self) -> Facts { + Facts::KEY + } + + fn writes(&self) -> Targets { + Targets::ADMISSION + } + + fn run<'a>( + &'a self, + _request: &'a Arc, + _edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { + self.0.wait().await; + Ok(None) + }) + } + } + + #[tokio::test] + async fn hooks_run_at_once() { + let barrier = Arc::new(tokio::sync::Barrier::new(3)); + let pipeline = Pipeline::new( + (0..3) + .map(|_| Arc::new(Meets(barrier.clone())) as Arc) + .collect(), + Observers::start(Vec::new(), 8), + Arc::default(), + ); + let admitted = + tokio::time::timeout(std::time::Duration::from_secs(5), pipeline.admit(request())) + .await + .expect("hooks waited on each other"); + assert!(admitted.is_ok()); } }