refactor(rust): make Python hooks async fns

Co-Authored-By: Yujong Lee <yujong@berri.ai>

litellm-rust/Cargo.lock

@@ -3631,9 +3631,11 @@ dependencies = [
 name = "litellm-callbacks-legacy-python"
 version = "0.1.0"
 dependencies = [
+ "futures-util",
  "litellm-auth",
  "litellm-host",
  "litellm-host-python",
+ "parking_lot",
  "proptest",
  "pyo3",
  "rstest",

litellm-rust/crates/callbacks-legacy-python/Cargo.toml

@@ -9,6 +9,9 @@ repository.workspace = true
 litellm-host.workspace = true
 litellm-host-python.workspace = true
 
+futures-util.workspace = true
+parking_lot.workspace = true
+
 pyo3.workspace = true
 strum.workspace = true
 serde_json.workspace = true

litellm-rust/crates/callbacks-legacy-python/src/lib.rs

@@ -14,7 +14,7 @@ mod deferred;
 mod logger;
 mod mapping;
 mod python;
-pub use adapter::LegacyLogging;
+pub use adapter::{LegacyHooks, LegacyLogging};
 pub use call::PublicCall;
 pub(crate) use callbacks::{LegacyCallbacks, is_internal_call};
 pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup};

litellm-rust/crates/callbacks-legacy-python/src/mapping.rs

@@ -1,12 +1,15 @@
+use std::future::{Future, ready};
+
+use futures_util::future::BoxFuture;
 use litellm_host::{
     hooks::CallHooks,
     interceptors::{RawResponse, RequestContext, WireRequest},
     lifecycle::{ExecutionEvent, FailureOrigin, Timing},
 };
-use litellm_host_python::{HookStep, PythonCallEvent, PythonRuntime};
-use pyo3::{prelude::*, types::PyDict};
+use litellm_host_python::{Awaiter, PythonCallEvent, PythonRuntime};
+use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
 
-use crate::LegacyLogging;
+use crate::LegacyHooks;
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum CallBoundary {
@@ -48,44 +51,62 @@ impl<H> Binding<H> {
     }
 }
 
-type Step<T> = PyResult<HookStep<LegacyLogging, T>>;
-type Prepare = fn(&mut LegacyLogging, Python<'_>, Py<PyDict>, f64) -> Step<Py<PyDict>>;
-type Before =
-    fn(&mut LegacyLogging, Python<'_>, Box<WireRequest>, &RequestContext) -> Step<Box<WireRequest>>;
-type After = fn(&mut LegacyLogging, Python<'_>, &RawResponse) -> Step<()>;
-type Transform = fn(&mut LegacyLogging, Python<'_>, Py<PyAny>, Timing) -> Step<Py<PyAny>>;
-type Success = fn(&mut LegacyLogging, Python<'_>, Timing, &Py<PyAny>) -> Step<()>;
-type Failure = fn(&mut LegacyLogging, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>;
-type Open = fn(&mut LegacyLogging, Python<'_>, &Py<PyAny>) -> PyResult<()>;
-type Chunk = fn(&mut LegacyLogging, Python<'_>, &Py<PyAny>) -> PyResult<()>;
+type Prepare = for<'a> fn(
+    &'a LegacyHooks,
+    &'a dyn Awaiter,
+    Py<PyDict>,
+    f64,
+) -> BoxFuture<'a, PyResult<Py<PyDict>>>;
+type Before = fn(&LegacyHooks, Box<WireRequest>, &RequestContext) -> PyResult<Box<WireRequest>>;
+type After = fn(&LegacyHooks, &RawResponse) -> PyResult<()>;
+type Transform = for<'a> fn(
+    &'a LegacyHooks,
+    &'a dyn Awaiter,
+    Py<PyAny>,
+    Timing,
+) -> BoxFuture<'a, PyResult<Py<PyAny>>>;
+type Success = fn(&LegacyHooks, Timing, &Py<PyAny>) -> PyResult<()>;
+type Failure = for<'a> fn(
+    &'a LegacyHooks,
+    &'a dyn Awaiter,
+    Timing,
+    FailureOrigin,
+    &'a Py<PyBaseException>,
+) -> BoxFuture<'a, PyResult<()>>;
+type Open = fn(&LegacyHooks, &Py<PyAny>) -> PyResult<()>;
+type Chunk = fn(&LegacyHooks, &Py<PyAny>) -> PyResult<()>;
 
 const PREPARE: Binding<Prepare> = Binding {
     boundary: CallBoundary::PrepareArguments,
-    invoke: LegacyLogging::prepare_call,
+    invoke: |hooks, task, arguments, started_at| {
+        Box::pin(hooks.prepare_call(task, arguments, started_at))
+    },
     callbacks: &["async_pre_call_deployment_hook"],
 };
 
 const BEFORE: Binding<Before> = Binding {
     boundary: CallBoundary::BeforeProviderRequest,
-    invoke: LegacyLogging::pre_call,
+    invoke: LegacyHooks::pre_call,
     callbacks: &["log_pre_api_call", "log_input_event"],
 };
 
 const AFTER: Binding<After> = Binding {
     boundary: CallBoundary::AfterProviderResponse,
-    invoke: LegacyLogging::post_call,
+    invoke: LegacyHooks::post_call,
     callbacks: &["log_post_api_call"],
 };
 
 const TRANSFORM: Binding<Transform> = Binding {
     boundary: CallBoundary::TransformResponse,
-    invoke: LegacyLogging::transform_public_response,
+    invoke: |hooks, task, response, timing| {
+        Box::pin(hooks.transform_public_response(task, response, timing))
+    },
     callbacks: &["async_post_call_success_deployment_hook"],
 };
 
 const SUCCESS: Binding<Success> = Binding {
     boundary: CallBoundary::Succeeded,
-    invoke: LegacyLogging::succeeded,
+    invoke: LegacyHooks::succeeded,
     callbacks: &[
         "log_success_event",
         "async_log_success_event",
@@ -99,7 +120,9 @@ const SUCCESS: Binding<Success> = Binding {
 
 const FAILURE: Binding<Failure> = Binding {
     boundary: CallBoundary::Failed,
-    invoke: LegacyLogging::failed,
+    invoke: |hooks, task, timing, origin, error| {
+        Box::pin(hooks.failed(task, timing, origin, error))
+    },
     callbacks: &[
         "async_post_call_failure_deployment_hook",
         "log_failure_event",
@@ -112,13 +135,13 @@ const FAILURE: Binding<Failure> = Binding {
 
 const OPEN: Binding<Open> = Binding {
     boundary: CallBoundary::StreamOpened,
-    invoke: LegacyLogging::stream_opened,
+    invoke: LegacyHooks::stream_opened,
     callbacks: &[],
 };
 
 const CHUNK: Binding<Chunk> = Binding {
     boundary: CallBoundary::StreamChunk,
-    invoke: LegacyLogging::stream_chunk,
+    invoke: LegacyHooks::stream_chunk,
     callbacks: &[],
 };
 
@@ -135,67 +158,69 @@ pub fn callback_mappings() -> impl Iterator<Item = CallbackMapping> {
         .chain(PYTHON_CALLBACKS.iter().copied())
 }
 
-impl CallHooks<PythonRuntime> for LegacyLogging {
-    fn prepare_arguments(
-        &mut self,
-        py: Python<'_>,
+impl CallHooks<PythonRuntime> for LegacyHooks {
+    fn prepare_arguments<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
         arguments: Py<PyDict>,
         started_at: f64,
-    ) -> Step<Py<PyDict>> {
-        (PREPARE.invoke)(self, py, arguments, started_at)
+    ) -> impl Future<Output = PyResult<Py<PyDict>>> + Send {
+        (PREPARE.invoke)(self, task, arguments, started_at)
     }
 
-    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
-        self.adopt_arguments(py, arguments);
+    fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
+        self.adopt_arguments(arguments);
         Ok(())
     }
 
-    fn before_provider_request(
-        &mut self,
-        py: Python<'_>,
+    fn before_provider_request<'a>(
+        &'a self,
+        _: &'a dyn Awaiter,
         wire: Box<WireRequest>,
         context: &RequestContext,
-    ) -> Step<Box<WireRequest>> {
-        (BEFORE.invoke)(self, py, wire, context)
+    ) -> impl Future<Output = PyResult<Box<WireRequest>>> + Send {
+        ready((BEFORE.invoke)(self, wire, context))
     }
 
-    fn transform_response(
-        &mut self,
-        py: Python<'_>,
+    fn transform_response<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
         response: Py<PyAny>,
         timing: Timing,
-    ) -> Step<Py<PyAny>> {
-        (TRANSFORM.invoke)(self, py, response, timing)
+    ) -> impl Future<Output = PyResult<Py<PyAny>>> + Send {
+        (TRANSFORM.invoke)(self, task, response, timing)
     }
 
-    fn on_event(&mut self, py: Python<'_>, event: PythonCallEvent<'_>) -> Step<()> {
+    async fn on_event<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
+        event: PythonCallEvent<'_>,
+    ) -> PyResult<()> {
         match event {
-            PythonCallEvent::Started { .. } | PythonCallEvent::Cancelled { .. } => {
-                Ok(HookStep::Ready(()))
-            }
+            PythonCallEvent::Started { .. } | PythonCallEvent::Cancelled { .. } => Ok(()),
             PythonCallEvent::Execution(ExecutionEvent::ResultReady { facts }) => {
-                self.result_ready(py, &facts)
+                self.result_ready(&facts)
             }
             PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
-                (AFTER.invoke)(self, py, raw)
+                (AFTER.invoke)(self, raw)
             }
             PythonCallEvent::Succeeded { timing, response } => {
-                (SUCCESS.invoke)(self, py, timing, response)
+                (SUCCESS.invoke)(self, timing, response)
             }
             PythonCallEvent::Failed {
                 timing,
                 origin,
                 error,
-            } => (FAILURE.invoke)(self, py, timing, origin, error),
+            } => (FAILURE.invoke)(self, task, timing, origin, error).await,
         }
     }
 
-    fn on_stream_open(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
-        (OPEN.invoke)(self, py, head)
+    fn on_stream_open(&self, head: &Py<PyAny>) -> PyResult<()> {
+        (OPEN.invoke)(self, head)
     }
 
-    fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
-        (CHUNK.invoke)(self, py, chunk)
+    fn on_stream_chunk(&self, chunk: &Py<PyAny>) -> PyResult<()> {
+        (CHUNK.invoke)(self, chunk)
     }
 }
 

litellm-rust/crates/callbacks-legacy-python/src/test_support.rs

@@ -1,9 +1,17 @@
+use std::collections::VecDeque;
 use std::ffi::CStr;
+use std::future::{Future, ready};
+use std::pin::pin;
+use std::task::{Context, Poll, Waker};
+
+use futures_util::future::BoxFuture;
+use litellm_host_python::Awaiter;
+use parking_lot::Mutex;
 
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyTuple};
 
-use crate::{LegacyLogging, PublicCall};
+use crate::{LegacyHooks, LegacyLogging, PublicCall};
 
 /// The parameters of every `callbacks_legacy_python` function, as the real module declares them.
 /// `tests/unit/rust_bridge/host/test_callbacks_legacy_python.py` pins this file to the Python
@@ -193,3 +201,51 @@ pub(crate) fn legacy_call(
     let call = PublicCall::capture(&dict("bound"), &PyTuple::empty(py), &dict("kwargs")).unwrap();
     LegacyLogging::new(py, crate::LoggingOperation::Ocr, call, asynchronous)
 }
+
+/// Answers each awaitable a hook hands it from `answers` in order, then with the awaitable
+/// itself, and counts what it was handed.
+#[derive(Default)]
+pub(crate) struct Scripted {
+    answers: Mutex<VecDeque<PyResult<Py<PyAny>>>>,
+    awaited: Mutex<usize>,
+}
+
+impl Scripted {
+    pub(crate) fn answering(answers: impl IntoIterator<Item = PyResult<Py<PyAny>>>) -> Self {
+        Self {
+            answers: Mutex::new(answers.into_iter().collect()),
+            awaited: Mutex::new(0),
+        }
+    }
+
+    pub(crate) fn count(&self) -> usize {
+        *self.awaited.lock()
+    }
+}
+
+impl Awaiter for Scripted {
+    fn awaited(&self, awaitable: Py<PyAny>) -> BoxFuture<'_, PyResult<Py<PyAny>>> {
+        *self.awaited.lock() += 1;
+        let answer = self.answers.lock().pop_front().unwrap_or(Ok(awaitable));
+        Box::pin(ready(answer))
+    }
+}
+
+/// Runs a hook whose every await `Scripted` answers inline.
+pub(crate) fn finish<T>(hook: impl Future<Output = T>) -> T {
+    match pin!(hook).poll(&mut Context::from_waker(Waker::noop())) {
+        Poll::Ready(value) => value,
+        Poll::Pending => panic!("a scripted hook never suspends"),
+    }
+}
+
+pub(crate) fn hooks(logging: LegacyLogging) -> LegacyHooks {
+    LegacyHooks::new(logging)
+}
+
+pub(crate) fn failure(locals: &Bound<'_, PyDict>) -> Py<pyo3::exceptions::PyBaseException> {
+    local(locals, "failure")
+        .cast_into::<pyo3::exceptions::PyBaseException>()
+        .unwrap()
+        .unbind()
+}

litellm-rust/crates/host-python/src/caller.rs

@@ -1,10 +1,11 @@
 use std::ops::ControlFlow;
 
+use futures_util::future::BoxFuture;
 use litellm_coroutine::{Abandoned, Co, Reply};
 use pyo3::prelude::*;
 
-use crate::Error;
 use crate::handle::Suspension;
+use crate::{Awaiter, Error};
 
 pub(crate) type Answer = PyResult<Py<PyAny>>;
 
@@ -56,3 +57,9 @@ impl PythonCaller {
         }
     }
 }
+
+impl Awaiter for PythonCaller {
+    fn awaited(&self, awaitable: Py<PyAny>) -> BoxFuture<'_, Answer> {
+        Box::pin(self.await_python(awaitable))
+    }
+}

litellm-rust/crates/host-python/src/driver.rs

@@ -20,7 +20,7 @@ use litellm_host::{
 use crate::PythonHostCalls;
 use crate::caller::{Answer, PythonCaller, Yielded};
 use crate::handle::{Execution, ExecutionBody, ExecutionStep, StepTypes};
-use crate::hooks::{HookStep, PythonCallEvent, PythonCallHooks};
+use crate::hooks::{PythonCallEvent, PythonCallHooks};
 use crate::native::{NativeMachine, NativePoll};
 use crate::{Error, InvokeError, PythonBinding};
 
@@ -91,18 +91,27 @@ impl Observation {
     }
 }
 
-/// What the running call shares with its handle: the binding and hooks the handle
-/// traverses and closes, and the Python values the call retains across a suspension.
-struct Parts<H, L> {
-    binding: H,
+/// Hooks lock their own state around Python work, so the handle can traverse them while a
+/// hook is suspended.
+struct Shared<H, L> {
     hooks: L,
+    parts: Mutex<Parts<H>>,
+}
+
+struct Parts<H> {
+    binding: H,
     arguments: Option<Py<PyDict>>,
     interrupted: Option<Py<PyBaseException>>,
-    selected: Option<Py<PyAny>>,
+    selected: Option<Selected>,
     observation: Observation,
 }
 
-type SharedParts<H, L> = Arc<Mutex<Parts<H, L>>>;
+/// The call's outcome while its terminal hooks run. The body borrows it for the event and
+/// holds no Python reference of its own, so the one reference is the traversed one.
+enum Selected {
+    Response(Arc<Py<PyAny>>),
+    Exception(Arc<Py<PyBaseException>>),
+}
 
 enum Setup<E> {
     Host(PyErr),
@@ -121,7 +130,7 @@ fn classified<H: PythonBinding>(py: Python<'_>, binding: &H, error: ErrorOf<H>)
     classifier_error
 }
 
-/// The call as one async body. Every hook step and host call is awaited in place on the
+/// The call as one async body. Every hook and host call is awaited in place on the
 /// caller's task; the parts lock is never held across an await.
 struct CallBody<H, M, L>
 where
@@ -129,7 +138,7 @@ where
     M: Machine,
 {
     caller: PythonCaller,
-    parts: SharedParts<H, L>,
+    shared: Arc<Shared<H, L>>,
     native: NativeMachine<M>,
     start: Option<StartMachine<ProtocolOf<H>, M>>,
     streaming: bool,
@@ -142,8 +151,12 @@ where
     M: Machine<Protocol = H::Protocol> + 'static,
     M::Complete: Into<HostedCompletion<ResponseOf<H>>>,
 {
-    fn parts(&self) -> MutexGuard<'_, Parts<H, L>> {
-        self.parts.lock()
+    fn parts(&self) -> MutexGuard<'_, Parts<H>> {
+        self.shared.parts.lock()
+    }
+
+    fn hooks(&self) -> &L {
+        &self.shared.hooks
     }
 
     async fn run(mut self) -> PyResult<Py<PyAny>> {
@@ -159,45 +172,39 @@ where
 
     async fn begin(&mut self) -> Result<(), Setup<ErrorOf<H>>> {
         let started_at = epoch_seconds();
-        let started = {
+        let started = PythonCallEvent::Started {
+            start_time: started_at,
+        };
+        {
             let mut parts = self.parts();
             parts.observation.started_at = started_at;
             parts.observation.observing = true;
-            let event = PythonCallEvent::Started {
-                start_time: started_at,
-            };
-            parts.observation.emit(&event);
-            Python::attach(|py| parts.hooks.on_event(py, event)).map_err(Setup::Host)?
-        };
-        self.settle(started).await.map_err(Setup::Host)?;
+            parts.observation.emit(&started);
+        }
+        self.hooks()
+            .on_event(&self.caller, started)
+            .await
+            .map_err(Setup::Host)?;
         let arguments = self
             .parts()
             .arguments
             .take()
             .ok_or_else(|| Setup::Host(Error::MissingState.into()))?;
-        let prepare = {
+        let prepared = self
+            .hooks()
+            .prepare_arguments(&self.caller, arguments, started_at)
+            .await
+            .map_err(Setup::Host)?;
+        self.hooks()
+            .arguments_prepared(&prepared)
+            .map_err(Setup::Host)?;
+        let decoded = {
             let mut parts = self.parts();
-            Python::attach(|py| parts.hooks.prepare_arguments(py, arguments, started_at))
-                .map_err(Setup::Host)?
-        };
-        let prepared = self.settle(prepare).await.map_err(Setup::Host)?;
-        let request = {
-            let mut parts = self.parts();
-            let Parts {
-                binding,
-                hooks,
-                arguments,
-                ..
-            } = &mut *parts;
-            Python::attach(|py| {
-                hooks
-                    .arguments_prepared(py, &prepared)
-                    .map_err(Setup::Host)?;
-                let decoded = binding.decode_request(py, prepared.bind(py));
-                *arguments = Some(prepared);
-                decoded.map_err(Setup::Decode)
-            })?
+            let decoded = Python::attach(|py| parts.binding.decode_request(py, prepared.bind(py)));
+            parts.arguments = Some(prepared);
+            decoded
         };
+        let request = decoded.map_err(Setup::Decode)?;
         let start = self
             .start
             .take()
@@ -214,22 +221,6 @@ where
         Ok(())
     }
 
-    /// Runs one hook step to its value, handing each awaitable it produces to the caller.
-    async fn settle<T>(&self, step: HookStep<L, T>) -> PyResult<T> {
-        let mut step = step;
-        loop {
-            let (awaitable, resume) = match step {
-                HookStep::Ready(value) => return Ok(value),
-                HookStep::Await(awaitable, resume) => (awaitable, resume),
-            };
-            let answer = self.caller.await_python(awaitable).await;
-            step = {
-                let mut parts = self.parts();
-                Python::attach(|py| resume(&mut parts.hooks, py, answer))?
-            };
-        }
-    }
-
     async fn execute(&mut self) -> PyResult<Py<PyAny>> {
         let mut interruption: Interruption<H> = None;
         loop {
@@ -286,11 +277,10 @@ where
                 context,
                 reply,
             }) => {
-                let step = {
-                    let mut parts = self.parts();
-                    Python::attach(|py| parts.hooks.before_provider_request(py, wire, &context))?
-                };
-                let wire = self.settle(step).await?;
+                let wire = self
+                    .hooks()
+                    .before_provider_request(&self.caller, wire, &context)
+                    .await?;
                 reply.send(*wire);
                 Ok(())
             }
@@ -314,26 +304,18 @@ where
                 self.streaming = true;
                 let head = {
                     let mut parts = self.parts();
-                    let Parts { binding, hooks, .. } = &mut *parts;
-                    Python::attach(|py| {
-                        let head = binding.encode_stream_head(py, head)?;
-                        hooks.on_stream_open(py, &head)?;
-                        Ok::<_, PyErr>(head)
-                    })?
+                    Python::attach(|py| parts.binding.encode_stream_head(py, head))?
                 };
+                self.hooks().on_stream_open(&head)?;
                 reply.send(self.caller.open_stream(head).await);
                 Ok(())
             }
             HostRequest::Stream(StreamDelivery::Chunk(chunk, reply)) => {
                 let chunk = {
                     let mut parts = self.parts();
-                    let Parts { binding, hooks, .. } = &mut *parts;
-                    Python::attach(|py| {
-                        let chunk = binding.encode_chunk(py, chunk)?;
-                        hooks.on_stream_chunk(py, &chunk)?;
-                        Ok::<_, PyErr>(chunk)
-                    })?
+                    Python::attach(|py| parts.binding.encode_chunk(py, chunk))?
                 };
+                self.hooks().on_stream_chunk(&chunk)?;
                 reply.send(self.caller.yield_chunk(chunk).await);
                 Ok(())
             }
@@ -359,12 +341,8 @@ where
     }
 
     async fn notify(&self, event: PythonCallEvent<'_>) -> PyResult<()> {
-        let step = {
-            let mut parts = self.parts();
-            parts.observation.emit(&event);
-            Python::attach(|py| parts.hooks.on_event(py, event))?
-        };
-        self.settle(step).await
+        self.parts().observation.emit(&event);
+        self.hooks().on_event(&self.caller, event).await
     }
 
     async fn completed(
@@ -390,16 +368,12 @@ where
         if self.streaming {
             return self.succeeded(public).await;
         }
-        let transform = {
-            let mut parts = self.parts();
-            let timing = parts.observation.timing();
-            Python::attach(|py| parts.hooks.transform_response(py, public, timing))
-        };
-        let step = match transform {
-            Ok(step) => step,
-            Err(error) => return self.failure(error, FailureOrigin::Host).await,
-        };
-        match self.settle(step).await {
+        let timing = self.parts().observation.timing();
+        match self
+            .hooks()
+            .transform_response(&self.caller, public, timing)
+            .await
+        {
             Ok(response) => self.succeeded(response).await,
             Err(error) => self.failure(error, FailureOrigin::Host).await,
         }
@@ -417,58 +391,60 @@ where
         self.failure(error, FailureOrigin::Call).await
     }
 
-    /// The terminal hooks may suspend. Until they finish, the selected response is held
-    /// only in the traversed parts, so a cycle through it stays collectable.
     async fn succeeded(&mut self, response: Py<PyAny>) -> PyResult<Py<PyAny>> {
-        let step = {
-            let mut parts = self.parts();
-            let timing = parts.observation.timing();
-            let event = PythonCallEvent::Succeeded {
-                timing,
-                response: &response,
-            };
-            parts.observation.retain_terminal(&event);
-            let step = Python::attach(|py| parts.hooks.on_event(py, event))?;
-            parts.selected = Some(response);
-            step
+        let response = Arc::new(response);
+        let timing = self.parts().observation.timing();
+        let event = PythonCallEvent::Succeeded {
+            timing,
+            response: &response,
         };
-        self.settle(step).await?;
-        let mut parts = self.parts();
-        parts.observation.finish();
-        parts
-            .selected
-            .take()
-            .ok_or_else(|| Error::MissingState.into())
+        {
+            let mut parts = self.parts();
+            parts.observation.retain_terminal(&event);
+            parts.selected = Some(Selected::Response(Arc::clone(&response)));
+        }
+        let notified = self.hooks().on_event(&self.caller, event).await;
+        {
+            let mut parts = self.parts();
+            parts.selected = None;
+            notified?;
+            parts.observation.finish();
+        }
+        Arc::into_inner(response).ok_or_else(|| Error::MissingState.into())
     }
 
-    /// As [`Self::succeeded`], with the selected exception held as its value and raised
-    /// again once the hooks finish.
     async fn failure(&mut self, error: PyErr, origin: FailureOrigin) -> PyResult<Py<PyAny>> {
-        let step = {
+        let timing = {
             let mut parts = self.parts();
             parts.observation.ended_at.get_or_insert_with(epoch_seconds);
-            let timing = parts.observation.timing();
-            Python::attach(|py| {
-                if is_cancellation(py, &error) {
-                    return Err(error);
-                }
-                let event = PythonCallEvent::Failed {
-                    timing,
-                    origin,
-                    error: &error,
-                };
-                parts.observation.retain_terminal(&event);
-                let step = parts.hooks.on_event(py, event)?;
-                parts.selected = Some(error.into_value(py).into_any());
-                Ok(step)
-            })?
+            parts.observation.timing()
         };
-        self.settle(step).await?;
-        let mut parts = self.parts();
-        parts.observation.finish();
-        let selected = parts.selected.take().ok_or(Error::MissingState)?;
+        let exception = Python::attach(|py| {
+            if is_cancellation(py, &error) {
+                return Err(error);
+            }
+            Ok(Arc::new(error.into_value(py)))
+        })?;
+        let event = PythonCallEvent::Failed {
+            timing,
+            origin,
+            error: &exception,
+        };
+        {
+            let mut parts = self.parts();
+            parts.observation.retain_terminal(&event);
+            parts.selected = Some(Selected::Exception(Arc::clone(&exception)));
+        }
+        let notified = self.hooks().on_event(&self.caller, event).await;
+        {
+            let mut parts = self.parts();
+            parts.selected = None;
+            notified?;
+            parts.observation.finish();
+        }
+        let exception = Arc::into_inner(exception).ok_or(Error::MissingState)?;
         Err(Python::attach(|py| {
-            PyErr::from_value(selected.into_bound(py))
+            PyErr::from_value(exception.into_bound(py).into_any())
         }))
     }
 }
@@ -480,7 +456,7 @@ where
 {
     body: Mutex<Body>,
     reply: Option<Reply<Answer>>,
-    parts: SharedParts<H, L>,
+    shared: Arc<Shared<H, L>>,
 }
 
 impl<H, L> CallDriver<H, L>
@@ -491,19 +467,20 @@ where
     /// A failure the body raised past its terminal hooks, or a cancellation, is the last
     /// thing the observers hear about the call.
     fn observe_host_failure(&self, error: &PyErr) {
-        let mut parts = self.parts.lock();
+        let mut parts = self.shared.parts.lock();
         if !parts.observation.observing {
             return;
         }
         let timing = parts.observation.timing();
         Python::attach(|py| {
+            let exception = error.value(py).clone().unbind();
             let event = if is_cancellation(py, error) {
                 PythonCallEvent::Cancelled { timing }
             } else {
                 PythonCallEvent::Failed {
                     timing,
                     origin: FailureOrigin::Host,
-                    error,
+                    error: &exception,
                 }
             };
             parts.observation.emit(&event);
@@ -550,14 +527,18 @@ where
     }
 
     fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
-        let Some(parts) = self.parts.try_lock() else {
+        self.shared.hooks.traverse(visit)?;
+        let Some(parts) = self.shared.parts.try_lock() else {
             return Ok(());
         };
         parts.binding.traverse(visit)?;
-        parts.hooks.traverse(visit)?;
         visit.call(&parts.arguments)?;
         visit.call(&parts.interrupted)?;
-        visit.call(&parts.selected)
+        match &parts.selected {
+            Some(Selected::Response(response)) => visit.call(&**response),
+            Some(Selected::Exception(exception)) => visit.call(&**exception),
+            None => Ok(()),
+        }
     }
 }
 
@@ -569,17 +550,22 @@ where
     fn drop(&mut self) {
         self.body.get_mut().cancel();
         self.reply = None;
-        let mut parts = self.parts.lock();
+        let Some(Shared { hooks, parts }) = Arc::get_mut(&mut self.shared) else {
+            return;
+        };
+        let Parts {
+            binding,
+            observation,
+            ..
+        } = parts.get_mut();
         Python::attach(|py| {
-            if parts.observation.observing {
-                let timing = parts.observation.timing();
-                parts
-                    .observation
-                    .emit(&PythonCallEvent::Cancelled { timing });
-                parts.observation.observing = false;
+            if observation.observing {
+                let timing = observation.timing();
+                observation.emit(&PythonCallEvent::Cancelled { timing });
+                observation.observing = false;
             }
-            parts.hooks.close(py);
-            parts.binding.close(py);
+            hooks.close(py);
+            binding.close(py);
         });
     }
 }
@@ -605,27 +591,29 @@ where
     M: Machine<Protocol = H::Protocol> + 'static,
     M::Complete: Into<HostedCompletion<ResponseOf<H>>>,
 {
-    let parts = Arc::new(Mutex::new(Parts {
-        binding,
+    let shared = Arc::new(Shared {
         hooks,
-        arguments: Some(arguments),
-        interrupted: None,
-        selected: None,
-        observation: Observation {
-            observers: options.observers,
-            observing: false,
-            started_at: 0.0,
-            ended_at: None,
-            terminal: None,
-        },
-    }));
+        parts: Mutex::new(Parts {
+            binding,
+            arguments: Some(arguments),
+            interrupted: None,
+            selected: None,
+            observation: Observation {
+                observers: options.observers,
+                observing: false,
+                started_at: 0.0,
+                ended_at: None,
+                terminal: None,
+            },
+        }),
+    });
     let body = Coroutine::new({
-        let parts = Arc::clone(&parts);
+        let shared = Arc::clone(&shared);
         let asynchronous = options.asynchronous;
         move |co| {
             CallBody {
                 caller: PythonCaller::new(co),
-                parts,
+                shared,
                 native: NativeMachine::new(asynchronous),
                 start: Some(Box::new(start)),
                 streaming: false,
@@ -636,7 +624,7 @@ where
     let driver = CallDriver {
         body: Mutex::new(body),
         reply: None,
-        parts,
+        shared,
     };
     Py::new(py, Execution::new(driver, options.steps)).map(Py::into_any)
 }

litellm-rust/crates/host-python/src/hooks.rs

@@ -2,32 +2,27 @@ mod chain;
 pub use chain::HookChain;
 
 use crate::PythonOwned;
+use futures_util::future::BoxFuture;
 use litellm_host::hooks::{CallHooks, HookRuntime, RuntimeCallEvent};
+use pyo3::exceptions::PyBaseException;
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
 
-/// What a hook step produced: either the value the driver asked for, or a Python
-/// awaitable the driver hands back to the caller's task before asking again.
-pub type HookResume<L, T> = fn(&mut L, Python<'_>, PyResult<Py<PyAny>>) -> PyResult<HookStep<L, T>>;
-
-pub enum HookStep<L, T> {
-    Await(Py<PyAny>, HookResume<L, T>),
-    Ready(T),
+/// Where a hook or host call awaits a Python awaitable: the caller's task, which the call
+/// suspends to and resumes from with the awaitable's result.
+pub trait Awaiter: Sync {
+    fn awaited(&self, awaitable: Py<PyAny>) -> BoxFuture<'_, PyResult<Py<PyAny>>>;
 }
 
 pub struct PythonRuntime;
 
 impl HookRuntime for PythonRuntime {
-    type Context<'a> = Python<'a>;
+    type Context<'a> = &'a dyn Awaiter;
     type Arguments = Py<PyDict>;
     type Response = Py<PyAny>;
     type Chunk = Py<PyAny>;
     type Error = PyErr;
-    type Step<H, T> = HookStep<H, T>;
-
-    fn ready<H, T>(value: T) -> Self::Step<H, T> {
-        HookStep::Ready(value)
-    }
+    type Failure = Py<PyBaseException>;
 }
 
 pub type PythonCallEvent<'a> = RuntimeCallEvent<'a, PythonRuntime>;

litellm-rust/crates/host-python/src/hooks/chain.rs

@@ -0,0 +1,279 @@
+use std::future::{Future, pending, poll_fn};
+use std::pin::pin;
+use std::task::Poll;
+
+use futures_util::future::BoxFuture;
+use litellm_host::{
+    hooks::CallHooks,
+    interceptors::{RequestContext, WireRequest},
+    lifecycle::{CallEvent, Timing},
+};
+use parking_lot::Mutex;
+use pyo3::{
+    exceptions::PyException,
+    gc::{PyTraverseError, PyVisit},
+    prelude::*,
+    types::PyDict,
+};
+
+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<()>;
+}
+
+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,
+        ))
+    }
+
+    fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
+        CallHooks::arguments_prepared(self, arguments)
+    }
+
+    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 is_cancellation(error: &PyErr) -> bool {
+    Python::attach(|py| !error.is_instance_of::<PyException>(py))
+}
+
+/// The caller's task as one chain stage hands it to its hooks. A cancellation the task
+/// resumes with ends the stage there: the hook that awaited it never runs again.
+struct Stage<'a> {
+    task: &'a dyn Awaiter,
+    cancelled: Mutex<Option<PyErr>>,
+}
+
+impl<'a> Stage<'a> {
+    fn new(task: &'a dyn Awaiter) -> Self {
+        Self {
+            task,
+            cancelled: Mutex::new(None),
+        }
+    }
+
+    async fn guard<T>(&self, hook: impl Future<Output = PyResult<T>>) -> PyResult<T> {
+        let mut hook = pin!(hook);
+        poll_fn(|context| match hook.as_mut().poll(context) {
+            Poll::Ready(result) => Poll::Ready(result),
+            Poll::Pending => self
+                .cancelled
+                .lock()
+                .take()
+                .map_or(Poll::Pending, |error| Poll::Ready(Err(error))),
+        })
+        .await
+    }
+}
+
+impl Awaiter for Stage<'_> {
+    fn awaited(&self, awaitable: Py<PyAny>) -> BoxFuture<'_, PyResult<Py<PyAny>>> {
+        Box::pin(async move {
+            match self.task.awaited(awaitable).await {
+                Err(error) if is_cancellation(&error) => {
+                    *self.cancelled.lock() = Some(error);
+                    pending().await
+                }
+                answer => answer,
+            }
+        })
+    }
+}
+
+fn is_terminal<Response, Error, Raw>(event: &CallEvent<Response, Error, Raw>) -> bool {
+    matches!(
+        event,
+        CallEvent::Succeeded { .. } | CallEvent::Failed { .. }
+    )
+}
+
+/// A terminal notification reaches every hook, so an ordinary error from one is reported
+/// and the next hook still runs.
+fn notification_result(terminal: bool, result: PyResult<()>) -> PyResult<()> {
+    match result {
+        Err(error) if terminal && !is_cancellation(&error) => {
+            Python::attach(|py| error.write_unraisable(py, None));
+            Ok(())
+        }
+        result => result,
+    }
+}
+
+#[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 {
+    async fn prepare_arguments<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
+        arguments: Py<PyDict>,
+        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)
+    }
+
+    fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
+        self.hooks
+            .iter()
+            .try_for_each(|hooks| hooks.arguments_prepared(arguments))
+    }
+
+    async fn before_provider_request<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
+        wire: Box<WireRequest>,
+        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)
+    }
+
+    async fn transform_response<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
+        response: Py<PyAny>,
+        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)
+    }
+
+    async fn on_event<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
+        event: PythonCallEvent<'_>,
+    ) -> 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(())
+    }
+
+    fn on_stream_open(&self, head: &Py<PyAny>) -> PyResult<()> {
+        self.hooks
+            .iter()
+            .try_for_each(|hooks| hooks.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))
+    }
+}
+
+impl PythonOwned for HookChain {
+    fn close(&mut self, py: Python<'_>) {
+        for hooks in &mut self.hooks {
+            hooks.close(py);
+        }
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        self.hooks
+            .iter()
+            .try_for_each(|hooks| hooks.traverse(visit))
+    }
+}

litellm-rust/crates/host-python/src/hooks/chain/adapter.rs

@@ -1,202 +0,0 @@
-use litellm_host::interceptors::{RequestContext, WireRequest};
-use litellm_host::lifecycle::Timing;
-use pyo3::{
-    gc::{PyTraverseError, PyVisit},
-    prelude::*,
-    types::PyDict,
-};
-
-use crate::{Error, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned};
-
-pub(super) enum ChainStep<T> {
-    Ready(T),
-    Await(Py<PyAny>),
-}
-
-enum Continuation<H> {
-    Arguments(HookResume<H, Py<PyDict>>),
-    Wire(HookResume<H, Box<WireRequest>>),
-    Response(HookResume<H, Py<PyAny>>),
-    Event(HookResume<H, ()>),
-}
-
-pub(super) struct HookAdapter<H> {
-    hooks: H,
-    continuation: Option<Continuation<H>>,
-}
-
-impl<H> HookAdapter<H> {
-    pub(super) fn new(hooks: H) -> Self {
-        Self {
-            hooks,
-            continuation: None,
-        }
-    }
-
-    fn step<T>(
-        &mut self,
-        step: HookStep<H, T>,
-        continuation: impl FnOnce(HookResume<H, T>) -> Continuation<H>,
-    ) -> ChainStep<T> {
-        match step {
-            HookStep::Ready(value) => ChainStep::Ready(value),
-            HookStep::Await(awaitable, resume) => {
-                self.continuation = Some(continuation(resume));
-                ChainStep::Await(awaitable)
-            }
-        }
-    }
-}
-
-pub(super) trait ChainHooks: PythonOwned {
-    fn prepare_arguments(
-        &mut self,
-        py: Python<'_>,
-        arguments: Py<PyDict>,
-        started_at: f64,
-    ) -> PyResult<ChainStep<Py<PyDict>>>;
-    fn resume_arguments(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<ChainStep<Py<PyDict>>>;
-    fn before_provider_request(
-        &mut self,
-        py: Python<'_>,
-        wire: Box<WireRequest>,
-        context: &RequestContext,
-    ) -> PyResult<ChainStep<Box<WireRequest>>>;
-    fn resume_wire(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<ChainStep<Box<WireRequest>>>;
-    fn transform_response(
-        &mut self,
-        py: Python<'_>,
-        response: Py<PyAny>,
-        timing: Timing,
-    ) -> PyResult<ChainStep<Py<PyAny>>>;
-    fn resume_response(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<ChainStep<Py<PyAny>>>;
-    fn on_event(&mut self, py: Python<'_>, event: PythonCallEvent<'_>) -> PyResult<ChainStep<()>>;
-    fn resume_event(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<ChainStep<()>>;
-    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()>;
-    fn on_stream_open(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()>;
-    fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()>;
-}
-
-impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
-    fn prepare_arguments(
-        &mut self,
-        py: Python<'_>,
-        arguments: Py<PyDict>,
-        started_at: f64,
-    ) -> PyResult<ChainStep<Py<PyDict>>> {
-        let step = self.hooks.prepare_arguments(py, arguments, started_at)?;
-        Ok(self.step(step, Continuation::Arguments))
-    }
-
-    fn resume_arguments(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<ChainStep<Py<PyDict>>> {
-        let Some(Continuation::Arguments(resume)) = self.continuation.take() else {
-            return Err(Error::MissingState.into());
-        };
-        let step = resume(&mut self.hooks, py, result)?;
-        Ok(self.step(step, Continuation::Arguments))
-    }
-
-    fn before_provider_request(
-        &mut self,
-        py: Python<'_>,
-        wire: Box<WireRequest>,
-        context: &RequestContext,
-    ) -> PyResult<ChainStep<Box<WireRequest>>> {
-        let step = self.hooks.before_provider_request(py, wire, context)?;
-        Ok(self.step(step, Continuation::Wire))
-    }
-
-    fn resume_wire(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<ChainStep<Box<WireRequest>>> {
-        let Some(Continuation::Wire(resume)) = self.continuation.take() else {
-            return Err(Error::MissingState.into());
-        };
-        let step = resume(&mut self.hooks, py, result)?;
-        Ok(self.step(step, Continuation::Wire))
-    }
-
-    fn transform_response(
-        &mut self,
-        py: Python<'_>,
-        response: Py<PyAny>,
-        timing: Timing,
-    ) -> PyResult<ChainStep<Py<PyAny>>> {
-        let step = self.hooks.transform_response(py, response, timing)?;
-        Ok(self.step(step, Continuation::Response))
-    }
-
-    fn resume_response(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<ChainStep<Py<PyAny>>> {
-        let Some(Continuation::Response(resume)) = self.continuation.take() else {
-            return Err(Error::MissingState.into());
-        };
-        let step = resume(&mut self.hooks, py, result)?;
-        Ok(self.step(step, Continuation::Response))
-    }
-
-    fn on_event(&mut self, py: Python<'_>, event: PythonCallEvent<'_>) -> PyResult<ChainStep<()>> {
-        let step = self.hooks.on_event(py, event)?;
-        Ok(self.step(step, Continuation::Event))
-    }
-
-    fn resume_event(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<ChainStep<()>> {
-        let Some(Continuation::Event(resume)) = self.continuation.take() else {
-            return Err(Error::MissingState.into());
-        };
-        let step = resume(&mut self.hooks, py, result)?;
-        Ok(self.step(step, Continuation::Event))
-    }
-
-    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
-        self.hooks.arguments_prepared(py, arguments)
-    }
-
-    fn on_stream_open(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
-        self.hooks.on_stream_open(py, head)
-    }
-
-    fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
-        self.hooks.on_stream_chunk(py, chunk)
-    }
-}
-
-impl<H: PythonOwned> PythonOwned for HookAdapter<H> {
-    fn close(&mut self, py: Python<'_>) {
-        self.continuation = None;
-        self.hooks.close(py);
-    }
-
-    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
-        self.hooks.traverse(visit)
-    }
-}

litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs

@@ -1,394 +0,0 @@
-use litellm_host::{
-    hooks::CallHooks,
-    interceptors::{RawResponse, RequestContext, WireRequest},
-    lifecycle::{CallEvent, ExecutionEvent, Timing},
-};
-use pyo3::{
-    exceptions::{PyBaseException, PyException},
-    gc::{PyTraverseError, PyVisit},
-    prelude::*,
-    types::PyDict,
-};
-
-use crate::{Error, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime};
-
-use super::adapter::{ChainHooks, ChainStep, HookAdapter};
-
-type OwnedEvent = CallEvent<Py<PyAny>, Py<PyBaseException>, RawResponse>;
-
-#[derive(Default)]
-pub struct HookChain {
-    hooks: Vec<Box<dyn ChainHooks>>,
-    arguments: Option<(usize, f64)>,
-    wire: Option<(usize, RequestContext)>,
-    response: Option<(usize, Timing)>,
-    event: Option<(usize, OwnedEvent)>,
-}
-
-impl HookChain {
-    pub fn new() -> Self {
-        Self::default()
-    }
-
-    pub fn with(mut self, hooks: impl PythonCallHooks + 'static) -> Self {
-        self.hooks.push(Box::new(HookAdapter::new(hooks)));
-        self
-    }
-
-    fn arguments_from(
-        &mut self,
-        py: Python<'_>,
-        index: usize,
-        arguments: Py<PyDict>,
-        started_at: f64,
-    ) -> PyResult<HookStep<Self, Py<PyDict>>> {
-        let Some(hooks) = self.hooks.get_mut(index) else {
-            return Ok(HookStep::Ready(arguments));
-        };
-        let step = hooks.prepare_arguments(py, arguments, started_at)?;
-        self.arguments_step(py, index, step, started_at)
-    }
-
-    fn arguments_step(
-        &mut self,
-        py: Python<'_>,
-        index: usize,
-        step: ChainStep<Py<PyDict>>,
-        started_at: f64,
-    ) -> PyResult<HookStep<Self, Py<PyDict>>> {
-        match step {
-            ChainStep::Ready(value) => self.arguments_from(py, index + 1, value, started_at),
-            ChainStep::Await(awaitable) => {
-                self.arguments = Some((index, started_at));
-                Ok(HookStep::Await(awaitable, Self::resume_arguments))
-            }
-        }
-    }
-
-    fn resume_arguments(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<HookStep<Self, Py<PyDict>>> {
-        let (index, context) = self.arguments.take().ok_or(Error::MissingState)?;
-        let result = resume_unless_cancelled(py, result)?;
-        let step = self.hooks[index].resume_arguments(py, result)?;
-        self.arguments_step(py, index, step, context)
-    }
-
-    fn wire_from(
-        &mut self,
-        py: Python<'_>,
-        index: usize,
-        wire: Box<WireRequest>,
-        context: &RequestContext,
-    ) -> PyResult<HookStep<Self, Box<WireRequest>>> {
-        let Some(hooks) = self.hooks.get_mut(index) else {
-            return Ok(HookStep::Ready(wire));
-        };
-        let step = hooks.before_provider_request(py, wire, context)?;
-        self.wire_step(py, index, step, context)
-    }
-
-    fn wire_step(
-        &mut self,
-        py: Python<'_>,
-        index: usize,
-        step: ChainStep<Box<WireRequest>>,
-        context: &RequestContext,
-    ) -> PyResult<HookStep<Self, Box<WireRequest>>> {
-        match step {
-            ChainStep::Ready(value) => self.wire_from(py, index + 1, value, context),
-            ChainStep::Await(awaitable) => {
-                self.wire = Some((index, context.clone()));
-                Ok(HookStep::Await(awaitable, Self::resume_wire))
-            }
-        }
-    }
-
-    fn resume_wire(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<HookStep<Self, Box<WireRequest>>> {
-        let (index, context) = self.wire.take().ok_or(Error::MissingState)?;
-        let result = resume_unless_cancelled(py, result)?;
-        let step = self.hooks[index].resume_wire(py, result)?;
-        self.wire_step(py, index, step, &context)
-    }
-
-    fn response_from(
-        &mut self,
-        py: Python<'_>,
-        index: usize,
-        response: Py<PyAny>,
-        timing: Timing,
-    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
-        let Some(hooks) = self.hooks.get_mut(index) else {
-            return Ok(HookStep::Ready(response));
-        };
-        let step = hooks.transform_response(py, response, timing)?;
-        self.response_step(py, index, step, timing)
-    }
-
-    fn response_step(
-        &mut self,
-        py: Python<'_>,
-        index: usize,
-        step: ChainStep<Py<PyAny>>,
-        timing: Timing,
-    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
-        match step {
-            ChainStep::Ready(value) => self.response_from(py, index + 1, value, timing),
-            ChainStep::Await(awaitable) => {
-                self.response = Some((index, timing));
-                Ok(HookStep::Await(awaitable, Self::resume_response))
-            }
-        }
-    }
-
-    fn resume_response(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
-        let (index, context) = self.response.take().ok_or(Error::MissingState)?;
-        let result = resume_unless_cancelled(py, result)?;
-        let step = self.hooks[index].resume_response(py, result)?;
-        self.response_step(py, index, step, context)
-    }
-
-    fn event_from(
-        &mut self,
-        py: Python<'_>,
-        index: usize,
-        event: OwnedEvent,
-    ) -> PyResult<HookStep<Self, ()>> {
-        let Some(hooks) = self.hooks.get_mut(index) else {
-            return Ok(HookStep::Ready(()));
-        };
-        let result = dispatch(py, hooks.as_mut(), &event);
-        let step = notification_result(py, is_terminal(&event), result)?;
-        self.event_step(py, index, step, event)
-    }
-
-    fn event_step(
-        &mut self,
-        py: Python<'_>,
-        index: usize,
-        step: ChainStep<()>,
-        event: OwnedEvent,
-    ) -> PyResult<HookStep<Self, ()>> {
-        match step {
-            ChainStep::Ready(()) => self.event_from(py, index + 1, event),
-            ChainStep::Await(awaitable) => {
-                self.event = Some((index, event));
-                Ok(HookStep::Await(awaitable, Self::resume_event))
-            }
-        }
-    }
-
-    fn resume_event(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<HookStep<Self, ()>> {
-        let (index, event) = self.event.take().ok_or(Error::MissingState)?;
-        let result = resume_unless_cancelled(py, result)?;
-        let result = self.hooks[index].resume_event(py, result);
-        let step = notification_result(py, is_terminal(&event), result)?;
-        self.event_step(py, index, step, event)
-    }
-}
-
-fn resume_unless_cancelled(
-    py: Python<'_>,
-    result: PyResult<Py<PyAny>>,
-) -> PyResult<PyResult<Py<PyAny>>> {
-    match result {
-        Err(error) if !error.is_instance_of::<PyException>(py) => Err(error),
-        result => Ok(result),
-    }
-}
-
-fn is_terminal<Response, Error, Raw>(event: &CallEvent<Response, Error, Raw>) -> bool {
-    matches!(
-        event,
-        CallEvent::Succeeded { .. } | CallEvent::Failed { .. }
-    )
-}
-
-fn notification_result(
-    py: Python<'_>,
-    terminal: bool,
-    result: PyResult<ChainStep<()>>,
-) -> PyResult<ChainStep<()>> {
-    match result {
-        Err(error) if terminal && error.is_instance_of::<PyException>(py) => {
-            error.write_unraisable(py, None);
-            Ok(ChainStep::Ready(()))
-        }
-        result => result,
-    }
-}
-
-fn retain_event(py: Python<'_>, event: PythonCallEvent<'_>) -> OwnedEvent {
-    match event {
-        CallEvent::Execution(ExecutionEvent::ResultReady { facts }) => {
-            CallEvent::Execution(ExecutionEvent::ResultReady { facts })
-        }
-        CallEvent::Started { start_time } => CallEvent::Started { start_time },
-        CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
-            CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw: raw.clone() })
-        }
-        CallEvent::Succeeded { timing, response } => CallEvent::Succeeded {
-            timing,
-            response: response.clone_ref(py),
-        },
-        CallEvent::Failed {
-            timing,
-            origin,
-            error,
-        } => CallEvent::Failed {
-            timing,
-            origin,
-            error: error.clone_ref(py).into_value(py),
-        },
-        CallEvent::Cancelled { timing } => CallEvent::Cancelled { timing },
-    }
-}
-
-fn dispatch(
-    py: Python<'_>,
-    hooks: &mut dyn ChainHooks,
-    event: &OwnedEvent,
-) -> PyResult<ChainStep<()>> {
-    match event {
-        CallEvent::Execution(ExecutionEvent::ResultReady { facts }) => hooks.on_event(
-            py,
-            CallEvent::Execution(ExecutionEvent::ResultReady {
-                facts: facts.clone(),
-            }),
-        ),
-        CallEvent::Started { start_time } => hooks.on_event(
-            py,
-            CallEvent::Started {
-                start_time: *start_time,
-            },
-        ),
-        CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => hooks.on_event(
-            py,
-            CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
-        ),
-        CallEvent::Succeeded { timing, response } => hooks.on_event(
-            py,
-            CallEvent::Succeeded {
-                timing: *timing,
-                response,
-            },
-        ),
-        CallEvent::Failed {
-            timing,
-            origin,
-            error,
-        } => hooks.on_event(
-            py,
-            CallEvent::Failed {
-                timing: *timing,
-                origin: *origin,
-                error: &PyErr::from_value(error.bind(py).clone().into_any()),
-            },
-        ),
-        CallEvent::Cancelled { timing } => {
-            hooks.on_event(py, CallEvent::Cancelled { timing: *timing })
-        }
-    }
-}
-
-impl CallHooks<PythonRuntime> for HookChain {
-    fn prepare_arguments(
-        &mut self,
-        py: Python<'_>,
-        arguments: Py<PyDict>,
-        started_at: f64,
-    ) -> PyResult<HookStep<Self, Py<PyDict>>> {
-        self.arguments_from(py, 0, arguments, started_at)
-    }
-
-    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
-        self.hooks
-            .iter_mut()
-            .try_for_each(|hooks| hooks.arguments_prepared(py, arguments))
-    }
-
-    fn before_provider_request(
-        &mut self,
-        py: Python<'_>,
-        wire: Box<WireRequest>,
-        context: &RequestContext,
-    ) -> PyResult<HookStep<Self, Box<WireRequest>>> {
-        self.wire_from(py, 0, wire, context)
-    }
-
-    fn transform_response(
-        &mut self,
-        py: Python<'_>,
-        response: Py<PyAny>,
-        timing: Timing,
-    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
-        self.response_from(py, 0, response, timing)
-    }
-
-    fn on_event(
-        &mut self,
-        py: Python<'_>,
-        event: PythonCallEvent<'_>,
-    ) -> PyResult<HookStep<Self, ()>> {
-        for (index, hooks) in self.hooks.iter_mut().enumerate() {
-            let result = hooks.on_event(py, event.clone());
-            match notification_result(py, is_terminal(&event), result)? {
-                ChainStep::Ready(()) => {}
-                ChainStep::Await(awaitable) => {
-                    self.event = Some((index, retain_event(py, event)));
-                    return Ok(HookStep::Await(awaitable, Self::resume_event));
-                }
-            }
-        }
-        Ok(HookStep::Ready(()))
-    }
-
-    fn on_stream_open(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
-        self.hooks
-            .iter_mut()
-            .try_for_each(|hooks| hooks.on_stream_open(py, head))
-    }
-
-    fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
-        self.hooks
-            .iter_mut()
-            .try_for_each(|hooks| hooks.on_stream_chunk(py, chunk))
-    }
-}
-
-impl PythonOwned for HookChain {
-    fn close(&mut self, py: Python<'_>) {
-        self.arguments = None;
-        self.wire = None;
-        self.response = None;
-        self.event = None;
-        for hooks in &mut self.hooks {
-            hooks.close(py);
-        }
-    }
-
-    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
-        for hooks in &self.hooks {
-            hooks.traverse(visit)?;
-        }
-        match &self.event {
-            Some((_, CallEvent::Succeeded { response, .. })) => visit.call(response),
-            Some((_, CallEvent::Failed { error, .. })) => visit.call(error),
-            _ => Ok(()),
-        }
-    }
-}

litellm-rust/crates/host-python/src/hooks/chain/mod.rs

@@ -1,4 +0,0 @@
-mod adapter;
-mod dispatch;
-
-pub use dispatch::HookChain;

litellm-rust/crates/host-python/src/lib.rs

@@ -29,7 +29,7 @@ pub use file_reader::{FileContent, PythonFileReader, py_bytes};
 pub use fork_gate::RuntimeAlreadyStarted;
 pub use gil::{PythonContext, attach_blocking, release_count, release_gil};
 pub use handle::{Execution, ExecutionBody, ExecutionStep, StepTypes, Suspension};
-pub use hooks::{HookChain, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonRuntime};
+pub use hooks::{Awaiter, HookChain, PythonCallEvent, PythonCallHooks, PythonRuntime};
 pub use marshal::{
     Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py,
 };

litellm-rust/crates/host-python/tests/driver.rs

@@ -535,11 +535,11 @@ enum ArgumentPolicy {
 }
 
 impl CallHooks<PythonRuntime> for ArgumentPolicy {
-    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
-        match self {
+    fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
+        Python::attach(|py| match self {
             Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"),
             Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())),
-        }
+        })
     }
 }
 

litellm-rust/crates/host-python/tests/handle.rs

@@ -2,9 +2,10 @@ mod support;
 
 use litellm_host::{hooks::CallHooks, lifecycle::Timing};
 use litellm_host_python::{
-    Execution, ExecutionBody, ExecutionStep, HookStep, PythonCallEvent, PythonOwned, PythonRuntime,
+    Awaiter, Execution, ExecutionBody, ExecutionStep, PythonCallEvent, PythonOwned, PythonRuntime,
     Suspension, run_call,
 };
+use parking_lot::Mutex;
 use pyo3::{
     exceptions::PyBaseException,
     gc::{PyTraverseError, PyVisit},
@@ -114,55 +115,59 @@ enum Holder {
 }
 
 struct PausingHooks {
-    retained: Option<Py<PyAny>>,
-    pause: Option<Py<PyAny>>,
+    retained: Mutex<Option<Py<PyAny>>>,
+    pause: Mutex<Option<Py<PyAny>>>,
     holder: Holder,
 }
 
 impl CallHooks<PythonRuntime> for PausingHooks {
-    fn transform_response(
-        &mut self,
-        py: Python<'_>,
+    async fn transform_response<'a>(
+        &'a self,
+        _: &'a dyn Awaiter,
         response: Py<PyAny>,
         _: Timing,
-    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
+    ) -> PyResult<Py<PyAny>> {
+        let retained = || self.retained.lock().take().unwrap();
         match self.holder {
-            Holder::Hooks => Ok(HookStep::Ready(response)),
-            Holder::Response => Ok(HookStep::Ready(self.retained.take().unwrap())),
-            Holder::Exception => Err(PyErr::from_value(
-                self.retained.take().unwrap().into_bound(py),
-            )),
+            Holder::Hooks => Ok(response),
+            Holder::Response => Ok(retained()),
+            Holder::Exception => Err(Python::attach(|py| {
+                PyErr::from_value(retained().into_bound(py))
+            })),
         }
     }
 
-    fn on_event(
-        &mut self,
-        _: Python<'_>,
+    async fn on_event<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
         event: PythonCallEvent<'_>,
-    ) -> PyResult<HookStep<Self, ()>> {
+    ) -> PyResult<()> {
         let pauses_here = match self.holder {
             Holder::Hooks => true,
             Holder::Response => matches!(event, PythonCallEvent::Succeeded { .. }),
             Holder::Exception => matches!(event, PythonCallEvent::Failed { .. }),
         };
-        Ok(match self.pause.take_if(|_| pauses_here) {
-            Some(pause) => {
-                HookStep::Await(pause, |_, _, answer| answer.map(|_| HookStep::Ready(())))
-            }
-            None => HookStep::Ready(()),
-        })
+        let pause = self.pause.lock().take_if(|_| pauses_here);
+        if let Some(pause) = pause {
+            task.awaited(pause).await?;
+        }
+        Ok(())
     }
 }
 
 impl PythonOwned for PausingHooks {
     fn close(&mut self, _: Python<'_>) {
-        self.retained = None;
-        self.pause = None;
+        *self.retained.get_mut() = None;
+        *self.pause.get_mut() = None;
     }
 
     fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
-        visit.call(&self.retained)?;
-        visit.call(&self.pause)
+        for slot in [&self.retained, &self.pause] {
+            if let Some(value) = slot.try_lock() {
+                visit.call(&*value)?;
+            }
+        }
+        Ok(())
     }
 }
 
@@ -180,8 +185,8 @@ fn paused_call(
             move |_, _, request| Ok(machine(request)),
             binding(&Log::default(), OpScript::Answer),
             PausingHooks {
-                retained: Some(retained),
-                pause: Some(pause),
+                retained: Mutex::new(Some(retained)),
+                pause: Mutex::new(Some(pause)),
                 holder,
             },
             PyDict::new(py).unbind(),

litellm-rust/crates/host-python/tests/hook_chain.rs

@@ -1,74 +1,57 @@
+mod support;
+
+use std::future::{Future, ready};
+
+use futures_util::{FutureExt, future::BoxFuture};
 use litellm_host::{
     hooks::CallHooks,
     interceptors::{RequestContext, WireRequest},
     lifecycle::{FailureOrigin, Timing},
 };
-use litellm_host_python::{HookChain, HookStep, PythonCallEvent, PythonOwned, PythonRuntime};
+use litellm_host_python::{
+    Awaiter, HookChain, PythonCallEvent, PythonOwned, PythonRuntime, run_call,
+};
 use pyo3::{
     gc::{PyTraverseError, PyVisit},
     prelude::*,
     types::{PyDict, PyTuple},
 };
 use rstest::{fixture, rstest};
+use support::{
+    Interpreter, Log, OpScript, binding, call_options, handed, interpreter, success_machine,
+};
 
 struct ScriptHooks {
     object: Py<PyAny>,
     asynchronous: bool,
-    wire: Option<Box<WireRequest>>,
 }
 
 impl ScriptHooks {
-    fn invoke(&self, py: Python<'_>, name: &str, value: Py<PyAny>) -> PyResult<Py<PyAny>> {
+    async fn invoke(
+        &self,
+        task: &dyn Awaiter,
+        name: &str,
+        value: impl FnOnce(Python<'_>) -> PyResult<Py<PyAny>>,
+    ) -> PyResult<Py<PyAny>> {
+        let result = Python::attach(|py| {
+            let value = value(py)?;
+            if self.asynchronous {
+                self.object.call_method1(py, "invoke", (name, value))
+            } else {
+                self.object.call_method1(py, name, (value,))
+            }
+        })?;
         if self.asynchronous {
-            self.object.call_method1(py, "invoke", (name, value))
+            task.awaited(result).await
         } else {
-            self.object.call_method1(py, name, (value,))
+            Ok(result)
         }
     }
-
-    fn arguments(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<HookStep<Self, Py<PyDict>>> {
-        Ok(HookStep::Ready(
-            result?.into_bound(py).cast_into::<PyDict>()?.unbind(),
-        ))
-    }
-
-    fn response(
-        &mut self,
-        _: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
-        result.map(HookStep::Ready)
-    }
-
-    fn notification(
-        &mut self,
-        _: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<HookStep<Self, ()>> {
-        result.map(|_| HookStep::Ready(()))
-    }
-
-    fn request(
-        &mut self,
-        py: Python<'_>,
-        result: PyResult<Py<PyAny>>,
-    ) -> PyResult<HookStep<Self, Box<WireRequest>>> {
-        let wire = self.wire.take().unwrap();
-        Ok(HookStep::Ready(Box::new(WireRequest {
-            url: result?.extract(py)?,
-            ..*wire
-        })))
-    }
 }
 
 impl PythonOwned for ScriptHooks {
     fn close(&mut self, py: Python<'_>) {
         self.object = py.None();
-        self.wire = None;
     }
 
     fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
@@ -77,98 +60,105 @@ impl PythonOwned for ScriptHooks {
 }
 
 impl CallHooks<PythonRuntime> for ScriptHooks {
-    fn prepare_arguments(
-        &mut self,
-        py: Python<'_>,
+    async fn prepare_arguments<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
         arguments: Py<PyDict>,
         _: f64,
-    ) -> PyResult<HookStep<Self, Py<PyDict>>> {
-        let value = self.invoke(py, "prepare", arguments.into_any())?;
-        if self.asynchronous {
-            Ok(HookStep::Await(value, Self::arguments))
-        } else {
-            self.arguments(py, Ok(value))
-        }
+    ) -> PyResult<Py<PyDict>> {
+        let value = self
+            .invoke(task, "prepare", |_| Ok(arguments.into_any()))
+            .await?;
+        Python::attach(|py| Ok(value.into_bound(py).cast_into::<PyDict>()?.unbind()))
     }
 
-    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
-        self.object.call_method1(py, "adopt", (arguments,))?;
+    fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
+        Python::attach(|py| self.object.call_method1(py, "adopt", (arguments,)))?;
         Ok(())
     }
 
-    fn before_provider_request(
-        &mut self,
-        py: Python<'_>,
+    async fn before_provider_request<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
         wire: Box<WireRequest>,
         context: &RequestContext,
-    ) -> PyResult<HookStep<Self, Box<WireRequest>>> {
-        self.object.bind(py).setattr("model", &context.model)?;
-        let value = self.invoke(
-            py,
-            "before",
-            wire.url.clone().into_pyobject(py)?.into_any().unbind(),
-        )?;
-        self.wire = Some(wire);
-        if self.asynchronous {
-            Ok(HookStep::Await(value, Self::request))
-        } else {
-            self.request(py, Ok(value))
-        }
+    ) -> PyResult<Box<WireRequest>> {
+        Python::attach(|py| self.object.bind(py).setattr("model", &context.model))?;
+        let url = self
+            .invoke(task, "before", |py| {
+                Ok(wire.url.clone().into_pyobject(py)?.into_any().unbind())
+            })
+            .await?;
+        Ok(Box::new(WireRequest {
+            url: Python::attach(|py| url.extract(py))?,
+            ..*wire
+        }))
     }
 
-    fn transform_response(
-        &mut self,
-        py: Python<'_>,
+    async fn transform_response<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
         response: Py<PyAny>,
         _: Timing,
-    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
-        let value = self.invoke(py, "transform", response)?;
-        if self.asynchronous {
-            Ok(HookStep::Await(value, Self::response))
-        } else {
-            self.response(py, Ok(value))
-        }
+    ) -> PyResult<Py<PyAny>> {
+        self.invoke(task, "transform", |_| Ok(response)).await
     }
 
-    fn on_event(
-        &mut self,
-        py: Python<'_>,
+    async fn on_event<'a>(
+        &'a self,
+        task: &'a dyn Awaiter,
         event: PythonCallEvent<'_>,
-    ) -> PyResult<HookStep<Self, ()>> {
-        let (name, value) = match event {
-            PythonCallEvent::Succeeded { response, .. } => ("success", response.clone_ref(py)),
-            PythonCallEvent::Failed { error, .. } => {
-                ("failure", error.clone_ref(py).into_value(py).into_any())
-            }
-            PythonCallEvent::Started { .. } => ("started", py.None()),
-            PythonCallEvent::Cancelled { .. } => ("cancelled", py.None()),
-            PythonCallEvent::Execution(_) => ("provider", py.None()),
-        };
-        let value = self.invoke(
-            py,
-            "event",
-            PyTuple::new(py, [name.into_pyobject(py)?.into_any().unbind(), value])?
-                .into_any()
-                .unbind(),
-        )?;
-        if self.asynchronous {
-            Ok(HookStep::Await(value, Self::notification))
-        } else {
-            self.notification(py, Ok(value))
-        }
+    ) -> PyResult<()> {
+        self.invoke(task, "event", |py| {
+            let (name, value) = match event {
+                PythonCallEvent::Succeeded { response, .. } => ("success", response.clone_ref(py)),
+                PythonCallEvent::Failed { error, .. } => {
+                    ("failure", error.clone_ref(py).into_any())
+                }
+                PythonCallEvent::Started { .. } => ("started", py.None()),
+                PythonCallEvent::Cancelled { .. } => ("cancelled", py.None()),
+                PythonCallEvent::Execution(_) => ("provider", py.None()),
+            };
+            Ok(
+                PyTuple::new(py, [name.into_pyobject(py)?.into_any().unbind(), value])?
+                    .into_any()
+                    .unbind(),
+            )
+        })
+        .await
+        .map(drop)
     }
 
-    fn on_stream_open(&mut self, py: Python<'_>, _head: &Py<PyAny>) -> PyResult<()> {
-        self.object.call_method1(py, "stream", (py.None(),))?;
+    fn on_stream_open(&self, _head: &Py<PyAny>) -> PyResult<()> {
+        Python::attach(|py| self.object.call_method1(py, "stream", (py.None(),)))?;
         Ok(())
     }
 
-    fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
-        self.object.call_method1(py, "stream", (chunk,))?;
+    fn on_stream_chunk(&self, chunk: &Py<PyAny>) -> PyResult<()> {
+        Python::attach(|py| self.object.call_method1(py, "stream", (chunk,)))?;
         Ok(())
     }
 }
 
+/// Runs every awaitable a hook hands it to completion on a fresh event loop.
+struct RunToCompletion;
+
+impl Awaiter for RunToCompletion {
+    fn awaited(&self, awaitable: Py<PyAny>) -> BoxFuture<'_, PyResult<Py<PyAny>>> {
+        let result = Python::attach(|py| {
+            py.import("asyncio")?
+                .call_method1("run", (awaitable,))
+                .map(Bound::unbind)
+        });
+        Box::pin(ready(result))
+    }
+}
+
+fn finish<T>(hook: impl Future<Output = PyResult<T>>) -> PyResult<T> {
+    hook.now_or_never()
+        .expect("every awaitable runs to completion inline")
+}
+
 #[fixture]
 fn scripts() -> Py<PyDict> {
     Python::initialize();
@@ -223,25 +213,10 @@ fn chain(py: Python<'_>, scripts: &Py<PyDict>, asynchronous: bool) -> HookChain
     let hook = |name| ScriptHooks {
         object: scripts.bind(py).get_item(name).unwrap().unwrap().unbind(),
         asynchronous,
-        wire: None,
     };
     HookChain::new().with(hook("first")).with(hook("second"))
 }
 
-fn finish<H, T>(py: Python<'_>, hooks: &mut H, step: HookStep<H, T>) -> PyResult<T> {
-    match step {
-        HookStep::Ready(value) => Ok(value),
-        HookStep::Await(awaitable, resume) => {
-            let value = py
-                .import("asyncio")?
-                .call_method1("run", (awaitable,))
-                .map(Bound::unbind);
-            let next = resume(hooks, py, value)?;
-            finish(py, hooks, next)
-        }
-    }
-}
-
 const TIMING: Timing = Timing {
     start_time: 0.0,
     end_time: 1.0,
@@ -250,8 +225,8 @@ const TIMING: Timing = Timing {
 struct PreparedPolicy;
 
 impl CallHooks<PythonRuntime> for PreparedPolicy {
-    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
-        arguments.bind(py).set_item("policy", "configured")
+    fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
+        Python::attach(|py| arguments.bind(py).set_item("policy", "configured"))
     }
 }
 
@@ -271,13 +246,11 @@ fn transformations_feed_each_other_and_notifications_share_final_values(
     #[case] asynchronous: bool,
 ) {
     Python::attach(|py| {
-        let mut hooks = chain(py, &scripts, asynchronous).with(PreparedPolicy);
+        let hooks = chain(py, &scripts, asynchronous).with(PreparedPolicy);
         let original = PyDict::new(py).unbind();
-        let step = hooks
-            .prepare_arguments(py, original.clone_ref(py), 0.0)
-            .unwrap();
-        let arguments = finish(py, &mut hooks, step).unwrap();
-        hooks.arguments_prepared(py, &arguments).unwrap();
+        let arguments =
+            finish(hooks.prepare_arguments(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap();
+        hooks.arguments_prepared(&arguments).unwrap();
         let context = RequestContext {
             model: "model".into(),
             custom_llm_provider: "provider".into(),
@@ -290,25 +263,22 @@ fn transformations_feed_each_other_and_notifications_share_final_values(
             headers: vec![],
             body: serde_json::Value::Null,
         });
-        let step = hooks.before_provider_request(py, wire, &context).unwrap();
-        assert_eq!(finish(py, &mut hooks, step).unwrap().url, "urlab");
+        let wire = finish(hooks.before_provider_request(&RunToCompletion, wire, &context)).unwrap();
+        assert_eq!(wire.url, "urlab");
         let response = PyDict::new(py).unbind().into_any();
-        let step = hooks
-            .transform_response(py, response.clone_ref(py), TIMING)
-            .unwrap();
-        let final_response = finish(py, &mut hooks, step).unwrap();
-        let step = hooks
-            .on_event(
-                py,
-                PythonCallEvent::Succeeded {
-                    timing: TIMING,
-                    response: &final_response,
-                },
-            )
-            .unwrap();
-        finish(py, &mut hooks, step).unwrap();
-        hooks.on_stream_open(py, &py.None()).unwrap();
-        hooks.on_stream_chunk(py, &response).unwrap();
+        let final_response =
+            finish(hooks.transform_response(&RunToCompletion, response.clone_ref(py), TIMING))
+                .unwrap();
+        finish(hooks.on_event(
+            &RunToCompletion,
+            PythonCallEvent::Succeeded {
+                timing: TIMING,
+                response: &final_response,
+            },
+        ))
+        .unwrap();
+        hooks.on_stream_open(&py.None()).unwrap();
+        hooks.on_stream_chunk(&response).unwrap();
         let locals = scripts.bind(py);
         locals.set_item("arguments", arguments).unwrap();
         locals.set_item("original", original).unwrap();
@@ -348,7 +318,7 @@ fn terminal_failure_does_not_skip_later_hooks(
     #[values("first", "second")] failing_hook: &str,
 ) {
     Python::attach(|py| {
-        let mut hooks = chain(py, &scripts, asynchronous).with(ScriptHooks {
+        let hooks = chain(py, &scripts, asynchronous).with(ScriptHooks {
             object: scripts
                 .bind(py)
                 .get_item("third")
@@ -356,7 +326,6 @@ fn terminal_failure_does_not_skip_later_hooks(
                 .unwrap()
                 .unbind(),
             asynchronous,
-            wire: None,
         });
         let locals = scripts.bind(py);
         locals
@@ -368,7 +337,7 @@ fn terminal_failure_does_not_skip_later_hooks(
                 pyo3::exceptions::PyRuntimeError::new_err("callback failed").into_value(py),
             )
             .unwrap();
-        let error = pyo3::exceptions::PyValueError::new_err("provider failed");
+        let error = pyo3::exceptions::PyValueError::new_err("provider failed").into_value(py);
         let response = py.None();
         let event = if failed {
             PythonCallEvent::Failed {
@@ -382,17 +351,9 @@ fn terminal_failure_does_not_skip_later_hooks(
                 response: &response,
             }
         };
-        let step = hooks.on_event(py, event).unwrap();
-        finish(py, &mut hooks, step).unwrap();
+        finish(hooks.on_event(&RunToCompletion, event)).unwrap();
         locals
-            .set_item(
-                "selected",
-                if failed {
-                    error.into_value(py).into_any()
-                } else {
-                    response
-                },
-            )
+            .set_item("selected", if failed { error.into_any() } else { response })
             .unwrap();
         py.run(
             c"
@@ -412,7 +373,7 @@ assert all(value is selected for name, kind, value in log)
 #[case::async_(true)]
 fn transformation_failure_stops_the_chain(scripts: Py<PyDict>, #[case] asynchronous: bool) {
     Python::attach(|py| {
-        let mut hooks = chain(py, &scripts, asynchronous);
+        let hooks = chain(py, &scripts, asynchronous);
         let locals = scripts.bind(py);
         py.run(
             c"first.error = ValueError('prepare failed')",
@@ -420,9 +381,8 @@ fn transformation_failure_stops_the_chain(scripts: Py<PyDict>, #[case] asynchron
             Some(locals),
         )
         .unwrap();
-        let result = hooks
-            .prepare_arguments(py, PyDict::new(py).unbind(), 0.0)
-            .and_then(|step| finish(py, &mut hooks, step));
+        let result =
+            finish(hooks.prepare_arguments(&RunToCompletion, PyDict::new(py).unbind(), 0.0));
         assert!(
             result.unwrap_err().value(py).is(locals
                 .get_item("first")
@@ -445,7 +405,7 @@ fn transformation_failure_stops_the_chain(scripts: Py<PyDict>, #[case] asynchron
 #[case::async_(true)]
 fn cancellation_stops_notification_dispatch(scripts: Py<PyDict>, #[case] asynchronous: bool) {
     Python::attach(|py| {
-        let mut hooks = chain(py, &scripts, asynchronous);
+        let hooks = chain(py, &scripts, asynchronous);
         let locals = scripts.bind(py);
         py.run(
             c"first.error = asyncio.CancelledError()",
@@ -454,15 +414,13 @@ fn cancellation_stops_notification_dispatch(scripts: Py<PyDict>, #[case] asynchr
         )
         .unwrap();
         let response = py.None();
-        let result = hooks
-            .on_event(
-                py,
-                PythonCallEvent::Succeeded {
-                    timing: TIMING,
-                    response: &response,
-                },
-            )
-            .and_then(|step| finish(py, &mut hooks, step));
+        let result = finish(hooks.on_event(
+            &RunToCompletion,
+            PythonCallEvent::Succeeded {
+                timing: TIMING,
+                response: &response,
+            },
+        ));
         assert!(
             result
                 .unwrap_err()
@@ -477,56 +435,11 @@ fn cancellation_stops_notification_dispatch(scripts: Py<PyDict>, #[case] asynchr
     });
 }
 
-struct Preparing {
-    hooks: HookChain,
-    arguments: Option<Py<PyDict>>,
-    resume: Option<litellm_host_python::HookResume<HookChain, Py<PyDict>>>,
-}
-
-impl litellm_host_python::ExecutionBody for Preparing {
-    fn resume(
-        &mut self,
-        result: Option<PyResult<Py<PyAny>>>,
-    ) -> PyResult<litellm_host_python::ExecutionStep> {
-        Python::attach(|py| {
-            let step = match self.resume.take() {
-                Some(resume) => resume(&mut self.hooks, py, result.unwrap())?,
-                None => self
-                    .hooks
-                    .prepare_arguments(py, self.arguments.take().unwrap(), 0.0)?,
-            };
-            match step {
-                HookStep::Ready(arguments) => {
-                    self.hooks.arguments_prepared(py, &arguments)?;
-                    Ok(litellm_host_python::ExecutionStep::Return(
-                        arguments.into_any(),
-                    ))
-                }
-                HookStep::Await(awaitable, resume) => {
-                    self.resume = Some(resume);
-                    Ok(litellm_host_python::ExecutionStep::Suspended(
-                        litellm_host_python::Suspension::Await(awaitable),
-                    ))
-                }
-            }
-        })
-    }
-
-    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
-        visit.call(&self.arguments)?;
-        self.hooks.traverse(visit)
-    }
-}
-
-fn lifecycle(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
-    let source =
-        std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/execution.py"))
-            .unwrap();
-    PyModule::from_code(py, &source, c"execution.py", c"hook_chain_test_lifecycle")
-}
-
 #[rstest]
-fn suspended_hooks_keep_the_callers_task_and_context(scripts: Py<PyDict>) {
+fn suspended_hooks_keep_the_callers_task_and_context(
+    #[from(interpreter)] _python: Interpreter,
+    scripts: Py<PyDict>,
+) {
     Python::attach(|py| {
         let locals = scripts.bind(py);
         py.run(
@@ -537,8 +450,9 @@ state = contextvars.ContextVar('state')
 async def invoke(name, value):
     assert asyncio.current_task() is caller
     assert threading.get_ident() == thread
-    if name == 'prepare':
-        state.set(state.get() + 'x')
+    if name != 'prepare':
+        return value
+    state.set(state.get() + 'x')
     await asyncio.sleep(0)
     assert asyncio.current_task() is caller
     return {**value, 'context': state.get()}
@@ -549,24 +463,20 @@ second = Hooks('b', invoke)
             Some(locals),
         )
         .unwrap();
-        let hooks = chain(py, &scripts, true);
-        let execution = litellm_host_python::Execution::new(
-            Preparing {
-                hooks,
-                arguments: Some(PyDict::new(py).unbind()),
-                resume: None,
-            },
-            lifecycle,
-        );
-        let module = lifecycle(py).unwrap();
-        let call = module
-            .getattr("drive")
-            .unwrap()
-            .call1((
-                Py::new(py, execution).unwrap(),
-                module.getattr("AsyncStream").unwrap(),
-            ))
-            .unwrap();
+        let machine = success_machine();
+        let call = handed(
+            py,
+            run_call(
+                py,
+                move |_, _, request| Ok(machine(request)),
+                binding(&Log::default(), OpScript::Answer),
+                chain(py, &scripts, true),
+                PyDict::new(py).unbind(),
+                call_options(true),
+            ),
+            true,
+        )
+        .unwrap();
         locals.set_item("call", call).unwrap();
         py.run(
             c"
@@ -575,10 +485,10 @@ async def exercise():
     caller = asyncio.current_task()
     thread = threading.get_ident()
     state.set('caller')
-    result = await call
-    assert result['context'] == 'callerxx'
+    await call
+    assert first.adopted['context'] == 'callerxx'
     assert state.get() == 'callerxx'
-    assert first.adopted is result and second.adopted is result
+    assert second.adopted is first.adopted
 asyncio.run(exercise())
 ",
             Some(locals),
@@ -588,106 +498,6 @@ asyncio.run(exercise())
     });
 }
 
-#[pyclass(weakref)]
-struct HookOwner {
-    hooks: HookChain,
-}
-
-#[pymethods]
-impl HookOwner {
-    fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
-        self.hooks.traverse(&visit)
-    }
-
-    fn __clear__(&mut self, py: Python<'_>) {
-        self.hooks.close(py);
-    }
-}
-
-#[rstest]
-#[case::first_response(false, false)]
-#[case::first_exception(true, false)]
-#[case::second_response(false, true)]
-#[case::second_exception(true, true)]
-fn suspended_notification_cycles_are_collectable(
-    scripts: Py<PyDict>,
-    #[case] failed: bool,
-    #[case] first_ready: bool,
-) {
-    Python::attach(|py| {
-        let hook = |name, asynchronous| ScriptHooks {
-            object: scripts.bind(py).get_item(name).unwrap().unwrap().unbind(),
-            asynchronous,
-            wire: None,
-        };
-        let owner = Py::new(
-            py,
-            HookOwner {
-                hooks: HookChain::new()
-                    .with(hook("first", !first_ready))
-                    .with(hook("second", true)),
-            },
-        )
-        .unwrap();
-        let locals = scripts.bind(py);
-        locals.set_item("owner", &owner).unwrap();
-        py.run(
-            c"
-import gc
-import weakref
-class Payload(Exception):
-    pass
-payload = Payload()
-payload.owner = owner
-owner_ref = weakref.ref(owner)
-payload_ref = weakref.ref(payload)
-",
-            Some(locals),
-            Some(locals),
-        )
-        .unwrap();
-        let payload = locals.get_item("payload").unwrap().unwrap().unbind();
-        let step = if failed {
-            owner.borrow_mut(py).hooks.on_event(
-                py,
-                PythonCallEvent::Failed {
-                    timing: TIMING,
-                    origin: FailureOrigin::Call,
-                    error: &PyErr::from_value(payload.bind(py).clone()),
-                },
-            )
-        } else {
-            owner.borrow_mut(py).hooks.on_event(
-                py,
-                PythonCallEvent::Succeeded {
-                    timing: TIMING,
-                    response: &payload,
-                },
-            )
-        }
-        .unwrap();
-        let HookStep::Await(awaitable, _) = step else {
-            panic!("notification must suspend")
-        };
-        awaitable.call_method0(py, "close").unwrap();
-        drop(awaitable);
-        drop(payload);
-        drop(owner);
-        py.run(
-            c"
-log.clear()
-del owner, payload
-gc.collect()
-assert owner_ref() is None
-assert payload_ref() is None
-",
-            Some(locals),
-            Some(locals),
-        )
-        .unwrap();
-    });
-}
-
 #[rstest]
 #[case::empty(0, "")]
 #[case::single(1, "a")]
@@ -701,21 +511,18 @@ fn builder_runs_hooks_in_append_order(
     #[values(false, true)] asynchronous: bool,
 ) {
     Python::attach(|py| {
-        let mut hooks = ["first", "second", "third", "fourth"]
+        let hooks = ["first", "second", "third", "fourth"]
             .into_iter()
             .take(count)
             .fold(HookChain::new(), |chain, name| {
                 chain.with(ScriptHooks {
                     object: scripts.bind(py).get_item(name).unwrap().unwrap().unbind(),
                     asynchronous,
-                    wire: None,
                 })
             });
         let original = PyDict::new(py).unbind();
-        let step = hooks
-            .prepare_arguments(py, original.clone_ref(py), 0.0)
-            .unwrap();
-        let result = finish(py, &mut hooks, step).unwrap();
+        let result =
+            finish(hooks.prepare_arguments(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap();
         if count == 0 {
             assert!(result.bind(py).is(original.bind(py)));
         } else {

litellm-rust/crates/host-python/tests/support/mod.rs

@@ -11,7 +11,7 @@ use litellm_host::{
     protocol::{Protocol, Reply},
 };
 use litellm_host_python::{
-    CallOptions, HookStep, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks,
+    Awaiter, CallOptions, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks,
     PythonHostCalls, PythonOwned, PythonRuntime, run_call,
 };
 use pyo3::{
@@ -335,12 +335,12 @@ pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks {
 }
 
 impl CallHooks<PythonRuntime> for SyntheticHooks {
-    fn prepare_arguments(
-        &mut self,
-        _: Python<'_>,
+    async fn prepare_arguments<'a>(
+        &'a self,
+        _: &'a dyn Awaiter,
         arguments: Py<PyDict>,
         _: f64,
-    ) -> PyResult<HookStep<Self, Py<PyDict>>> {
+    ) -> PyResult<Py<PyDict>> {
         self.log.push("begin");
         if matches!(self.script, HookScript::FailBegin) {
             return Err(PyValueError::new_err("begin failed"));
@@ -352,57 +352,54 @@ impl CallHooks<PythonRuntime> for SyntheticHooks {
                 prepared.set_item("api_key", "hook-key")?;
                 Ok(prepared.unbind())
             })?;
-            return Ok(HookStep::Ready(prepared));
+            return Ok(prepared);
         }
-        Ok(HookStep::Ready(arguments))
+        Ok(arguments)
     }
 
-    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
+    fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
         if matches!(self.script, HookScript::ObserveArguments) {
-            let value: String = arguments
-                .bind(py)
-                .get_item("prepared")?
-                .unwrap()
-                .extract()?;
+            let value: String =
+                Python::attach(|py| arguments.bind(py).get_item("prepared")?.unwrap().extract())?;
             self.log.push(format!("adopted:{value}"));
         }
         Ok(())
     }
 
-    fn before_provider_request(
-        &mut self,
-        _: Python<'_>,
+    async fn before_provider_request<'a>(
+        &'a self,
+        _: &'a dyn Awaiter,
         wire: Box<WireRequest>,
         _: &RequestContext,
-    ) -> PyResult<HookStep<Self, Box<WireRequest>>> {
+    ) -> PyResult<Box<WireRequest>> {
         self.log.push("before_provider_request");
-        Ok(HookStep::Ready(Box::new(WireRequest {
+        Ok(Box::new(WireRequest {
             url: "rewritten".into(),
             ..*wire
-        })))
+        }))
     }
 
-    fn transform_response(
-        &mut self,
-        py: Python<'_>,
+    async fn transform_response<'a>(
+        &'a self,
+        _: &'a dyn Awaiter,
         response: Py<PyAny>,
         _: Timing,
-    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
+    ) -> PyResult<Py<PyAny>> {
         self.log.push("after_success");
         match self.script {
-            HookScript::ReplaceResponse => Ok(HookStep::Ready(
-                "replaced".into_pyobject(py)?.into_any().unbind(),
-            )),
+            HookScript::ReplaceResponse => {
+                Python::attach(|py| Ok("replaced".into_pyobject(py)?.into_any().unbind()))
+            }
             HookScript::FailAfterSuccess => Err(PyValueError::new_err("after_success failed")),
-            _ => Ok(HookStep::Ready(response)),
+            _ => Ok(response),
         }
     }
 
-    fn on_event(
-        &mut self,
-        py: Python<'_>,
+    async fn on_event<'a>(
+        &'a self,
+        _: &'a dyn Awaiter,
         event: PythonCallEvent<'_>,
-    ) -> PyResult<HookStep<Self, ()>> {
+    ) -> PyResult<()> {
         let terminal = matches!(
             event,
             PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. }
@@ -413,7 +410,7 @@ impl CallHooks<PythonRuntime> for SyntheticHooks {
         if terminal && matches!(self.script, HookScript::FailTerminal) {
             return Err(PyValueError::new_err("callback failed"));
         }
-        self.log.push(match event {
+        self.log.push(Python::attach(|py| match event {
             PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "cache_hit".into(),
             PythonCallEvent::Started { .. } => "started".into(),
             PythonCallEvent::Cancelled { .. } => "cancelled".into(),
@@ -424,18 +421,18 @@ impl CallHooks<PythonRuntime> for SyntheticHooks {
                 format!("succeeded:{}", response.bind(py))
             }
             PythonCallEvent::Failed { origin, error, .. } => {
-                format!("failed:{origin:?}:{}", error.value(py))
+                format!("failed:{origin:?}:{}", error.bind(py))
             }
-        });
-        Ok(HookStep::Ready(()))
+        }));
+        Ok(())
     }
 
-    fn on_stream_open(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
+    fn on_stream_open(&self, _: &Py<PyAny>) -> PyResult<()> {
         self.log.push("opened");
         Ok(())
     }
 
-    fn on_stream_chunk(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
+    fn on_stream_chunk(&self, _: &Py<PyAny>) -> PyResult<()> {
         self.log.push("delivered");
         Ok(())
     }

litellm-rust/crates/host/src/hooks.rs

@@ -1,79 +1,67 @@
+use std::future::{Future, ready};
+
 use crate::{
     interceptors::{RawResponse, RequestContext, WireRequest},
     lifecycle::{CallEvent, Timing},
 };
 
 pub trait HookRuntime {
-    type Context<'a>;
-    type Arguments;
-    type Response;
+    type Context<'a>: Send;
+    type Arguments: Send;
+    type Response: Send + Sync;
     type Chunk;
-    type Error;
-    type Step<H, T>;
-
-    fn ready<H, T>(value: T) -> Self::Step<H, T>;
+    type Error: Send + Sync;
+    type Failure: Sync;
 }
 
 pub type RuntimeCallEvent<'a, R> =
-    CallEvent<&'a <R as HookRuntime>::Response, &'a <R as HookRuntime>::Error, &'a RawResponse>;
+    CallEvent<&'a <R as HookRuntime>::Response, &'a <R as HookRuntime>::Failure, &'a RawResponse>;
 
-pub trait CallHooks<R: HookRuntime>: Sized {
-    fn prepare_arguments(
-        &mut self,
-        _runtime: R::Context<'_>,
+pub trait CallHooks<R: HookRuntime>: Sync {
+    fn prepare_arguments<'a>(
+        &'a self,
+        _runtime: R::Context<'a>,
         arguments: R::Arguments,
         _started_at: f64,
-    ) -> Result<R::Step<Self, R::Arguments>, R::Error> {
-        Ok(R::ready(arguments))
+    ) -> impl Future<Output = Result<R::Arguments, R::Error>> + Send {
+        ready(Ok(arguments))
     }
 
-    fn arguments_prepared(
-        &mut self,
-        _runtime: R::Context<'_>,
-        _arguments: &R::Arguments,
-    ) -> Result<(), R::Error> {
+    fn arguments_prepared(&self, _arguments: &R::Arguments) -> Result<(), R::Error> {
         Ok(())
     }
 
-    fn before_provider_request(
-        &mut self,
-        _runtime: R::Context<'_>,
+    fn before_provider_request<'a>(
+        &'a self,
+        _runtime: R::Context<'a>,
         wire: Box<WireRequest>,
         _context: &RequestContext,
-    ) -> Result<R::Step<Self, Box<WireRequest>>, R::Error> {
-        Ok(R::ready(wire))
+    ) -> impl Future<Output = Result<Box<WireRequest>, R::Error>> + Send {
+        ready(Ok(wire))
     }
 
-    fn transform_response(
-        &mut self,
-        _runtime: R::Context<'_>,
+    fn transform_response<'a>(
+        &'a self,
+        _runtime: R::Context<'a>,
         response: R::Response,
         _timing: Timing,
-    ) -> Result<R::Step<Self, R::Response>, R::Error> {
-        Ok(R::ready(response))
+    ) -> impl Future<Output = Result<R::Response, R::Error>> + Send {
+        ready(Ok(response))
     }
 
-    fn on_event(
-        &mut self,
-        _runtime: R::Context<'_>,
+    fn on_event<'a>(
+        &'a self,
+        _runtime: R::Context<'a>,
         _event: RuntimeCallEvent<'_, R>,
-    ) -> Result<R::Step<Self, ()>, R::Error> {
-        Ok(R::ready(()))
+    ) -> impl Future<Output = Result<(), R::Error>> + Send {
+        ready(Ok(()))
     }
 
-    fn on_stream_open(
-        &mut self,
-        _runtime: R::Context<'_>,
-        _head: &R::Response,
-    ) -> Result<(), R::Error> {
+    fn on_stream_open(&self, _head: &R::Response) -> Result<(), R::Error> {
         Ok(())
     }
 
-    fn on_stream_chunk(
-        &mut self,
-        _runtime: R::Context<'_>,
-        _chunk: &R::Chunk,
-    ) -> Result<(), R::Error> {
+    fn on_stream_chunk(&self, _chunk: &R::Chunk) -> Result<(), R::Error> {
         Ok(())
     }
 }

litellm-rust/crates/python-bridge/src/preflight.rs

@@ -36,14 +36,16 @@ pub(crate) const PYTHON_CONTRACT: &str = include_str!("../preflight_contract.jso
 pub(crate) struct SdkPolicy;
 
 impl CallHooks<PythonRuntime> for SdkPolicy {
-    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
-        inherit_credentials(py, arguments.bind(py), || {
-            Ok(PythonPreflight::CredentialList
-                .call(py, ())?
-                .cast_into::<PyList>()?)
-        })?;
-        PythonPreflight::CheckLimits.call(py, (arguments,))?;
-        Ok(())
+    fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
+        Python::attach(|py| {
+            inherit_credentials(py, arguments.bind(py), || {
+                Ok(PythonPreflight::CredentialList
+                    .call(py, ())?
+                    .cast_into::<PyList>()?)
+            })?;
+            PythonPreflight::CheckLimits.call(py, (arguments,))?;
+            Ok(())
+        })
     }
 }
 
@@ -426,7 +428,7 @@ arguments = {'litellm_credential_name': 'missing'}
 ",
             );
             SdkPolicy
-                .arguments_prepared(py, &argument_dict(&locals).unbind())
+                .arguments_prepared(&argument_dict(&locals).unbind())
                 .unwrap();
             py.run(
                 c"
@@ -464,7 +466,7 @@ arguments = {'litellm_credential_name': 'ocr-test'}
 ",
             );
             let error = SdkPolicy
-                .arguments_prepared(py, &argument_dict(&locals).unbind())
+                .arguments_prepared(&argument_dict(&locals).unbind())
                 .unwrap_err();
             assert!(
                 error

litellm-rust/crates/python-bridge/src/routes/mod.rs

@@ -9,7 +9,7 @@ pub(crate) mod responses;
 pub(crate) mod token_counter;
 pub(crate) mod traces;
 
-use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCall};
+use litellm_callbacks_legacy_python::{LegacyHooks, LegacyLogging, LoggingOperation, PublicCall};
 use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
 use litellm_host_python::{Error as HostError, present};
 use litellm_host_python::{
@@ -125,7 +125,7 @@ fn call_hooks(
     let arguments = call.arguments(py);
     Ok((
         arguments,
-        LegacyLogging::new(py, operation, call, asynchronous),
+        LegacyHooks::new(LegacyLogging::new(py, operation, call, asynchronous)),
     ))
 }