refactor(rust): compose async hooks in the shared host

This commit is contained in:
Yujong Lee 2026-10-13 13:54:59 -07:00
parent ab864fe02c
commit a3a58c55db
6 changed files with 501 additions and 133 deletions

View file

@ -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

View file

@ -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<PyDict>,
started_at: f64,
) -> BoxFuture<'a, PyResult<Py<PyDict>>>;
fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()>;
fn before_provider_request<'a>(
&'a self,
task: &'a dyn Awaiter,
wire: Box<WireRequest>,
context: &'a RequestContext,
) -> BoxFuture<'a, PyResult<Box<WireRequest>>>;
fn transform_response<'a>(
&'a self,
task: &'a dyn Awaiter,
response: Py<PyAny>,
timing: Timing,
) -> BoxFuture<'a, PyResult<Py<PyAny>>>;
fn on_event<'a>(
&'a self,
task: &'a dyn Awaiter,
event: PythonCallEvent<'a>,
) -> BoxFuture<'a, PyResult<()>>;
fn on_stream_open(&self, head: &Py<PyAny>) -> PyResult<()>;
fn on_stream_chunk(&self, chunk: &Py<PyAny>) -> PyResult<()>;
pub type HookChain = litellm_host::hooks::HookChain<PythonRuntime>;
pub trait PythonHookGroup: DynCallHooks<PythonRuntime> + PythonOwned {}
impl<H: DynCallHooks<PythonRuntime> + PythonOwned> PythonHookGroup for H {}
impl ChainRuntime for PythonRuntime {
type Group = dyn PythonHookGroup;
}
impl<H: PythonCallHooks> ChainHooks for H {
fn prepare_arguments<'a>(
&'a self,
task: &'a dyn Awaiter,
arguments: Py<PyDict>,
started_at: f64,
) -> BoxFuture<'a, PyResult<Py<PyDict>>> {
Box::pin(CallHooks::prepare_arguments(
self, task, arguments, started_at,
))
impl<H: PythonCallHooks + 'static> HookGroupAdapter<H> for PythonRuntime {
fn adapt_group(hooks: H) -> Box<dyn PythonHookGroup> {
Box::new(PythonHookAdapter(hooks))
}
}
fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
CallHooks::arguments_prepared(self, arguments)
struct PythonHookAdapter<H>(H);
impl<H: PythonOwned> PythonOwned for PythonHookAdapter<H> {
fn close(&mut self, py: Python<'_>) {
self.0.close(py);
}
fn before_provider_request<'a>(
&'a self,
task: &'a dyn Awaiter,
wire: Box<WireRequest>,
context: &'a RequestContext,
) -> BoxFuture<'a, PyResult<Box<WireRequest>>> {
Box::pin(CallHooks::before_provider_request(
self, task, wire, context,
))
}
fn transform_response<'a>(
&'a self,
task: &'a dyn Awaiter,
response: Py<PyAny>,
timing: Timing,
) -> BoxFuture<'a, PyResult<Py<PyAny>>> {
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<PyAny>) -> PyResult<()> {
CallHooks::on_stream_open(self, head)
}
fn on_stream_chunk(&self, chunk: &Py<PyAny>) -> 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<Box<dyn ChainHooks>>,
}
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<PythonRuntime> for HookChain {
impl<H: PythonCallHooks> CallHooks<PythonRuntime> for PythonHookAdapter<H> {
async fn prepare_arguments<'a>(
&'a self,
task: &'a dyn Awaiter,
@ -190,19 +118,13 @@ impl CallHooks<PythonRuntime> for HookChain {
started_at: f64,
) -> PyResult<Py<PyDict>> {
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<PyDict>) -> 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<PythonRuntime> for HookChain {
context: &RequestContext,
) -> PyResult<Box<WireRequest>> {
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<PythonRuntime> for HookChain {
timing: Timing,
) -> PyResult<Py<PyAny>> {
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<PythonRuntime> 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<PyAny>) -> 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<PyAny>) -> 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))
}
}

View file

@ -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`

View file

@ -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::{

View file

@ -0,0 +1,91 @@
use futures_util::future::BoxFuture;
use super::{CallHooks, HookRuntime, RuntimeCallEvent};
use crate::{
interceptors::{RequestContext, WireRequest},
lifecycle::Timing,
};
pub trait DynCallHooks<R: HookRuntime + 'static>: Sync {
fn prepare_arguments<'a>(
&'a self,
runtime: R::Context<'a>,
arguments: R::Arguments,
started_at: f64,
) -> BoxFuture<'a, Result<R::Arguments, R::Error>>;
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<WireRequest>,
context: &'b RequestContext,
) -> BoxFuture<'b, Result<Box<WireRequest>, R::Error>>;
fn transform_response<'a>(
&'a self,
runtime: R::Context<'a>,
response: R::Response,
timing: Timing,
) -> BoxFuture<'a, Result<R::Response, R::Error>>;
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<R: HookRuntime + 'static, H: CallHooks<R>> DynCallHooks<R> for H {
fn prepare_arguments<'a>(
&'a self,
runtime: R::Context<'a>,
arguments: R::Arguments,
started_at: f64,
) -> BoxFuture<'a, Result<R::Arguments, R::Error>> {
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<WireRequest>,
context: &'b RequestContext,
) -> BoxFuture<'b, Result<Box<WireRequest>, 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<R::Response, R::Error>> {
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)
}
}

View file

@ -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<Self> + ?Sized;
}
pub trait HookGroupAdapter<H>: ChainRuntime {
fn adapt_group(hooks: H) -> Box<Self::Group>;
}
pub struct HookChain<R: ChainRuntime> {
hooks: Vec<Box<R::Group>>,
}
impl<R: ChainRuntime> HookChain<R> {
pub fn groups(&self) -> impl Iterator<Item = &R::Group> {
self.hooks.iter().map(Box::as_ref)
}
pub fn groups_mut(&mut self) -> impl Iterator<Item = &mut R::Group> {
self.hooks.iter_mut().map(Box::as_mut)
}
pub fn new() -> Self {
Self::default()
}
pub fn with<H>(mut self, hooks: H) -> Self
where
R: HookGroupAdapter<H>,
{
self.hooks.push(R::adapt_group(hooks));
self
}
}
impl<R: ChainRuntime> CallHooks<R> for HookChain<R>
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<R::Arguments, R::Error> {
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<WireRequest>,
context: &RequestContext,
) -> Result<Box<WireRequest>, 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<R::Response, R::Error> {
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<R: ChainRuntime> Default for HookChain<R> {
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<Vec<String>>;
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<Self>;
}
impl<H: CallHooks<TestRuntime> + 'static> HookGroupAdapter<H> for TestRuntime {
fn adapt_group(hooks: H) -> Box<Self::Group> {
Box::new(hooks)
}
}
struct Group {
name: &'static str,
suspensions: u8,
reject: bool,
}
impl Group {
async fn value(
&self,
trace: &Mutex<Vec<String>>,
stage: &str,
value: String,
) -> Result<String, &'static str> {
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<TestRuntime> for Group {
async fn prepare_arguments<'a>(
&'a self,
trace: &'a Mutex<Vec<String>>,
arguments: String,
_: f64,
) -> Result<String, &'static str> {
self.value(trace, "prepare", arguments).await
}
async fn before_provider_request<'a>(
&'a self,
trace: &'a Mutex<Vec<String>>,
wire: Box<WireRequest>,
_: &RequestContext,
) -> Result<Box<WireRequest>, &'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<Vec<String>>,
response: String,
_: Timing,
) -> Result<String, &'static str> {
self.value(trace, "response", response).await
}
async fn on_event<'a>(
&'a self,
trace: &'a Mutex<Vec<String>>,
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<Vec<String>> {
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<Vec<String>>,
timing: Timing,
context: RequestContext,
#[values(0, 1, 2)] suspensions: u8,
) {
let hooks = HookChain::<TestRuntime>::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<Vec<String>>,
timing: Timing,
#[values(0, 2)] suspensions: u8,
#[values(false, true)] terminal: bool,
) {
let hooks = HookChain::<TestRuntime>::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<Vec<String>>, timing: Timing) {
let hooks = HookChain::<TestRuntime>::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());
}
}