From a3a58c55db6925db437934576dad029a4e92330f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 13 Oct 2026 13:54:59 -0700 Subject: [PATCH] refactor(rust): compose async hooks in the shared host --- litellm-rust/crates/host-python/AGENTS.md | 2 +- .../crates/host-python/src/hooks/chain.rs | 167 ++------ litellm-rust/crates/host/AGENTS.md | 2 +- litellm-rust/crates/host/src/hooks.rs | 6 + litellm-rust/crates/host/src/hooks/adapter.rs | 91 +++++ litellm-rust/crates/host/src/hooks/chain.rs | 366 ++++++++++++++++++ 6 files changed, 501 insertions(+), 133 deletions(-) create mode 100644 litellm-rust/crates/host/src/hooks/adapter.rs create mode 100644 litellm-rust/crates/host/src/hooks/chain.rs diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index 04e8dab997c..093871421f2 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -14,7 +14,7 @@ The driver owns ordering: start, argument preparation, prepared-argument hooks, Boundary tests exercise behavior with a supplied lifecycle binding without importing the LiteLLM Python package. Pin inline awaiting, awaitable final values, exception identity, cancellation, re-entry and release of retained objects, rather than module names or source layout -`HookChain` composes Python runtime hooks in order. Each argument, wire-request and response transformation feeds its result to the next hook. After all argument transformations, the driver calls `arguments_prepared` on every hook in order. Retained callback views must adopt that dictionary before later policy hooks can mutate or reject it. SDK policy is supplied by bridge composition as a hook, never a separate driver phase or parameter. Hooks implement only the stages they need. Default stages preserve the supplied values +`HookChain` specializes the shared host chain for Python runtime hooks. The Python adapter owns cancellation guards, notification error policy, close and GC traversal. Each argument, wire-request and response transformation feeds its result to the next hook. After all argument transformations, the driver calls `arguments_prepared` on every hook in order. Retained callback views must adopt that dictionary before later policy hooks can mutate or reject it. SDK policy is supplied by bridge composition as a hook, never a separate driver phase or parameter. Hooks implement only the stages they need. Default stages preserve the supplied values Terminal notifications share the selected response or exception. An ordinary notification error is reported as unraisable and does not skip the next hook or replace the selected outcome. Preparation, interception and transformation errors stop the chain. Cancellation stops all further hook dispatch. Suspensions stay inline in the existing driver, and the chain traverses retained event values for GC diff --git a/litellm-rust/crates/host-python/src/hooks/chain.rs b/litellm-rust/crates/host-python/src/hooks/chain.rs index a29be6b40f6..5a46c817eb3 100644 --- a/litellm-rust/crates/host-python/src/hooks/chain.rs +++ b/litellm-rust/crates/host-python/src/hooks/chain.rs @@ -4,7 +4,7 @@ use std::task::Poll; use futures_util::future::BoxFuture; use litellm_host::{ - hooks::CallHooks, + hooks::{CallHooks, ChainRuntime, DynCallHooks, HookGroupAdapter}, interceptors::{RequestContext, WireRequest}, lifecycle::{CallEvent, Timing}, }; @@ -18,85 +18,29 @@ use pyo3::{ use crate::{Awaiter, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime}; -trait ChainHooks: PythonOwned { - fn prepare_arguments<'a>( - &'a self, - task: &'a dyn Awaiter, - arguments: Py, - started_at: f64, - ) -> BoxFuture<'a, PyResult>>; - fn arguments_prepared(&self, arguments: &Py) -> PyResult<()>; - fn before_provider_request<'a>( - &'a self, - task: &'a dyn Awaiter, - wire: Box, - context: &'a RequestContext, - ) -> BoxFuture<'a, PyResult>>; - fn transform_response<'a>( - &'a self, - task: &'a dyn Awaiter, - response: Py, - timing: Timing, - ) -> BoxFuture<'a, PyResult>>; - fn on_event<'a>( - &'a self, - task: &'a dyn Awaiter, - event: PythonCallEvent<'a>, - ) -> BoxFuture<'a, PyResult<()>>; - fn on_stream_open(&self, head: &Py) -> PyResult<()>; - fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()>; +pub type HookChain = litellm_host::hooks::HookChain; + +pub trait PythonHookGroup: DynCallHooks + PythonOwned {} +impl + PythonOwned> PythonHookGroup for H {} + +impl ChainRuntime for PythonRuntime { + type Group = dyn PythonHookGroup; } -impl ChainHooks for H { - fn prepare_arguments<'a>( - &'a self, - task: &'a dyn Awaiter, - arguments: Py, - started_at: f64, - ) -> BoxFuture<'a, PyResult>> { - Box::pin(CallHooks::prepare_arguments( - self, task, arguments, started_at, - )) +impl HookGroupAdapter for PythonRuntime { + fn adapt_group(hooks: H) -> Box { + Box::new(PythonHookAdapter(hooks)) } +} - fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { - CallHooks::arguments_prepared(self, arguments) +struct PythonHookAdapter(H); + +impl PythonOwned for PythonHookAdapter { + fn close(&mut self, py: Python<'_>) { + self.0.close(py); } - - fn before_provider_request<'a>( - &'a self, - task: &'a dyn Awaiter, - wire: Box, - context: &'a RequestContext, - ) -> BoxFuture<'a, PyResult>> { - Box::pin(CallHooks::before_provider_request( - self, task, wire, context, - )) - } - - fn transform_response<'a>( - &'a self, - task: &'a dyn Awaiter, - response: Py, - timing: Timing, - ) -> BoxFuture<'a, PyResult>> { - Box::pin(CallHooks::transform_response(self, task, response, timing)) - } - - fn on_event<'a>( - &'a self, - task: &'a dyn Awaiter, - event: PythonCallEvent<'a>, - ) -> BoxFuture<'a, PyResult<()>> { - Box::pin(CallHooks::on_event(self, task, event)) - } - - fn on_stream_open(&self, head: &Py) -> PyResult<()> { - CallHooks::on_stream_open(self, head) - } - - fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()> { - CallHooks::on_stream_chunk(self, chunk) + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.0.traverse(visit) } } @@ -166,23 +110,7 @@ fn notification_result(terminal: bool, result: PyResult<()>) -> PyResult<()> { } } -#[derive(Default)] -pub struct HookChain { - hooks: Vec>, -} - -impl HookChain { - pub fn new() -> Self { - Self::default() - } - - pub fn with(mut self, hooks: impl PythonCallHooks + 'static) -> Self { - self.hooks.push(Box::new(hooks)); - self - } -} - -impl CallHooks for HookChain { +impl CallHooks for PythonHookAdapter { async fn prepare_arguments<'a>( &'a self, task: &'a dyn Awaiter, @@ -190,19 +118,13 @@ impl CallHooks for HookChain { started_at: f64, ) -> PyResult> { let stage = Stage::new(task); - let mut arguments = arguments; - for hooks in &self.hooks { - arguments = stage - .guard(hooks.prepare_arguments(&stage, arguments, started_at)) - .await?; - } - Ok(arguments) + stage + .guard(self.0.prepare_arguments(&stage, arguments, started_at)) + .await } fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { - self.hooks - .iter() - .try_for_each(|hooks| hooks.arguments_prepared(arguments)) + self.0.arguments_prepared(arguments) } async fn before_provider_request<'a>( @@ -212,13 +134,9 @@ impl CallHooks for HookChain { context: &RequestContext, ) -> PyResult> { let stage = Stage::new(task); - let mut wire = wire; - for hooks in &self.hooks { - wire = stage - .guard(hooks.before_provider_request(&stage, wire, context)) - .await?; - } - Ok(wire) + stage + .guard(self.0.before_provider_request(&stage, wire, context)) + .await } async fn transform_response<'a>( @@ -228,13 +146,9 @@ impl CallHooks for HookChain { timing: Timing, ) -> PyResult> { let stage = Stage::new(task); - let mut response = response; - for hooks in &self.hooks { - response = stage - .guard(hooks.transform_response(&stage, response, timing)) - .await?; - } - Ok(response) + stage + .guard(self.0.transform_response(&stage, response, timing)) + .await } async fn on_event<'a>( @@ -244,36 +158,27 @@ impl CallHooks for HookChain { ) -> PyResult<()> { let stage = Stage::new(task); let terminal = is_terminal(&event); - for hooks in &self.hooks { - let result = stage.guard(hooks.on_event(&stage, event.clone())).await; - notification_result(terminal, result)?; - } - Ok(()) + let result = stage.guard(self.0.on_event(&stage, event)).await; + notification_result(terminal, result) } fn on_stream_open(&self, head: &Py) -> PyResult<()> { - self.hooks - .iter() - .try_for_each(|hooks| hooks.on_stream_open(head)) + self.0.on_stream_open(head) } fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()> { - self.hooks - .iter() - .try_for_each(|hooks| hooks.on_stream_chunk(chunk)) + self.0.on_stream_chunk(chunk) } } impl PythonOwned for HookChain { fn close(&mut self, py: Python<'_>) { - for hooks in &mut self.hooks { + for hooks in self.groups_mut() { hooks.close(py); } } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.hooks - .iter() - .try_for_each(|hooks| hooks.traverse(visit)) + self.groups().try_for_each(|hooks| hooks.traverse(visit)) } } diff --git a/litellm-rust/crates/host/AGENTS.md b/litellm-rust/crates/host/AGENTS.md index 0f06f0e57e1..f3442e9864c 100644 --- a/litellm-rust/crates/host/AGENTS.md +++ b/litellm-rust/crates/host/AGENTS.md @@ -16,7 +16,7 @@ Each driver owns terminal dispatch. Hooks can change or fail execution. Passive `interceptors.rs` owns `Interceptors` and its request/response payload types. `lifecycle.rs` owns `CallObserver`, `CallEvent`, `ExecutionEvent`, timing, failure origin, and the observation wrappers. Event payloads are generic so a runtime can retain its own response, exception and raw-response references without introducing a language dependency. `snapshot()` projects them into the owned observation contract without retaining runtime objects. Pass interceptors and observers separately at direct route and HTTP entrypoints. Routes publish execution events independently of interception -`hooks.rs` owns the call-stage interface and its runtime-associated types. It contains no Python types or legacy callback policy. A runtime supplies its context and continuation representation through `HookRuntime` +`hooks.rs` owns the async call-stage interface and its runtime-associated types. `hooks::HookChain` composes stages in order, using `DynCallHooks` for boxed storage and `CallHooks` for callback authors. A runtime supplies its context through `HookRuntime` and its stored hook type through `ChainRuntime`. Python ownership, cancellation and notification error policy stay in the Python adapter `protocol.rs` owns `Protocol` and suspension messages, including `InterceptRequest` and `StreamDelivery`. `call.rs` owns route outputs and their adaptation into a hosted machine. Rust service handling belongs in `host-native::services`. Coroutine channel handles stay in `machine/context.rs` diff --git a/litellm-rust/crates/host/src/hooks.rs b/litellm-rust/crates/host/src/hooks.rs index 7f78c059ed2..bfb1ef1cf83 100644 --- a/litellm-rust/crates/host/src/hooks.rs +++ b/litellm-rust/crates/host/src/hooks.rs @@ -1,3 +1,9 @@ +mod adapter; +mod chain; + +pub use adapter::DynCallHooks; +pub use chain::{ChainRuntime, HookChain, HookGroupAdapter}; + use std::future::{Future, ready}; use crate::{ diff --git a/litellm-rust/crates/host/src/hooks/adapter.rs b/litellm-rust/crates/host/src/hooks/adapter.rs new file mode 100644 index 00000000000..3c0b0e295f0 --- /dev/null +++ b/litellm-rust/crates/host/src/hooks/adapter.rs @@ -0,0 +1,91 @@ +use futures_util::future::BoxFuture; + +use super::{CallHooks, HookRuntime, RuntimeCallEvent}; +use crate::{ + interceptors::{RequestContext, WireRequest}, + lifecycle::Timing, +}; + +pub trait DynCallHooks: Sync { + fn prepare_arguments<'a>( + &'a self, + runtime: R::Context<'a>, + arguments: R::Arguments, + started_at: f64, + ) -> BoxFuture<'a, Result>; + fn arguments_prepared(&self, arguments: &R::Arguments) -> Result<(), R::Error>; + fn before_provider_request<'a: 'b, 'b>( + &'a self, + runtime: R::Context<'a>, + wire: Box, + context: &'b RequestContext, + ) -> BoxFuture<'b, Result, R::Error>>; + fn transform_response<'a>( + &'a self, + runtime: R::Context<'a>, + response: R::Response, + timing: Timing, + ) -> BoxFuture<'a, Result>; + fn on_event<'a: 'b, 'b>( + &'a self, + runtime: R::Context<'a>, + event: RuntimeCallEvent<'b, R>, + ) -> BoxFuture<'b, Result<(), R::Error>>; + fn on_stream_open(&self, head: &R::Response) -> Result<(), R::Error>; + fn on_stream_chunk(&self, chunk: &R::Chunk) -> Result<(), R::Error>; +} + +impl> DynCallHooks for H { + fn prepare_arguments<'a>( + &'a self, + runtime: R::Context<'a>, + arguments: R::Arguments, + started_at: f64, + ) -> BoxFuture<'a, Result> { + Box::pin(CallHooks::prepare_arguments( + self, runtime, arguments, started_at, + )) + } + + fn arguments_prepared(&self, arguments: &R::Arguments) -> Result<(), R::Error> { + CallHooks::arguments_prepared(self, arguments) + } + + fn before_provider_request<'a: 'b, 'b>( + &'a self, + runtime: R::Context<'a>, + wire: Box, + context: &'b RequestContext, + ) -> BoxFuture<'b, Result, R::Error>> { + Box::pin(CallHooks::before_provider_request( + self, runtime, wire, context, + )) + } + + fn transform_response<'a>( + &'a self, + runtime: R::Context<'a>, + response: R::Response, + timing: Timing, + ) -> BoxFuture<'a, Result> { + Box::pin(CallHooks::transform_response( + self, runtime, response, timing, + )) + } + + fn on_event<'a: 'b, 'b>( + &'a self, + runtime: R::Context<'a>, + event: RuntimeCallEvent<'b, R>, + ) -> BoxFuture<'b, Result<(), R::Error>> { + Box::pin(CallHooks::on_event(self, runtime, event)) + } + + fn on_stream_open(&self, head: &R::Response) -> Result<(), R::Error> { + CallHooks::on_stream_open(self, head) + } + + fn on_stream_chunk(&self, chunk: &R::Chunk) -> Result<(), R::Error> { + CallHooks::on_stream_chunk(self, chunk) + } +} diff --git a/litellm-rust/crates/host/src/hooks/chain.rs b/litellm-rust/crates/host/src/hooks/chain.rs new file mode 100644 index 00000000000..d270888e43d --- /dev/null +++ b/litellm-rust/crates/host/src/hooks/chain.rs @@ -0,0 +1,366 @@ +use super::{CallHooks, DynCallHooks, HookRuntime, RuntimeCallEvent}; +use crate::{ + interceptors::{RequestContext, WireRequest}, + lifecycle::Timing, +}; + +pub trait ChainRuntime: HookRuntime + Sized + 'static { + type Group: DynCallHooks + ?Sized; +} + +pub trait HookGroupAdapter: ChainRuntime { + fn adapt_group(hooks: H) -> Box; +} + +pub struct HookChain { + hooks: Vec>, +} + +impl HookChain { + pub fn groups(&self) -> impl Iterator { + self.hooks.iter().map(Box::as_ref) + } + + pub fn groups_mut(&mut self) -> impl Iterator { + self.hooks.iter_mut().map(Box::as_mut) + } + + pub fn new() -> Self { + Self::default() + } + + pub fn with(mut self, hooks: H) -> Self + where + R: HookGroupAdapter, + { + self.hooks.push(R::adapt_group(hooks)); + self + } +} + +impl CallHooks for HookChain +where + for<'a> R::Context<'a>: Copy, +{ + async fn prepare_arguments<'a>( + &'a self, + runtime: R::Context<'a>, + arguments: R::Arguments, + started_at: f64, + ) -> Result { + let mut arguments = arguments; + for hooks in &self.hooks { + arguments = hooks + .prepare_arguments(runtime, arguments, started_at) + .await?; + } + Ok(arguments) + } + + fn arguments_prepared(&self, arguments: &R::Arguments) -> Result<(), R::Error> { + self.hooks + .iter() + .try_for_each(|hooks| hooks.arguments_prepared(arguments)) + } + + async fn before_provider_request<'a>( + &'a self, + runtime: R::Context<'a>, + wire: Box, + context: &RequestContext, + ) -> Result, R::Error> { + let mut wire = wire; + for hooks in &self.hooks { + wire = hooks + .before_provider_request(runtime, wire, context) + .await?; + } + Ok(wire) + } + + async fn transform_response<'a>( + &'a self, + runtime: R::Context<'a>, + response: R::Response, + timing: Timing, + ) -> Result { + let mut response = response; + for hooks in &self.hooks { + response = hooks.transform_response(runtime, response, timing).await?; + } + Ok(response) + } + + async fn on_event<'a>( + &'a self, + runtime: R::Context<'a>, + event: RuntimeCallEvent<'_, R>, + ) -> Result<(), R::Error> { + for hooks in &self.hooks { + hooks.on_event(runtime, event.clone()).await?; + } + Ok(()) + } + + fn on_stream_open(&self, head: &R::Response) -> Result<(), R::Error> { + self.hooks + .iter() + .try_for_each(|hooks| hooks.on_stream_open(head)) + } + + fn on_stream_chunk(&self, chunk: &R::Chunk) -> Result<(), R::Error> { + self.hooks + .iter() + .try_for_each(|hooks| hooks.on_stream_chunk(chunk)) + } +} + +impl Default for HookChain { + fn default() -> Self { + Self { hooks: Vec::new() } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use rstest::{fixture, rstest}; + + use super::*; + use crate::lifecycle::CallEvent; + + struct TestRuntime; + + impl HookRuntime for TestRuntime { + type Context<'a> = &'a Mutex>; + type Arguments = String; + type Response = String; + type Chunk = u32; + type Error = &'static str; + type Failure = &'static str; + } + + impl ChainRuntime for TestRuntime { + type Group = dyn DynCallHooks; + } + + impl + 'static> HookGroupAdapter for TestRuntime { + fn adapt_group(hooks: H) -> Box { + Box::new(hooks) + } + } + + struct Group { + name: &'static str, + suspensions: u8, + reject: bool, + } + + impl Group { + async fn value( + &self, + trace: &Mutex>, + stage: &str, + value: String, + ) -> Result { + trace.lock().unwrap().push(format!("{stage}:{}", self.name)); + for _ in 0..self.suspensions { + tokio::task::yield_now().await; + } + if self.reject { + return Err("rejected"); + } + Ok(format!("{value}{}", self.name)) + } + } + + impl CallHooks for Group { + async fn prepare_arguments<'a>( + &'a self, + trace: &'a Mutex>, + arguments: String, + _: f64, + ) -> Result { + self.value(trace, "prepare", arguments).await + } + + async fn before_provider_request<'a>( + &'a self, + trace: &'a Mutex>, + wire: Box, + _: &RequestContext, + ) -> Result, &'static str> { + let url = self.value(trace, "wire", wire.url.clone()).await?; + Ok(Box::new(WireRequest { url, ..*wire })) + } + + async fn transform_response<'a>( + &'a self, + trace: &'a Mutex>, + response: String, + _: Timing, + ) -> Result { + self.value(trace, "response", response).await + } + + async fn on_event<'a>( + &'a self, + trace: &'a Mutex>, + event: RuntimeCallEvent<'_, TestRuntime>, + ) -> Result<(), &'static str> { + let CallEvent::Succeeded { response, .. } = event else { + panic!("expected success") + }; + self.value(trace, response, String::new()).await.map(|_| ()) + } + } + + #[fixture] + fn trace() -> Mutex> { + Mutex::new(Vec::new()) + } + + #[fixture] + fn timing() -> Timing { + Timing { + start_time: 0.0, + end_time: 1.0, + } + } + + #[fixture] + fn context() -> RequestContext { + RequestContext { + model: String::new(), + custom_llm_provider: String::new(), + optional_params: serde_json::Value::Null, + secret_fields: Vec::new(), + api_key: None, + } + } + + #[rstest] + #[tokio::test] + async fn groups_compose_without_python( + trace: Mutex>, + timing: Timing, + context: RequestContext, + #[values(0, 1, 2)] suspensions: u8, + ) { + let hooks = HookChain::::new() + .with(Group { + name: "a", + suspensions, + reject: false, + }) + .with(Group { + name: "b", + suspensions, + reject: false, + }); + let arguments = CallHooks::prepare_arguments(&hooks, &trace, "input".into(), 0.0) + .await + .unwrap(); + assert_eq!(arguments, "inputab"); + let wire = Box::new(WireRequest { + url: "url".into(), + headers: Vec::new(), + body: serde_json::Value::Null, + }); + let rewritten = CallHooks::before_provider_request(&hooks, &trace, wire, &context) + .await + .unwrap(); + assert_eq!(rewritten.url, "urlab"); + let response = CallHooks::transform_response(&hooks, &trace, "output".into(), timing) + .await + .unwrap(); + assert_eq!(response, "outputab"); + CallHooks::on_event( + &hooks, + &trace, + CallEvent::Succeeded { + timing, + response: &response, + }, + ) + .await + .unwrap(); + assert_eq!( + *trace.lock().unwrap(), + [ + "prepare:a", + "prepare:b", + "wire:a", + "wire:b", + "response:a", + "response:b", + "outputab:a", + "outputab:b", + ] + ); + } + + #[rstest] + #[tokio::test] + async fn failure_stops_later_groups( + trace: Mutex>, + timing: Timing, + #[values(0, 2)] suspensions: u8, + #[values(false, true)] terminal: bool, + ) { + let hooks = HookChain::::new() + .with(Group { + name: "a", + suspensions, + reject: false, + }) + .with(Group { + name: "b", + suspensions, + reject: true, + }) + .with(Group { + name: "c", + suspensions, + reject: false, + }); + if terminal { + let response = "selected".to_owned(); + let result = CallHooks::on_event( + &hooks, + &trace, + CallEvent::Succeeded { + timing, + response: &response, + }, + ) + .await; + assert_eq!(result, Err("rejected")); + assert_eq!(*trace.lock().unwrap(), ["selected:a", "selected:b"]); + return; + } + let result = CallHooks::prepare_arguments(&hooks, &trace, String::new(), 0.0).await; + assert_eq!(result, Err("rejected")); + assert_eq!(*trace.lock().unwrap(), ["prepare:a", "prepare:b"]); + } + + #[rstest] + #[tokio::test] + async fn empty_chain_preserves_values(trace: Mutex>, timing: Timing) { + let hooks = HookChain::::new(); + assert_eq!( + CallHooks::prepare_arguments(&hooks, &trace, "input".into(), 0.0) + .await + .unwrap(), + "input" + ); + assert_eq!( + CallHooks::transform_response(&hooks, &trace, "output".into(), timing) + .await + .unwrap(), + "output" + ); + assert!(trace.lock().unwrap().is_empty()); + } +}