refactor(host): split call hooks into facets with the policy in the return type
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled

One trait per payload replaces CallHooks and its on_event bag: ArgumentHooks,
ProviderHooks, ResponseHooks and OutcomeHooks over a Surface type bag. Transforms
fold through the groups and stop on the first error, outcomes reach every group
and only a Cancelled group stops the rest, so ChainRuntime, HookGroupAdapter,
DynCallHooks and the runtime terminal match are gone. Hooks<S> composes mounted
groups, request stages in mount order, response and terminal stages in reverse

host-python keeps PythonSurface, the per-stage cancellation guard behind
PythonSurface::group, and report_unless_cancelled as the one place an ordinary
outcome error becomes unraisable. The legacy adapter implements each facet in
its own file under facets/, which also lists the Python callbacks each stage
reaches, and mapping.rs only collects them. SdkPolicy does its work in prepare,
so the logging-before-policy ordering rule is no longer load-bearing
This commit is contained in:
Yujong Lee 2026-10-14 10:16:21 -07:00
parent 9380590298
commit c63d46becc
27 changed files with 1505 additions and 1129 deletions

View file

@ -1,8 +1,8 @@
- Target invariants, not completion claims
- This crate owns compatibility for all existing Python callbacks and loggers, including `CustomLogger`. `mapping.rs` owns the executable call bindings and the inventory of Python-owned hooks. A Python-owned entry records an existing path, never permission to invoke it a second time. The native call adapter preserves the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on)
- This crate owns compatibility for all existing Python callbacks and loggers, including `CustomLogger`. `facets/` holds `LegacyHooks` on each facet of the call, one file per facet, each listing the Python callbacks its stages reach. `mapping.rs` collects those lists and the inventory of Python-owned hooks. A Python-owned entry records an existing path, never permission to invoke it a second time. The native call adapter preserves the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on)
- Smell test: if a future callback host (`callbacks-v1-python`, WASM, in-process Rust) could share a piece of this crate, it does not belong here
- SDK request policy (credential inheritance, the budget and retry-count limits) is a separate hook supplied by `python-bridge`. Compose it after this adapter so logging adopts the final keyword view before policy mutates or rejects it
- The driver in `litellm-host-python`, the routes and core see one `PythonCallHooks` using the shared `CallEvent`. They never learn which Python objects consume a call
- SDK request policy (credential inheritance, the budget and retry-count limits) is a separate argument group supplied by `python-bridge`
- The driver in `litellm-host-python`, the routes and core see one `PythonGroup`. They never learn which Python objects consume a call
- Every litellm Python internal Rust still borrows is a variant of `LegacyPython`, grouped by subsystem, with its signature pinned in `python_contract.json`
- The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it
- Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython`

View file

@ -9,7 +9,7 @@ use litellm_host::{
interceptors::{RawResponse, RequestContext, WireRequest},
lifecycle::{FailureOrigin, Timing, epoch_seconds},
};
use litellm_host_python::{Awaiter, Error as HostError, from_py, to_py};
use litellm_host_python::{Error as HostError, from_py, to_py};
use parking_lot::Mutex;
use pyo3::{
exceptions::{PyBaseException, PyException},
@ -64,7 +64,7 @@ fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
PyDateTime::from_timestamp(py, epoch_seconds, None).map(|value| value.into_any().unbind())
}
fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {
pub(crate) fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {
!error.is_instance_of::<PyException>(py)
}
@ -76,14 +76,6 @@ pub(crate) enum FailureStage {
Handlers(Option<Py<PyAny>>),
}
/// An awaited failure handler ends the call however it ends, unless it was cancelled.
fn async_failure_result(result: PyResult<Py<PyAny>>) -> PyResult<()> {
match result {
Err(error) if Python::attach(|py| is_cancellation(py, &error)) => Err(error),
_ => Ok(()),
}
}
impl LegacyLogging {
pub fn new(
py: Python<'_>,
@ -158,14 +150,14 @@ impl LegacyLogging {
/// The keyword view the rest of the call reads: a copy, so the deployment hook's own
/// dict is left as the hook returned it, carrying the logger as `@client` injects it.
fn prepare(&mut self, py: Python<'_>) -> PyResult<Py<PyDict>> {
pub(crate) fn prepare(&mut self, py: Python<'_>) -> PyResult<Py<PyDict>> {
let prepared = self.call.kwargs().bind(py).copy()?;
prepared.set_item("litellm_logging_obj", self.logger()?.object(py))?;
self.call.set_kwargs(prepared.unbind());
Ok(self.call.kwargs().clone_ref(py))
}
fn finalize(&mut self, py: Python<'_>) -> PyResult<Py<PyAny>> {
pub(crate) fn finalize(&mut self, py: Python<'_>) -> PyResult<Py<PyAny>> {
finalize(
py,
&self.response,
@ -290,7 +282,7 @@ impl LegacyLogging {
/// The sync failure handler, then the async one for async calls. Ordinary handler
/// errors never replace the selected failure or suppress the other family; a
/// cancellation does end the call.
fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
pub(crate) fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
let (Some(logger), Some(error)) = (&self.logger, &self.error) else {
return Ok(None);
};
@ -314,7 +306,7 @@ impl LegacyLogging {
}
impl LegacyLogging {
fn adopt_hooked_arguments(
pub(crate) fn adopt_hooked_arguments(
&mut self,
py: Python<'_>,
result: PyResult<Py<PyAny>>,
@ -324,7 +316,7 @@ impl LegacyLogging {
Ok(())
}
fn adopt_hooked_response(&mut self, result: PyResult<Py<PyAny>>) -> PyResult<()> {
pub(crate) fn adopt_hooked_response(&mut self, result: PyResult<Py<PyAny>>) -> PyResult<()> {
self.response = Some(result?);
Ok(())
}
@ -332,7 +324,7 @@ impl LegacyLogging {
impl LegacyLogging {
/// Sets the call up, returning the deployment pre-call hook to await when it runs.
fn begin_call(
pub(crate) fn begin_call(
&mut self,
py: Python<'_>,
arguments: Py<PyDict>,
@ -405,7 +397,7 @@ impl LegacyLogging {
/// Takes the public response, returning the deployment post-call hook to await when
/// it runs.
fn respond(
pub(crate) fn respond(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
@ -489,7 +481,7 @@ impl LegacyLogging {
}
}
fn failed(
pub(crate) fn failed(
&mut self,
py: Python<'_>,
timing: Timing,
@ -536,100 +528,16 @@ impl LegacyLogging {
/// [`LegacyLogging`] as the call's hooks. The state is locked only around Python work and
/// never across an await, so the call's handle can traverse it while a hook is suspended.
pub struct LegacyHooks(Mutex<LegacyLogging>);
pub struct LegacyHooks(pub(crate) Mutex<LegacyLogging>);
impl LegacyHooks {
pub fn new(logging: LegacyLogging) -> Self {
Self(Mutex::new(logging))
}
fn with<T>(&self, step: impl FnOnce(&mut LegacyLogging, Python<'_>) -> T) -> T {
pub(crate) fn with<T>(&self, step: impl FnOnce(&mut LegacyLogging, Python<'_>) -> T) -> T {
Python::attach(|py| step(&mut self.0.lock(), py))
}
pub(crate) async fn prepare_call(
&self,
task: &dyn Awaiter,
arguments: Py<PyDict>,
started_at: f64,
) -> PyResult<Py<PyDict>> {
let hook = self.with(|logging, py| logging.begin_call(py, arguments, started_at))?;
if let Some(hook) = hook {
let hooked = task.awaited(hook).await;
self.with(|logging, py| logging.adopt_hooked_arguments(py, hooked))?;
}
self.with(|logging, py| logging.prepare(py))
}
pub(crate) fn adopt_arguments(&self, arguments: &Py<PyDict>) {
self.with(|logging, py| logging.adopt_arguments(py, arguments));
}
pub(crate) fn pre_call(
&self,
wire: Box<WireRequest>,
context: &RequestContext,
) -> PyResult<Box<WireRequest>> {
self.with(|logging, py| logging.pre_call(py, wire, context))
}
pub(crate) async fn transform_public_response(
&self,
task: &dyn Awaiter,
response: Py<PyAny>,
timing: Timing,
) -> PyResult<Py<PyAny>> {
let hook = self.with(|logging, py| logging.respond(py, response, timing))?;
if let Some(hook) = hook {
let hooked = task.awaited(hook).await;
self.with(|logging, _| logging.adopt_hooked_response(hooked))?;
}
self.with(|logging, py| logging.finalize(py))
}
pub(crate) fn result_ready(
&self,
facts: &litellm_host::interceptors::ExecutionFacts,
) -> PyResult<()> {
self.with(|logging, py| logging.result_ready(py, facts))
}
pub(crate) fn post_call(&self, raw: &RawResponse) -> PyResult<()> {
self.with(|logging, py| logging.post_call(py, raw))
}
pub(crate) fn succeeded(&self, timing: Timing, response: &Py<PyAny>) -> PyResult<()> {
self.with(|logging, py| logging.succeeded(py, timing, response))
}
pub(crate) async fn failed(
&self,
task: &dyn Awaiter,
timing: Timing,
origin: FailureOrigin,
error: &Py<PyBaseException>,
) -> PyResult<()> {
let stage = self.with(|logging, py| logging.failed(py, timing, origin, error))?;
let handler = match stage {
FailureStage::DeploymentHook(hook) => {
drop(task.awaited(hook).await);
self.with(|logging, py| logging.dispatch_failure(py))?
}
FailureStage::Handlers(handler) => handler,
};
match handler {
Some(handler) => async_failure_result(task.awaited(handler).await),
None => Ok(()),
}
}
pub(crate) fn stream_opened(&self, head: &Py<PyAny>) -> PyResult<()> {
self.with(|logging, py| logging.stream_opened(py, head))
}
pub(crate) fn stream_chunk(&self, chunk: &Py<PyAny>) -> PyResult<()> {
self.with(|logging, py| logging.stream_chunk(py, chunk))
}
}
impl PythonOwned for LegacyHooks {
@ -680,9 +588,8 @@ impl PythonOwned for LegacyLogging {
mod deployment_hooks_tests {
use std::ffi::CStr;
use litellm_host::hooks::CallHooks;
use litellm_host::hooks::{ArgumentHooks, OutcomeHooks, ResponseHooks};
use litellm_host::lifecycle::{FailureOrigin, Timing};
use litellm_host_python::PythonCallEvent;
use pyo3::exceptions::asyncio::CancelledError;
use pyo3::prelude::*;
use pyo3::types::PyDict;
@ -718,7 +625,7 @@ kwargs = {'logger': logger, 'document': document}
task: &Scripted,
) -> (LegacyHooks, PyResult<Py<PyDict>>) {
let hooks = hooks(legacy_call(py, locals, asynchronous));
let prepared = finish(hooks.prepare_arguments(task, kwargs(locals), 0.0));
let prepared = finish(hooks.prepare(task, kwargs(locals), 0.0));
(hooks, prepared)
}
@ -744,7 +651,7 @@ kwargs = {'logger': logger, 'document': document}
..legacy_call(py, &locals, asynchronous)
});
let task = Scripted::answering([Ok(kwargs(&locals).into_any())]);
finish(hooks.prepare_arguments(&task, kwargs(&locals), 0.0)).unwrap();
finish(hooks.prepare(&task, kwargs(&locals), 0.0)).unwrap();
assert_eq!(task.count() == 1, asynchronous);
locals.set_item("expected", expected).unwrap();
locals.set_item("asynchronous", asynchronous).unwrap();
@ -859,12 +766,9 @@ logger.hooks = {'pre': lambda kwargs: kwargs}
]);
let (hooks, prepared) = begin(py, &locals, true, &task);
prepared.unwrap();
let returned = finish(hooks.transform_response(
&task,
local(&locals, "response").unbind(),
TIMING,
))
.unwrap();
let returned =
finish(hooks.transform(&task, local(&locals, "response").unbind(), TIMING))
.unwrap();
assert_eq!(task.count(), 2);
assert!(returned.bind(py).is(local(&locals, "replacement")));
run(
@ -896,7 +800,7 @@ assert finalized is replacement
let (hooks, prepared) = begin(py, &locals, true, &task);
let error = if post_call {
prepared.unwrap();
finish(hooks.transform_response(&task, local(&locals, "response").unbind(), TIMING))
finish(hooks.transform(&task, local(&locals, "response").unbind(), TIMING))
.unwrap_err()
} else {
prepared.unwrap_err()
@ -934,12 +838,7 @@ assert finalized is replacement
let (hooks, prepared) = begin(py, &locals, true, &task);
prepared.unwrap();
let failure = failure(&locals);
let failed = PythonCallEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Call,
error: &failure,
};
finish(hooks.on_event(&task, failed)).unwrap();
finish(hooks.failed(&task, TIMING, FailureOrigin::Call, &failure)).unwrap();
assert_eq!(task.count(), 3);
run(
py,
@ -1718,9 +1617,9 @@ def check():
mod terminal_tests {
use std::ffi::CStr;
use litellm_host::hooks::CallHooks;
use litellm_host::lifecycle::{FailureOrigin, Timing};
use litellm_host_python::{PythonCallEvent, PythonOwned};
use litellm_host::hooks::{Cancelled, OutcomeHooks, ResponseHooks};
use litellm_host::lifecycle::{CallEvent, FailureOrigin, Timing};
use litellm_host_python::PythonOwned;
use pyo3::exceptions::PyRuntimeError;
use pyo3::exceptions::asyncio::CancelledError;
use pyo3::prelude::*;
@ -1750,28 +1649,15 @@ mod terminal_tests {
fn succeed(locals: &Bound<'_, PyDict>, hooks: &LegacyHooks) -> usize {
let response = local(locals, "response").unbind();
let task = Scripted::default();
finish(hooks.on_event(
&task,
PythonCallEvent::Succeeded {
timing: TIMING,
response: &response,
},
))
.unwrap();
finish(hooks.succeeded(&task, TIMING, &response)).unwrap();
task.count()
}
/// Reports the call's `failure`, answering any awaited handler from `task`.
fn fail(locals: &Bound<'_, PyDict>, hooks: &LegacyHooks, task: &Scripted) -> PyResult<()> {
let failure = failure(locals);
finish(hooks.on_event(
task,
PythonCallEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Host,
error: &failure,
},
))
finish(hooks.failed(task, TIMING, FailureOrigin::Host, &failure))
.map_err(|Cancelled(error)| error)
}
#[rstest]
@ -1821,17 +1707,16 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy
);
let hooks = hooks(logged(py, &locals, true));
let response = local(&locals, "response").unbind();
let event = PythonCallEvent::Succeeded {
timing: TIMING,
response: &response,
};
let (sender, receiver) = litellm_host::observation::observation_channel(
std::num::NonZeroUsize::new(1).unwrap(),
);
drop(receiver);
sender.emit(event.snapshot());
sender.emit(CallEvent::Succeeded {
timing: TIMING,
response: (),
});
let task = Scripted::default();
finish(hooks.on_event(&task, event)).unwrap();
finish(hooks.succeeded(&task, TIMING, &response)).unwrap();
assert_eq!(task.count(), 0);
assert_eq!(sender.dropped_events(), 1);
run(py, &locals, c"
@ -1855,13 +1740,13 @@ assert logger.calls[1][1] is response
..logged(py, &locals, true)
});
hooks
.on_stream_open(&local(&locals, "head").unbind())
.stream_opened(&local(&locals, "head").unbind())
.unwrap();
hooks
.on_stream_chunk(&local(&locals, "first").unbind())
.stream_chunk(&local(&locals, "first").unbind())
.unwrap();
hooks
.on_stream_chunk(&local(&locals, "last").unbind())
.stream_chunk(&local(&locals, "last").unbind())
.unwrap();
assert_eq!(succeed(&locals, &hooks), 0);
run(

View file

@ -0,0 +1,33 @@
use litellm_host::hooks::ArgumentHooks;
use litellm_host_python::{Awaiter, PythonSurface};
use pyo3::{prelude::*, types::PyDict};
use crate::{CallBoundary, LegacyHooks};
pub(crate) const CALLBACKS: &[(CallBoundary, &[&str])] = &[(
CallBoundary::PrepareArguments,
&["async_pre_call_deployment_hook"],
)];
impl ArgumentHooks<PythonSurface> for LegacyHooks {
/// `function_setup`, then the deployment pre-call hook, whose returned kwargs become
/// the keyword view the rest of the call reads.
async fn prepare<'a>(
&'a self,
task: &'a dyn Awaiter,
arguments: Py<PyDict>,
started_at: f64,
) -> PyResult<Py<PyDict>> {
let hook = self.with(|logging, py| logging.begin_call(py, arguments, started_at))?;
if let Some(hook) = hook {
let hooked = task.awaited(hook).await;
self.with(|logging, py| logging.adopt_hooked_arguments(py, hooked))?;
}
self.with(|logging, py| logging.prepare(py))
}
fn prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
self.with(|logging, py| logging.adopt_arguments(py, arguments));
Ok(())
}
}

View file

@ -0,0 +1,13 @@
//! [`LegacyHooks`](crate::LegacyHooks) on each facet of the call, one file per facet, so the
//! stage-to-callback mapping is the file layout. Every facet file also lists the Python
//! callbacks its stages reach, which `mapping.rs` collects.
mod arguments;
mod outcome;
mod provider;
mod response;
pub(crate) use arguments::CALLBACKS as ARGUMENT_CALLBACKS;
pub(crate) use outcome::CALLBACKS as OUTCOME_CALLBACKS;
pub(crate) use provider::CALLBACKS as PROVIDER_CALLBACKS;
pub(crate) use response::CALLBACKS as RESPONSE_CALLBACKS;

View file

@ -0,0 +1,89 @@
use litellm_host::{
hooks::{Cancelled, OutcomeHooks},
lifecycle::{FailureOrigin, Timing},
};
use litellm_host_python::{Awaiter, PythonSurface, report_unless_cancelled};
use pyo3::{exceptions::PyBaseException, prelude::*};
use crate::{CallBoundary, LegacyHooks, adapter::FailureStage, is_cancellation};
pub(crate) const CALLBACKS: &[(CallBoundary, &[&str])] = &[
(
CallBoundary::Succeeded,
&[
"log_success_event",
"async_log_success_event",
"logging_hook",
"async_logging_hook",
"redact_standard_logging_payload_from_model_call_details",
"log_event",
"async_log_event",
],
),
(
CallBoundary::Failed,
&[
"async_post_call_failure_deployment_hook",
"log_failure_event",
"async_log_failure_event",
"log_model_group_rate_limit_error",
"log_event",
"async_log_event",
],
),
];
/// An awaited failure handler ends the call however it ends, unless it was cancelled.
fn async_failure_result(result: PyResult<Py<PyAny>>) -> PyResult<()> {
match result {
Err(error) if Python::attach(|py| is_cancellation(py, &error)) => Err(error),
_ => Ok(()),
}
}
impl LegacyHooks {
async fn fail(
&self,
task: &dyn Awaiter,
timing: Timing,
origin: FailureOrigin,
error: &Py<PyBaseException>,
) -> PyResult<()> {
let stage = self.with(|logging, py| logging.failed(py, timing, origin, error))?;
let handler = match stage {
FailureStage::DeploymentHook(hook) => {
drop(task.awaited(hook).await);
self.with(|logging, py| logging.dispatch_failure(py))?
}
FailureStage::Handlers(handler) => handler,
};
match handler {
Some(handler) => async_failure_result(task.awaited(handler).await),
None => Ok(()),
}
}
}
impl OutcomeHooks<PythonSurface> for LegacyHooks {
/// The sync and async success fan-out with the selected response.
async fn succeeded<'a>(
&'a self,
_: &'a dyn Awaiter,
timing: Timing,
response: &Py<PyAny>,
) -> Result<(), Cancelled<PyErr>> {
report_unless_cancelled(self.with(|logging, py| logging.succeeded(py, timing, response)))
}
/// The deployment failure hook, then the sync and async failure fan-out with the
/// selected exception.
async fn failed<'a>(
&'a self,
task: &'a dyn Awaiter,
timing: Timing,
origin: FailureOrigin,
error: &Py<PyBaseException>,
) -> Result<(), Cancelled<PyErr>> {
report_unless_cancelled(self.fail(task, timing, origin, error).await)
}
}

View file

@ -0,0 +1,46 @@
use litellm_host::{
hooks::ProviderHooks,
interceptors::{ExecutionFacts, RawResponse, RequestContext, WireRequest},
};
use litellm_host_python::{Awaiter, PythonSurface};
use pyo3::prelude::*;
use crate::{CallBoundary, LegacyHooks};
pub(crate) const CALLBACKS: &[(CallBoundary, &[&str])] = &[
(
CallBoundary::BeforeProviderRequest,
&["log_pre_api_call", "log_input_event"],
),
(CallBoundary::AfterProviderResponse, &["log_post_api_call"]),
];
impl ProviderHooks<PythonSurface> for LegacyHooks {
/// `Logging.pre_call` over the wire request, which the callbacks may edit in place.
async fn before_provider_request<'a>(
&'a self,
_: &'a dyn Awaiter,
wire: Box<WireRequest>,
context: &RequestContext,
) -> PyResult<Box<WireRequest>> {
self.with(|logging, py| logging.pre_call(py, wire, context))
}
/// `Logging.post_call` with the raw provider response.
async fn after_provider_response<'a>(
&'a self,
_: &'a dyn Awaiter,
raw: &RawResponse,
) -> PyResult<()> {
self.with(|logging, py| logging.post_call(py, raw))
}
/// The provider, cost and cache facts the spend record is built from.
async fn result_ready<'a>(
&'a self,
_: &'a dyn Awaiter,
facts: &ExecutionFacts,
) -> PyResult<()> {
self.with(|logging, py| logging.result_ready(py, facts))
}
}

View file

@ -0,0 +1,39 @@
use litellm_host::{hooks::ResponseHooks, lifecycle::Timing};
use litellm_host_python::{Awaiter, PythonSurface};
use pyo3::prelude::*;
use crate::{CallBoundary, LegacyHooks};
pub(crate) const CALLBACKS: &[(CallBoundary, &[&str])] = &[
(
CallBoundary::TransformResponse,
&["async_post_call_success_deployment_hook"],
),
(CallBoundary::StreamOpened, &[]),
(CallBoundary::StreamChunk, &[]),
];
impl ResponseHooks<PythonSurface> for LegacyHooks {
/// The deployment success hook, whose returned response is what the caller receives.
async fn transform<'a>(
&'a self,
task: &'a dyn Awaiter,
response: Py<PyAny>,
timing: Timing,
) -> PyResult<Py<PyAny>> {
let hook = self.with(|logging, py| logging.respond(py, response, timing))?;
if let Some(hook) = hook {
let hooked = task.awaited(hook).await;
self.with(|logging, _| logging.adopt_hooked_response(hooked))?;
}
self.with(|logging, py| logging.finalize(py))
}
fn stream_opened(&self, head: &Py<PyAny>) -> PyResult<()> {
self.with(|logging, py| logging.stream_opened(py, head))
}
fn stream_chunk(&self, chunk: &Py<PyAny>) -> PyResult<()> {
self.with(|logging, py| logging.stream_chunk(py, chunk))
}
}

View file

@ -1,8 +1,8 @@
//! The legacy `@client` wrapper as the native call sees it: litellm's `Logging` object, the
//! sync and async callback registries it fans out to, the deployment hooks and the deferred
//! proxy release. All of it sits behind one
//! [`PythonCallHooks`](litellm_host_python::PythonCallHooks), so the driver, the routes and
//! core never learn which Python object is on the other end.
//! proxy release. All of it sits behind one [`LegacyHooks`] group on the Python surface, one
//! facet per file under `facets/`, so the driver, the routes and core never learn which
//! Python object is on the other end.
//!
//! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`]
//! is where those objects live.
@ -11,9 +11,11 @@ mod adapter;
mod call;
mod callbacks;
mod deferred;
mod facets;
mod logger;
mod mapping;
mod python;
pub(crate) use adapter::is_cancellation;
pub use adapter::{LegacyHooks, LegacyLogging};
pub use call::PublicCall;
pub(crate) use callbacks::{LegacyCallbacks, is_internal_call};

View file

@ -1,15 +1,9 @@
use std::future::{Future, ready};
//! Which Python callbacks each neutral call stage reaches, and which callbacks Python
//! still invokes itself.
use futures_util::future::BoxFuture;
use litellm_host::{
hooks::CallHooks,
interceptors::{RawResponse, RequestContext, WireRequest},
lifecycle::{ExecutionEvent, FailureOrigin, Timing},
use crate::facets::{
ARGUMENT_CALLBACKS, OUTCOME_CALLBACKS, PROVIDER_CALLBACKS, RESPONSE_CALLBACKS,
};
use litellm_host_python::{Awaiter, PythonCallEvent, PythonRuntime};
use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
use crate::LegacyHooks;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CallBoundary {
@ -36,194 +30,24 @@ pub struct CallbackMapping {
pub dispatch: Dispatch,
}
struct Binding<H> {
boundary: CallBoundary,
invoke: H,
callbacks: &'static [&'static str],
}
impl<H> Binding<H> {
fn mappings(&self) -> impl Iterator<Item = CallbackMapping> {
self.callbacks.iter().map(|callback| CallbackMapping {
fn call_mappings(
stages: &'static [(CallBoundary, &'static [&'static str])],
) -> impl Iterator<Item = CallbackMapping> {
stages.iter().flat_map(|(boundary, callbacks)| {
callbacks.iter().map(|callback| CallbackMapping {
callback,
dispatch: Dispatch::Call(self.boundary),
dispatch: Dispatch::Call(*boundary),
})
}
})
}
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: |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: LegacyHooks::pre_call,
callbacks: &["log_pre_api_call", "log_input_event"],
};
const AFTER: Binding<After> = Binding {
boundary: CallBoundary::AfterProviderResponse,
invoke: LegacyHooks::post_call,
callbacks: &["log_post_api_call"],
};
const TRANSFORM: Binding<Transform> = Binding {
boundary: CallBoundary::TransformResponse,
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: LegacyHooks::succeeded,
callbacks: &[
"log_success_event",
"async_log_success_event",
"logging_hook",
"async_logging_hook",
"redact_standard_logging_payload_from_model_call_details",
"log_event",
"async_log_event",
],
};
const FAILURE: Binding<Failure> = Binding {
boundary: CallBoundary::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",
"async_log_failure_event",
"log_model_group_rate_limit_error",
"log_event",
"async_log_event",
],
};
const OPEN: Binding<Open> = Binding {
boundary: CallBoundary::StreamOpened,
invoke: LegacyHooks::stream_opened,
callbacks: &[],
};
const CHUNK: Binding<Chunk> = Binding {
boundary: CallBoundary::StreamChunk,
invoke: LegacyHooks::stream_chunk,
callbacks: &[],
};
pub fn callback_mappings() -> impl Iterator<Item = CallbackMapping> {
PREPARE
.mappings()
.chain(BEFORE.mappings())
.chain(AFTER.mappings())
.chain(TRANSFORM.mappings())
.chain(SUCCESS.mappings())
.chain(FAILURE.mappings())
.chain(OPEN.mappings())
.chain(CHUNK.mappings())
call_mappings(ARGUMENT_CALLBACKS)
.chain(call_mappings(PROVIDER_CALLBACKS))
.chain(call_mappings(RESPONSE_CALLBACKS))
.chain(call_mappings(OUTCOME_CALLBACKS))
.chain(PYTHON_CALLBACKS.iter().copied())
}
impl CallHooks<PythonRuntime> for LegacyHooks {
fn prepare_arguments<'a>(
&'a self,
task: &'a dyn Awaiter,
arguments: Py<PyDict>,
started_at: f64,
) -> impl Future<Output = PyResult<Py<PyDict>>> + Send {
(PREPARE.invoke)(self, task, arguments, started_at)
}
fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
self.adopt_arguments(arguments);
Ok(())
}
fn before_provider_request<'a>(
&'a self,
_: &'a dyn Awaiter,
wire: Box<WireRequest>,
context: &RequestContext,
) -> impl Future<Output = PyResult<Box<WireRequest>>> + Send {
ready((BEFORE.invoke)(self, wire, context))
}
fn transform_response<'a>(
&'a self,
task: &'a dyn Awaiter,
response: Py<PyAny>,
timing: Timing,
) -> impl Future<Output = PyResult<Py<PyAny>>> + Send {
(TRANSFORM.invoke)(self, task, response, timing)
}
async fn on_event<'a>(
&'a self,
task: &'a dyn Awaiter,
event: PythonCallEvent<'_>,
) -> PyResult<()> {
match event {
PythonCallEvent::Started { .. } | PythonCallEvent::Cancelled { .. } => Ok(()),
PythonCallEvent::Execution(ExecutionEvent::ResultReady { facts }) => {
self.result_ready(&facts)
}
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
(AFTER.invoke)(self, raw)
}
PythonCallEvent::Succeeded { timing, response } => {
(SUCCESS.invoke)(self, timing, response)
}
PythonCallEvent::Failed {
timing,
origin,
error,
} => (FAILURE.invoke)(self, task, timing, origin, error).await,
}
}
fn on_stream_open(&self, head: &Py<PyAny>) -> PyResult<()> {
(OPEN.invoke)(self, head)
}
fn on_stream_chunk(&self, chunk: &Py<PyAny>) -> PyResult<()> {
(CHUNK.invoke)(self, chunk)
}
}
macro_rules! python_callbacks {
($($dispatch:expr => [$($callback:literal),* $(,)?]),* $(,)?) => {
const PYTHON_CALLBACKS: &[CallbackMapping] = &[

View file

@ -14,19 +14,19 @@ The driver owns ordering: start, argument preparation, prepared-argument hooks,
Boundary tests exercise behavior with a supplied lifecycle binding without importing the LiteLLM Python package. Pin inline awaiting, awaitable final values, exception identity, cancellation, re-entry and release of retained objects, rather than module names or source layout
`HookChain` specializes the shared host chain for Python runtime hooks. The Python adapter owns cancellation guards, close and GC traversal. `PythonRuntime` reports an ordinary terminal notification error as unraisable and lets cancellation stop the chain. Each argument, wire-request and response transformation feeds its result to the next hook. After all argument transformations, the driver calls `arguments_prepared` on every hook in order. Retained callback views must adopt that dictionary before later policy hooks can mutate or reject it. SDK policy is supplied by bridge composition as a hook, never a separate driver phase or parameter. Hooks implement only the stages they need. Default stages preserve the supplied values
`PythonSurface` is the shared `Surface` for Python: `Py<PyDict>` arguments, `Py<PyAny>` responses and chunks, `PyErr` errors, the caller's task as `&dyn Awaiter`. A group implements the four facets (`ArgumentHooks`, `ProviderHooks`, `ResponseHooks`, `OutcomeHooks`) plus `PythonOwned` and is mounted with `PythonSurface::group`, which guards every stage: a cancellation the task resumes a hook with ends that stage there. `Hooks` is the shared chain on this surface, and the bridge composes it; SDK policy is a mounted group, never a driver phase. A group implements only the stages it needs, default stages preserve the supplied values. The driver calls `prepared` on every group after all `prepare` stages, with the final keyword view
Terminal notifications share the selected response or exception. An ordinary notification error is reported as unraisable and does not skip the next hook or replace the selected outcome. Preparation, interception and transformation errors stop the chain. Cancellation stops all further hook dispatch. Suspensions stay inline in the existing driver, and the chain traverses retained event values for GC
Terminal outcomes share the selected response or exception. A Python outcome hook settles its own result through `report_unless_cancelled`: an ordinary exception is written unraisable and does not skip the next group or replace the selected outcome, anything else is `Cancelled` and stops all further dispatch. Preparation, interception and transformation errors stop the chain. Suspensions stay inline in the existing driver
## Existing runtime invariants
- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonBinding`, `PythonHostCalls` and `PythonOwned` traits, and the `PythonRuntime` specialization of `host::hooks::CallHooks`
- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonBinding`, `PythonHostCalls` and `PythonOwned` traits, and the `PythonSurface` specialization of the `host::hooks` facets
- No LiteLLM domain dependencies beyond `litellm-host`: no route types, no `Logging` policy, no public API registration, no cdylib build features
- `PythonCallHooks` only constrains the shared call-stage interface to `PythonRuntime` and Python ownership. It must not redeclare the stages
- `PythonCallEvent` is a specialization of the shared `CallEvent`, never a separately defined lifecycle. The driver emits `Succeeded` or `Failed` exactly once and never dispatches callbacks after a cancellation. Which Python objects consume those events is the legacy adapter's business
- `PythonGroup` only constrains the shared facets to `PythonSurface` and Python ownership. It must not redeclare the stages
- The driver delivers `succeeded` or `failed` exactly once and never dispatches callbacks after a cancellation. Which Python objects consume those outcomes is the legacy adapter's business
- `CallOptions` can publish snapshots independently of callback delivery. Terminal snapshots follow completed hook dispatch, and cancellation never calls a Python callback. For Python-driven calls, leave the machine's observation publisher unset so the driver is the sole publisher of intercepted provider-response snapshots
- `PythonBinding::decode_request` receives the keyword view returned by `prepare_arguments` and updated by `arguments_prepared`, not the caller's dict. A binding that decodes from it inherits all composed argument rewrites
- A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the binding's `map_error`. A Python exception raised inside the call, and a failure in `prepare_arguments` or `transform_response`, is raised as is
- `PythonBinding::decode_request` receives the keyword view returned by `prepare` and seen by `prepared`, not the caller's dict. A binding that decodes from it inherits all composed argument rewrites
- A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the binding's `map_error`. A Python exception raised inside the call, and a failure in `prepare` or `transform`, is raised as is
- A failing `map_error` is raised with the native error's text as its `__context__`, never swallowed
- Use standard PyO3 ownership and conversion APIs
- Prefer `Bound<'py, T>` for attached operations/results, `Py<T>` for retention. Binding/unbinding does not copy payloads

View file

@ -11,7 +11,7 @@ pub trait PythonBinding: PythonOwned {
/// raises it.
type Failure: Into<PyErr>;
/// Decodes the keyword view returned by `prepare_arguments`, including preflight rewrites.
/// Decodes the keyword view returned by `prepare`, including preflight rewrites.
fn decode_request(
&self,
py: Python<'_>,

View file

@ -3,7 +3,7 @@ use std::sync::Arc;
use std::task::{Context, Poll, Waker};
use parking_lot::{Mutex, MutexGuard};
use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError};
use pyo3::exceptions::{PyBaseException, PyRuntimeError};
use pyo3::gc::{PyTraverseError, PyVisit};
use pyo3::prelude::*;
use pyo3::types::PyDict;
@ -11,6 +11,8 @@ use pyo3::types::PyDict;
use litellm_coroutine::{Coroutine, CoroutineState};
use litellm_host::{
call::HostedCompletion,
hooks::Cancelled,
interceptors::RawResponse,
lifecycle::{CallEvent, ExecutionEvent, FailureOrigin, Timing, epoch_seconds},
machine::{HostFailure, Machine, MachineStep},
observation::ObservationSender,
@ -20,7 +22,7 @@ use litellm_host::{
use crate::PythonHostCalls;
use crate::caller::{Answer, PythonCaller, Yielded};
use crate::handle::{Execution, ExecutionBody, ExecutionStep, StepTypes};
use crate::hooks::{PythonCallEvent, PythonCallHooks};
use crate::hooks::{PythonGroup, is_cancellation};
use crate::native::{NativeMachine, NativePoll};
use crate::{Error, InvokeError, PythonBinding};
@ -34,6 +36,8 @@ type StartMachine<P, M> = Box<
+ Sync,
>;
type Body = Coroutine<Yielded, PyResult<Py<PyAny>>>;
/// The call's events as the observers see them, borrowing the Python values the call holds.
type PythonCallEvent<'a> = CallEvent<&'a Py<PyAny>, &'a Py<PyBaseException>, &'a RawResponse>;
pub struct CallOptions {
pub asynchronous: bool,
@ -51,10 +55,6 @@ impl CallOptions {
}
}
fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {
!error.is_instance_of::<PyException>(py)
}
struct Observation {
observers: Option<ObservationSender>,
observing: bool,
@ -148,7 +148,7 @@ where
impl<H, M, L> CallBody<H, M, L>
where
L: PythonCallHooks + 'static,
L: PythonGroup + 'static,
H: PythonBinding + PythonHostCalls<H::Protocol> + 'static,
M: Machine<Protocol = H::Protocol> + 'static,
M::Complete: Into<HostedCompletion<ResponseOf<H>>>,
@ -188,9 +188,9 @@ where
parts.observation.emit(&started);
}
self.hooks()
.on_event(&self.caller, started)
.started(&self.caller, started_at)
.await
.map_err(Setup::Host)?;
.map_err(|Cancelled(error)| Setup::Host(error))?;
let arguments = self
.parts()
.arguments
@ -198,12 +198,10 @@ where
.ok_or_else(|| Setup::Host(Error::MissingState.into()))?;
let prepared = self
.hooks()
.prepare_arguments(&self.caller, arguments, started_at)
.prepare(&self.caller, arguments, started_at)
.await
.map_err(Setup::Host)?;
self.hooks()
.arguments_prepared(&prepared)
.map_err(Setup::Host)?;
self.hooks().prepared(&prepared).map_err(Setup::Host)?;
let decoded = Python::attach(|py| self.binding().decode_request(py, prepared.bind(py)));
self.parts().arguments = Some(prepared);
let request = decoded.map_err(Setup::Decode)?;
@ -289,42 +287,41 @@ where
Ok(())
}
HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => {
self.notify(PythonCallEvent::Execution(ExecutionEvent::ResultReady {
facts,
}))
.await?;
self.parts()
.observation
.emit(&CallEvent::Execution(ExecutionEvent::ResultReady {
facts: facts.clone(),
}));
self.hooks().result_ready(&self.caller, &facts).await?;
reply.send(());
Ok(())
}
HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => {
self.notify(PythonCallEvent::Execution(
self.parts().observation.emit(&CallEvent::Execution(
ExecutionEvent::ProviderResponseReceived { raw: &raw },
))
.await?;
));
self.hooks()
.after_provider_response(&self.caller, &raw)
.await?;
reply.send(());
Ok(())
}
HostRequest::Stream(StreamDelivery::Open(head, reply)) => {
self.streaming = true;
let head = Python::attach(|py| self.binding().encode_stream_head(py, head))?;
self.hooks().on_stream_open(&head)?;
self.hooks().stream_opened(&head)?;
reply.send(self.caller.open_stream(head).await);
Ok(())
}
HostRequest::Stream(StreamDelivery::Chunk(chunk, reply)) => {
let chunk = Python::attach(|py| self.binding().encode_chunk(py, chunk))?;
self.hooks().on_stream_chunk(&chunk)?;
self.hooks().stream_chunk(&chunk)?;
reply.send(self.caller.yield_chunk(chunk).await);
Ok(())
}
}
}
async fn notify(&self, event: PythonCallEvent<'_>) -> PyResult<()> {
self.parts().observation.emit(&event);
self.hooks().on_event(&self.caller, event).await
}
async fn completed(
&mut self,
completion: HostedCompletion<ResponseOf<H>>,
@ -345,11 +342,7 @@ where
return self.succeeded(public).await;
}
let timing = self.parts().observation.timing();
match self
.hooks()
.transform_response(&self.caller, public, timing)
.await
{
match self.hooks().transform(&self.caller, public, timing).await {
Ok(response) => self.succeeded(response).await,
Err(error) => self.failure(error, FailureOrigin::Host).await,
}
@ -379,11 +372,14 @@ where
parts.observation.retain_terminal(&event);
parts.selected = Some(Selected::Response(Arc::clone(&response)));
}
let notified = self.hooks().on_event(&self.caller, event).await;
let notified = self
.hooks()
.succeeded(&self.caller, timing, &response)
.await;
{
let mut parts = self.parts();
parts.selected = None;
notified?;
notified.map_err(|Cancelled(error)| error)?;
parts.observation.finish();
}
Arc::into_inner(response).ok_or_else(|| Error::MissingState.into())
@ -411,11 +407,14 @@ where
parts.observation.retain_terminal(&event);
parts.selected = Some(Selected::Exception(Arc::clone(&exception)));
}
let notified = self.hooks().on_event(&self.caller, event).await;
let notified = self
.hooks()
.failed(&self.caller, timing, origin, &exception)
.await;
{
let mut parts = self.parts();
parts.selected = None;
notified?;
notified.map_err(|Cancelled(error)| error)?;
parts.observation.finish();
}
let exception = Arc::into_inner(exception).ok_or(Error::MissingState)?;
@ -428,7 +427,7 @@ where
struct CallDriver<H, L>
where
H: PythonBinding,
L: PythonCallHooks,
L: PythonGroup,
{
body: Mutex<Body>,
reply: Option<Reply<Answer>>,
@ -438,7 +437,7 @@ where
impl<H, L> CallDriver<H, L>
where
H: PythonBinding + 'static,
L: PythonCallHooks + 'static,
L: PythonGroup + 'static,
{
/// A failure the body raised past its terminal hooks, or a cancellation, is the last
/// thing the observers hear about the call.
@ -469,7 +468,7 @@ where
impl<H, L> ExecutionBody for CallDriver<H, L>
where
H: PythonBinding + 'static,
L: PythonCallHooks + 'static,
L: PythonGroup + 'static,
{
fn resume(&mut self, answer: Option<Answer>) -> PyResult<ExecutionStep> {
match (self.reply.take(), answer) {
@ -521,7 +520,7 @@ where
impl<H, L> Drop for CallDriver<H, L>
where
H: PythonBinding,
L: PythonCallHooks,
L: PythonGroup,
{
fn drop(&mut self) {
self.body.get_mut().cancel();
@ -563,7 +562,7 @@ pub fn run_call<H, M, L>(
options: CallOptions,
) -> PyResult<Py<PyAny>>
where
L: PythonCallHooks + 'static,
L: PythonGroup + 'static,
H: PythonBinding + PythonHostCalls<H::Protocol> + 'static,
M: Machine<Protocol = H::Protocol> + 'static,
M::Complete: Into<HostedCompletion<ResponseOf<H>>>,

View file

@ -1,32 +1,55 @@
mod chain;
pub use chain::HookChain;
mod guard;
use crate::PythonOwned;
use futures_util::future::BoxFuture;
use litellm_host::hooks::{CallHooks, HookRuntime, RuntimeCallEvent};
use pyo3::exceptions::PyBaseException;
use litellm_host::hooks::{AnyGroup, Cancelled, Surface};
use pyo3::exceptions::{PyBaseException, PyException};
use pyo3::prelude::*;
use pyo3::types::PyDict;
use crate::PythonOwned;
/// 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;
/// The Python public surface: keyword arguments in, Python objects out, `PyErr` for every
/// failure, and the caller's task to await on.
pub struct PythonSurface;
impl HookRuntime for PythonRuntime {
type Context<'a> = &'a dyn Awaiter;
impl Surface for PythonSurface {
type Task<'a> = &'a dyn Awaiter;
type Arguments = Py<PyDict>;
type Response = Py<PyAny>;
type Chunk = Py<PyAny>;
type Error = PyErr;
type Failure = Py<PyBaseException>;
type Group = dyn PythonGroup;
}
pub type PythonCallEvent<'a> = RuntimeCallEvent<'a, PythonRuntime>;
/// A group mounted on the Python surface: the four facets plus Python ownership.
pub trait PythonGroup: AnyGroup<PythonSurface> + PythonOwned {}
pub trait PythonCallHooks: CallHooks<PythonRuntime> + PythonOwned {}
impl<H: AnyGroup<PythonSurface> + PythonOwned> PythonGroup for H {}
impl<H: CallHooks<PythonRuntime> + PythonOwned> PythonCallHooks for H {}
pub type Hooks = litellm_host::hooks::Hooks<PythonSurface>;
pub(crate) fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {
!error.is_instance_of::<PyException>(py)
}
/// How a Python outcome hook settles its own result: an ordinary exception is reported as
/// unraisable and the call goes on, anything else is a cancellation that stops it.
pub fn report_unless_cancelled(result: PyResult<()>) -> Result<(), Cancelled<PyErr>> {
let Err(error) = result else {
return Ok(());
};
Python::attach(|py| {
if is_cancellation(py, &error) {
return Err(Cancelled(error));
}
error.write_unraisable(py, None);
Ok(())
})
}

View file

@ -1,171 +0,0 @@
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, ChainRuntime, DynCallHooks, HookGroupAdapter},
interceptors::{RequestContext, WireRequest},
lifecycle::Timing,
};
use parking_lot::Mutex;
use pyo3::{
exceptions::PyException,
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
};
use crate::{Awaiter, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime};
pub type HookChain = litellm_host::hooks::HookChain<PythonRuntime>;
pub trait PythonHookGroup: DynCallHooks<PythonRuntime> + PythonOwned {}
impl<H: DynCallHooks<PythonRuntime> + PythonOwned> PythonHookGroup for H {}
impl ChainRuntime for PythonRuntime {
type Group = dyn PythonHookGroup;
fn terminal_notification_failed(error: PyErr) -> PyResult<()> {
if is_cancellation(&error) {
return Err(error);
}
Python::attach(|py| error.write_unraisable(py, None));
Ok(())
}
}
impl<H: PythonCallHooks + 'static> HookGroupAdapter<H> for PythonRuntime {
fn adapt_group(hooks: H) -> Box<dyn PythonHookGroup> {
Box::new(PythonHookAdapter(hooks))
}
}
struct PythonHookAdapter<H>(H);
impl<H: PythonOwned> PythonOwned for PythonHookAdapter<H> {
fn close(&mut self, py: Python<'_>) {
self.0.close(py);
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.0.traverse(visit)
}
}
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,
}
})
}
}
impl<H: PythonCallHooks> CallHooks<PythonRuntime> for PythonHookAdapter<H> {
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);
stage
.guard(self.0.prepare_arguments(&stage, arguments, started_at))
.await
}
fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
self.0.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);
stage
.guard(self.0.before_provider_request(&stage, wire, context))
.await
}
async fn transform_response<'a>(
&'a self,
task: &'a dyn Awaiter,
response: Py<PyAny>,
timing: Timing,
) -> PyResult<Py<PyAny>> {
let stage = Stage::new(task);
stage
.guard(self.0.transform_response(&stage, response, timing))
.await
}
async fn on_event<'a>(
&'a self,
task: &'a dyn Awaiter,
event: PythonCallEvent<'_>,
) -> PyResult<()> {
let stage = Stage::new(task);
stage.guard(self.0.on_event(&stage, event)).await
}
fn on_stream_open(&self, head: &Py<PyAny>) -> PyResult<()> {
self.0.on_stream_open(head)
}
fn on_stream_chunk(&self, chunk: &Py<PyAny>) -> PyResult<()> {
self.0.on_stream_chunk(chunk)
}
}
impl PythonOwned for HookChain {
fn close(&mut self, py: Python<'_>) {
for hooks in self.groups_mut() {
hooks.close(py);
}
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.groups().try_for_each(|hooks| hooks.traverse(visit))
}
}

View file

@ -0,0 +1,226 @@
use std::future::{Future, pending, poll_fn};
use std::pin::pin;
use std::task::Poll;
use futures_util::future::BoxFuture;
use litellm_host::{
hooks::{ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks},
interceptors::{ExecutionFacts, RawResponse, RequestContext, WireRequest},
lifecycle::{FailureOrigin, Timing},
};
use parking_lot::Mutex;
use pyo3::{
exceptions::PyBaseException,
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
};
use crate::hooks::is_cancellation;
use crate::{Awaiter, Hooks, PythonGroup, PythonOwned, PythonSurface};
impl PythonSurface {
/// Mounts a group on the Python surface. A cancellation the caller's task resumes a
/// hook with ends that hook's stage there, whatever the hook would have done with it.
pub fn group(hooks: impl PythonGroup + 'static) -> Box<dyn PythonGroup> {
Box::new(Guarded(hooks))
}
}
struct Guarded<H>(H);
impl<H: PythonOwned> PythonOwned for Guarded<H> {
fn close(&mut self, py: Python<'_>) {
self.0.close(py);
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.0.traverse(visit)
}
}
/// The caller's task as one stage hands it to its hook.
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, E>(
&self,
hook: impl Future<Output = Result<T, E>>,
cancelled: impl FnOnce(PyErr) -> E,
) -> Result<T, E> {
let mut hook = pin!(hook);
let mut cancelled = Some(cancelled);
poll_fn(|context| match hook.as_mut().poll(context) {
Poll::Ready(result) => Poll::Ready(result),
Poll::Pending => match self.cancelled.lock().take() {
Some(error) => Poll::Ready(Err(cancelled.take().expect("settled once")(error))),
None => Poll::Pending,
},
})
.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 Python::attach(|py| is_cancellation(py, &error)) => {
*self.cancelled.lock() = Some(error);
pending().await
}
answer => answer,
}
})
}
}
impl<H: PythonGroup> ArgumentHooks<PythonSurface> for Guarded<H> {
async fn prepare<'a>(
&'a self,
task: &'a dyn Awaiter,
arguments: Py<PyDict>,
started_at: f64,
) -> PyResult<Py<PyDict>> {
let stage = Stage::new(task);
stage
.guard(self.0.prepare(&stage, arguments, started_at), |error| error)
.await
}
fn prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
self.0.prepared(arguments)
}
}
impl<H: PythonGroup> ProviderHooks<PythonSurface> for Guarded<H> {
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);
stage
.guard(
self.0.before_provider_request(&stage, wire, context),
|error| error,
)
.await
}
async fn after_provider_response<'a>(
&'a self,
task: &'a dyn Awaiter,
raw: &RawResponse,
) -> PyResult<()> {
let stage = Stage::new(task);
stage
.guard(self.0.after_provider_response(&stage, raw), |error| error)
.await
}
async fn result_ready<'a>(
&'a self,
task: &'a dyn Awaiter,
facts: &ExecutionFacts,
) -> PyResult<()> {
let stage = Stage::new(task);
stage
.guard(self.0.result_ready(&stage, facts), |error| error)
.await
}
}
impl<H: PythonGroup> ResponseHooks<PythonSurface> for Guarded<H> {
async fn transform<'a>(
&'a self,
task: &'a dyn Awaiter,
response: Py<PyAny>,
timing: Timing,
) -> PyResult<Py<PyAny>> {
let stage = Stage::new(task);
stage
.guard(self.0.transform(&stage, response, timing), |error| error)
.await
}
fn stream_opened(&self, head: &Py<PyAny>) -> PyResult<()> {
self.0.stream_opened(head)
}
fn stream_chunk(&self, chunk: &Py<PyAny>) -> PyResult<()> {
self.0.stream_chunk(chunk)
}
}
impl<H: PythonGroup> OutcomeHooks<PythonSurface> for Guarded<H> {
async fn started<'a>(
&'a self,
task: &'a dyn Awaiter,
start_time: f64,
) -> Result<(), Cancelled<PyErr>> {
let stage = Stage::new(task);
stage
.guard(self.0.started(&stage, start_time), Cancelled)
.await
}
async fn succeeded<'a>(
&'a self,
task: &'a dyn Awaiter,
timing: Timing,
response: &Py<PyAny>,
) -> Result<(), Cancelled<PyErr>> {
let stage = Stage::new(task);
stage
.guard(self.0.succeeded(&stage, timing, response), Cancelled)
.await
}
async fn failed<'a>(
&'a self,
task: &'a dyn Awaiter,
timing: Timing,
origin: FailureOrigin,
error: &Py<PyBaseException>,
) -> Result<(), Cancelled<PyErr>> {
let stage = Stage::new(task);
stage
.guard(self.0.failed(&stage, timing, origin, error), Cancelled)
.await
}
async fn cancelled<'a>(
&'a self,
task: &'a dyn Awaiter,
timing: Timing,
) -> Result<(), Cancelled<PyErr>> {
let stage = Stage::new(task);
stage
.guard(self.0.cancelled(&stage, timing), Cancelled)
.await
}
}
impl PythonOwned for Hooks {
fn close(&mut self, py: Python<'_>) {
for group in self.groups_mut() {
group.close(py);
}
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.groups().try_for_each(|group| group.traverse(visit))
}
}

View file

@ -1,6 +1,6 @@
//! The CPython runtime adapter: value marshalling, interpreter detachment, the tokio and
//! asyncio glue, and the driver that runs a native [`Machine`](litellm_host::machine::Machine)
//! against a Python binding, host services and active call hooks. Everything here is Python-specific by
//! against a Python binding, host services and the mounted hook groups. Everything here is Python-specific by
//! construction; another host language gets its own crate of the same shape.
mod argument;
@ -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::{Awaiter, HookChain, PythonCallEvent, PythonCallHooks, PythonRuntime};
pub use hooks::{Awaiter, Hooks, PythonGroup, PythonSurface, report_unless_cancelled};
pub use marshal::{
Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py,
};

View file

@ -5,12 +5,12 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use litellm_host::{
hooks::CallHooks,
hooks::{ArgumentHooks, OutcomeHooks, ProviderHooks, ResponseHooks},
lifecycle::{CallEvent, ExecutionEvent, FailureOrigin},
machine::CallMachine,
observation::observation_channel,
};
use litellm_host_python::{CallOptions, HookChain, PythonOwned, PythonRuntime, run_call};
use litellm_host_python::{CallOptions, Hooks, PythonOwned, PythonSurface, run_call};
use pyo3::{
exceptions::{
PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError,
@ -79,10 +79,19 @@ fn composed_hooks_share_prepared_arguments_and_one_execution(
) {
Python::attach(|py| {
let log = Log::default();
let chain = HookChain::new()
.with(hooks(&log, HookScript::ObserveArguments))
.with(hooks(&log, HookScript::RewriteArguments))
.with(hooks(&log, HookScript::ReplaceResponse));
let chain = Hooks::new()
.with(PythonSurface::group(hooks(
&log,
HookScript::ObserveArguments,
)))
.with(PythonSurface::group(hooks(
&log,
HookScript::RewriteArguments,
)))
.with(PythonSurface::group(hooks(
&log,
HookScript::ReplaceResponse,
)));
let result = handed(
py,
run_call(
@ -183,7 +192,6 @@ fn observation_delivery_preserves_callback_order_and_response(
#[case::preparation(HookScript::FailBegin, false, 1)]
#[case::transformation(HookScript::FailAfterSuccess, false, 1)]
#[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)]
#[case::terminal_failure(HookScript::FailTerminal, false, 0)]
fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks(
#[from(interpreter)] _python: Interpreter,
#[case] script: HookScript,
@ -245,6 +253,51 @@ fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks(
});
}
#[rstest]
fn an_ordinary_terminal_hook_error_is_reported_and_the_call_still_succeeds(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let unraisable = PyDict::new(py);
py.run(
c"import sys\nseen = []\nsys.unraisablehook = lambda event: seen.append(event.exc_value)",
Some(&unraisable),
Some(&unraisable),
)
.unwrap();
let (options, mut events) = observed(4);
let log = Log::default();
let (result, log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::FailTerminal),
std::convert::identity,
CallOptions {
asynchronous,
..options
},
);
assert_eq!(
result.unwrap().extract::<String>(py).unwrap(),
"project:1|sign|rewritten"
);
assert!(matches!(events().last(), Some(CallEvent::Succeeded { .. })));
assert!(!log.iter().any(|entry| entry.starts_with("failed:")));
let seen: Vec<String> = unraisable
.get_item("seen")
.unwrap()
.unwrap()
.extract::<Vec<Bound<'_, PyAny>>>()
.unwrap()
.iter()
.map(|error| error.to_string())
.collect();
assert_eq!(seen, ["callback failed"]);
});
}
#[rstest]
#[case::preparation(HookScript::FailBegin, OpScript::Answer)]
#[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)]
@ -534,8 +587,8 @@ enum ArgumentPolicy {
Reject(Py<PyBaseException>),
}
impl CallHooks<PythonRuntime> for ArgumentPolicy {
fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
impl ArgumentHooks<PythonSurface> for ArgumentPolicy {
fn 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())),
@ -543,6 +596,10 @@ impl CallHooks<PythonRuntime> for ArgumentPolicy {
}
}
impl ProviderHooks<PythonSurface> for ArgumentPolicy {}
impl ResponseHooks<PythonSurface> for ArgumentPolicy {}
impl OutcomeHooks<PythonSurface> for ArgumentPolicy {}
impl PythonOwned for ArgumentPolicy {
fn close(&mut self, _: Python<'_>) {}
@ -581,9 +638,12 @@ fn resource_setup_uses_prepared_arguments_and_failures_are_terminal(
Ok(success_machine()(request))
},
binding(&log, OpScript::Answer),
HookChain::new()
.with(hooks(&log, HookScript::RewriteArguments))
.with(ArgumentPolicy::Inherit),
Hooks::new()
.with(PythonSurface::group(hooks(
&log,
HookScript::RewriteArguments,
)))
.with(PythonSurface::group(ArgumentPolicy::Inherit)),
arguments.clone().unbind(),
call_options(asynchronous),
),
@ -675,9 +735,11 @@ fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_start
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
|hooks| {
HookChain::new()
.with(hooks)
.with(ArgumentPolicy::Reject(raised.clone_ref(py)))
Hooks::new()
.with(PythonSurface::group(hooks))
.with(PythonSurface::group(ArgumentPolicy::Reject(
raised.clone_ref(py),
)))
},
call_options(asynchronous),
);
@ -707,7 +769,11 @@ fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote(
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
|hooks| HookChain::new().with(hooks).with(ArgumentPolicy::Inherit),
|hooks| {
Hooks::new()
.with(PythonSurface::group(hooks))
.with(PythonSurface::group(ArgumentPolicy::Inherit))
},
call_options(asynchronous),
);
assert_eq!(

View file

@ -1,9 +1,12 @@
mod support;
use litellm_host::{hooks::CallHooks, lifecycle::Timing};
use litellm_host::{
hooks::{ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks},
lifecycle::{FailureOrigin, Timing},
};
use litellm_host_python::{
Awaiter, Execution, ExecutionBody, ExecutionStep, PythonCallEvent, PythonOwned, PythonRuntime,
Suspension, run_call,
Awaiter, Execution, ExecutionBody, ExecutionStep, PythonOwned, PythonSurface, Suspension,
report_unless_cancelled, run_call,
};
use parking_lot::Mutex;
use pyo3::{
@ -120,8 +123,22 @@ struct PausingHooks {
holder: Holder,
}
impl CallHooks<PythonRuntime> for PausingHooks {
async fn transform_response<'a>(
impl PausingHooks {
/// Pauses the call on the test's awaitable, once, at the stage the holder names.
async fn pause_at(&self, task: &dyn Awaiter, here: bool) -> Result<(), Cancelled<PyErr>> {
let pause = self.pause.lock().take_if(|_| here);
match pause {
Some(pause) => report_unless_cancelled(task.awaited(pause).await.map(drop)),
None => Ok(()),
}
}
}
impl ArgumentHooks<PythonSurface> for PausingHooks {}
impl ProviderHooks<PythonSurface> for PausingHooks {}
impl ResponseHooks<PythonSurface> for PausingHooks {
async fn transform<'a>(
&'a self,
_: &'a dyn Awaiter,
response: Py<PyAny>,
@ -136,22 +153,33 @@ impl CallHooks<PythonRuntime> for PausingHooks {
})),
}
}
}
async fn on_event<'a>(
impl OutcomeHooks<PythonSurface> for PausingHooks {
async fn started<'a>(&'a self, task: &'a dyn Awaiter, _: f64) -> Result<(), Cancelled<PyErr>> {
self.pause_at(task, matches!(self.holder, Holder::Hooks))
.await
}
async fn succeeded<'a>(
&'a self,
task: &'a dyn Awaiter,
event: PythonCallEvent<'_>,
) -> PyResult<()> {
let pauses_here = match self.holder {
Holder::Hooks => true,
Holder::Response => matches!(event, PythonCallEvent::Succeeded { .. }),
Holder::Exception => matches!(event, PythonCallEvent::Failed { .. }),
};
let pause = self.pause.lock().take_if(|_| pauses_here);
if let Some(pause) = pause {
task.awaited(pause).await?;
}
Ok(())
_: Timing,
_: &Py<PyAny>,
) -> Result<(), Cancelled<PyErr>> {
self.pause_at(task, matches!(self.holder, Holder::Response))
.await
}
async fn failed<'a>(
&'a self,
task: &'a dyn Awaiter,
_: Timing,
_: FailureOrigin,
_: &Py<PyBaseException>,
) -> Result<(), Cancelled<PyErr>> {
self.pause_at(task, matches!(self.holder, Holder::Exception))
.await
}
}

View file

@ -4,14 +4,15 @@ use std::future::{Future, ready};
use futures_util::{FutureExt, future::BoxFuture};
use litellm_host::{
hooks::CallHooks,
hooks::{ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks},
interceptors::{RequestContext, WireRequest},
lifecycle::{FailureOrigin, Timing},
};
use litellm_host_python::{
Awaiter, HookChain, PythonCallEvent, PythonOwned, PythonRuntime, run_call,
Awaiter, Hooks, PythonGroup, PythonOwned, PythonSurface, report_unless_cancelled, run_call,
};
use pyo3::{
exceptions::PyBaseException,
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDict, PyTuple},
@ -59,8 +60,28 @@ impl PythonOwned for ScriptHooks {
}
}
impl CallHooks<PythonRuntime> for ScriptHooks {
async fn prepare_arguments<'a>(
impl ScriptHooks {
async fn event(
&self,
task: &dyn Awaiter,
name: &str,
value: impl FnOnce(Python<'_>) -> Py<PyAny>,
) -> Result<(), Cancelled<PyErr>> {
let result = self
.invoke(task, "event", |py| {
Ok(
PyTuple::new(py, [name.into_pyobject(py)?.into_any().unbind(), value(py)])?
.into_any()
.unbind(),
)
})
.await;
report_unless_cancelled(result.map(drop))
}
}
impl ArgumentHooks<PythonSurface> for ScriptHooks {
async fn prepare<'a>(
&'a self,
task: &'a dyn Awaiter,
arguments: Py<PyDict>,
@ -72,11 +93,13 @@ impl CallHooks<PythonRuntime> for ScriptHooks {
Python::attach(|py| Ok(value.into_bound(py).cast_into::<PyDict>()?.unbind()))
}
fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
fn prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
Python::attach(|py| self.object.call_method1(py, "adopt", (arguments,)))?;
Ok(())
}
}
impl ProviderHooks<PythonSurface> for ScriptHooks {
async fn before_provider_request<'a>(
&'a self,
task: &'a dyn Awaiter,
@ -94,8 +117,10 @@ impl CallHooks<PythonRuntime> for ScriptHooks {
..*wire
}))
}
}
async fn transform_response<'a>(
impl ResponseHooks<PythonSurface> for ScriptHooks {
async fn transform<'a>(
&'a self,
task: &'a dyn Awaiter,
response: Py<PyAny>,
@ -104,42 +129,44 @@ impl CallHooks<PythonRuntime> for ScriptHooks {
self.invoke(task, "transform", |_| Ok(response)).await
}
async fn on_event<'a>(
&'a self,
task: &'a dyn Awaiter,
event: PythonCallEvent<'_>,
) -> 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(&self, _head: &Py<PyAny>) -> PyResult<()> {
fn stream_opened(&self, _head: &Py<PyAny>) -> PyResult<()> {
Python::attach(|py| self.object.call_method1(py, "stream", (py.None(),)))?;
Ok(())
}
fn on_stream_chunk(&self, chunk: &Py<PyAny>) -> PyResult<()> {
fn stream_chunk(&self, chunk: &Py<PyAny>) -> PyResult<()> {
Python::attach(|py| self.object.call_method1(py, "stream", (chunk,)))?;
Ok(())
}
}
impl OutcomeHooks<PythonSurface> for ScriptHooks {
async fn started<'a>(&'a self, task: &'a dyn Awaiter, _: f64) -> Result<(), Cancelled<PyErr>> {
self.event(task, "started", |py| py.None()).await
}
async fn succeeded<'a>(
&'a self,
task: &'a dyn Awaiter,
_: Timing,
response: &Py<PyAny>,
) -> Result<(), Cancelled<PyErr>> {
self.event(task, "success", |py| response.clone_ref(py))
.await
}
async fn failed<'a>(
&'a self,
task: &'a dyn Awaiter,
_: Timing,
_: FailureOrigin,
error: &Py<PyBaseException>,
) -> Result<(), Cancelled<PyErr>> {
self.event(task, "failure", |py| error.clone_ref(py).into_any())
.await
}
}
/// Runs every awaitable a hook hands it to completion on a fresh event loop.
struct RunToCompletion;
@ -154,7 +181,7 @@ impl Awaiter for RunToCompletion {
}
}
fn finish<T>(hook: impl Future<Output = PyResult<T>>) -> PyResult<T> {
fn finish<T>(hook: impl Future<Output = T>) -> T {
hook.now_or_never()
.expect("every awaitable runs to completion inline")
}
@ -209,12 +236,33 @@ fourth = Hooks('d')
})
}
fn chain(py: Python<'_>, scripts: &Py<PyDict>, asynchronous: bool) -> HookChain {
let hook = |name| ScriptHooks {
fn script_hooks(
py: Python<'_>,
scripts: &Py<PyDict>,
name: &str,
asynchronous: bool,
) -> ScriptHooks {
ScriptHooks {
object: scripts.bind(py).get_item(name).unwrap().unwrap().unbind(),
asynchronous,
};
HookChain::new().with(hook("first")).with(hook("second"))
}
}
fn chain(py: Python<'_>, scripts: &Py<PyDict>, asynchronous: bool) -> Hooks {
["first", "second"]
.into_iter()
.fold(Hooks::new(), |hooks, name| {
hooks.with(PythonSurface::group(script_hooks(
py,
scripts,
name,
asynchronous,
)))
})
}
fn mounted(hooks: Hooks, group: impl PythonGroup + 'static) -> Hooks {
hooks.with(PythonSurface::group(group))
}
const TIMING: Timing = Timing {
@ -224,12 +272,16 @@ const TIMING: Timing = Timing {
struct PreparedPolicy;
impl CallHooks<PythonRuntime> for PreparedPolicy {
fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
impl ArgumentHooks<PythonSurface> for PreparedPolicy {
fn prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
Python::attach(|py| arguments.bind(py).set_item("policy", "configured"))
}
}
impl ProviderHooks<PythonSurface> for PreparedPolicy {}
impl ResponseHooks<PythonSurface> for PreparedPolicy {}
impl OutcomeHooks<PythonSurface> for PreparedPolicy {}
impl PythonOwned for PreparedPolicy {
fn close(&mut self, _: Python<'_>) {}
@ -241,16 +293,16 @@ impl PythonOwned for PreparedPolicy {
#[rstest]
#[case::sync(false)]
#[case::suspended(true)]
fn transformations_feed_each_other_and_notifications_share_final_values(
fn request_stages_feed_forward_and_response_stages_unwind_over_the_final_values(
scripts: Py<PyDict>,
#[case] asynchronous: bool,
) {
Python::attach(|py| {
let hooks = chain(py, &scripts, asynchronous).with(PreparedPolicy);
let hooks = mounted(chain(py, &scripts, asynchronous), PreparedPolicy);
let original = PyDict::new(py).unbind();
let arguments =
finish(hooks.prepare_arguments(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap();
hooks.arguments_prepared(&arguments).unwrap();
finish(hooks.prepare(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap();
hooks.prepared(&arguments).unwrap();
let context = RequestContext {
model: "model".into(),
custom_llm_provider: "provider".into(),
@ -267,18 +319,10 @@ fn transformations_feed_each_other_and_notifications_share_final_values(
assert_eq!(wire.url, "urlab");
let response = PyDict::new(py).unbind().into_any();
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();
finish(hooks.transform(&RunToCompletion, response.clone_ref(py), TIMING)).unwrap();
finish(hooks.succeeded(&RunToCompletion, TIMING, &final_response)).unwrap();
hooks.stream_opened(&py.None()).unwrap();
hooks.stream_chunk(&response).unwrap();
let locals = scripts.bind(py);
locals.set_item("arguments", arguments).unwrap();
locals.set_item("original", original).unwrap();
@ -291,10 +335,10 @@ assert original == {}
assert first.adopted is arguments and second.adopted is arguments
assert first.adopted['policy'] == second.adopted['policy'] == 'configured'
assert first.model == second.model == 'model'
assert final_response == ((response, 'a'), 'b')
assert final_response == ((response, 'b'), 'a')
assert [(name, kind) for name, kind, value in log] == [
('a', 'prepare'), ('b', 'prepare'), ('a', 'success'), ('b', 'success'),
('a', 'stream'), ('b', 'stream'), ('a', 'stream'), ('b', 'stream'),
('a', 'prepare'), ('b', 'prepare'), ('b', 'success'), ('a', 'success'),
('b', 'stream'), ('a', 'stream'), ('b', 'stream'), ('a', 'stream'),
]
assert log[2][2] is final_response and log[3][2] is final_response
assert log[6][2] is response and log[7][2] is response
@ -318,15 +362,10 @@ fn terminal_failure_does_not_skip_later_hooks(
#[values("first", "second")] failing_hook: &str,
) {
Python::attach(|py| {
let hooks = chain(py, &scripts, asynchronous).with(ScriptHooks {
object: scripts
.bind(py)
.get_item("third")
.unwrap()
.unwrap()
.unbind(),
asynchronous,
});
let hooks = mounted(
chain(py, &scripts, asynchronous),
script_hooks(py, &scripts, "third", asynchronous),
);
let locals = scripts.bind(py);
locals
.get_item(failing_hook)
@ -339,25 +378,17 @@ fn terminal_failure_does_not_skip_later_hooks(
.unwrap();
let error = pyo3::exceptions::PyValueError::new_err("provider failed").into_value(py);
let response = py.None();
let event = if failed {
PythonCallEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Call,
error: &error,
}
if failed {
finish(hooks.failed(&RunToCompletion, TIMING, FailureOrigin::Call, &error)).unwrap();
} else {
PythonCallEvent::Succeeded {
timing: TIMING,
response: &response,
}
};
finish(hooks.on_event(&RunToCompletion, event)).unwrap();
finish(hooks.succeeded(&RunToCompletion, TIMING, &response)).unwrap();
}
locals
.set_item("selected", if failed { error.into_any() } else { response })
.unwrap();
py.run(
c"
assert [name for name, kind, value in log] == ['a', 'b', 'c']
assert [name for name, kind, value in log] == ['c', 'b', 'a']
assert all(kind == log[0][1] for name, kind, value in log)
assert all(value is selected for name, kind, value in log)
",
@ -381,8 +412,7 @@ fn transformation_failure_stops_the_chain(scripts: Py<PyDict>, #[case] asynchron
Some(locals),
)
.unwrap();
let result =
finish(hooks.prepare_arguments(&RunToCompletion, PyDict::new(py).unbind(), 0.0));
let result = finish(hooks.prepare(&RunToCompletion, PyDict::new(py).unbind(), 0.0));
assert!(
result.unwrap_err().value(py).is(locals
.get_item("first")
@ -408,26 +438,17 @@ fn cancellation_stops_notification_dispatch(scripts: Py<PyDict>, #[case] asynchr
let hooks = chain(py, &scripts, asynchronous);
let locals = scripts.bind(py);
py.run(
c"first.error = asyncio.CancelledError()",
c"second.error = asyncio.CancelledError()",
Some(locals),
Some(locals),
)
.unwrap();
let response = py.None();
let result = finish(hooks.on_event(
&RunToCompletion,
PythonCallEvent::Succeeded {
timing: TIMING,
response: &response,
},
));
assert!(
result
.unwrap_err()
.is_instance_of::<pyo3::exceptions::asyncio::CancelledError>(py)
);
let Cancelled(error) =
finish(hooks.succeeded(&RunToCompletion, TIMING, &response)).unwrap_err();
assert!(error.is_instance_of::<pyo3::exceptions::asyncio::CancelledError>(py));
py.run(
c"assert len(log) == 1 and log[0][0] == 'a'",
c"assert len(log) == 1 and log[0][0] == 'b'",
Some(locals),
Some(locals),
)
@ -504,7 +525,7 @@ asyncio.run(exercise())
#[case::pair(2, "ab")]
#[case::three(3, "abc")]
#[case::four(4, "abcd")]
fn builder_runs_hooks_in_append_order(
fn argument_groups_run_in_mount_order(
scripts: Py<PyDict>,
#[case] count: usize,
#[case] expected: &str,
@ -514,15 +535,11 @@ fn builder_runs_hooks_in_append_order(
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,
})
.fold(Hooks::new(), |chain, name| {
mounted(chain, script_hooks(py, &scripts, name, asynchronous))
});
let original = PyDict::new(py).unbind();
let result =
finish(hooks.prepare_arguments(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap();
let result = finish(hooks.prepare(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap();
if count == 0 {
assert!(result.bind(py).is(original.bind(py)));
} else {

View file

@ -4,18 +4,18 @@ use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use litellm_host::{
call::{CallOutput, HostedMachine, hosted_call},
hooks::CallHooks,
interceptors::{Interceptors, RawResponse, RequestContext, WireRequest},
lifecycle::{ExecutionEvent, Timing},
hooks::{ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks},
interceptors::{ExecutionFacts, Interceptors, RawResponse, RequestContext, WireRequest},
lifecycle::{FailureOrigin, Timing},
machine::{CallMachine, MachineFault},
protocol::{Protocol, Reply},
};
use litellm_host_python::{
Awaiter, CallOptions, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks,
PythonHostCalls, PythonOwned, PythonRuntime, run_call,
Awaiter, CallOptions, InvokeError, PythonBinding, PythonGroup, PythonHostCalls, PythonOwned,
PythonSurface, report_unless_cancelled, run_call,
};
use pyo3::{
exceptions::{PyTypeError, PyValueError, asyncio::CancelledError},
exceptions::{PyBaseException, PyTypeError, PyValueError, asyncio::CancelledError},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDict, PyString},
@ -313,8 +313,27 @@ pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks {
}
}
impl CallHooks<PythonRuntime> for SyntheticHooks {
async fn prepare_arguments<'a>(
impl SyntheticHooks {
/// What a terminal hook does with the outcome it was handed: cancels the call, raises
/// an ordinary error that is only reported, or records it.
fn outcome(&self, entry: String) -> Result<(), Cancelled<PyErr>> {
match self.script {
HookScript::CancelTerminal => {
Err(Cancelled(CancelledError::new_err("callback cancelled")))
}
HookScript::FailTerminal => {
report_unless_cancelled(Err(PyValueError::new_err("callback failed")))
}
_ => {
self.log.push(entry);
Ok(())
}
}
}
}
impl ArgumentHooks<PythonSurface> for SyntheticHooks {
async fn prepare<'a>(
&'a self,
_: &'a dyn Awaiter,
arguments: Py<PyDict>,
@ -336,7 +355,7 @@ impl CallHooks<PythonRuntime> for SyntheticHooks {
Ok(arguments)
}
fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
fn prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
if matches!(self.script, HookScript::ObserveArguments) {
let value: String =
Python::attach(|py| arguments.bind(py).get_item("prepared")?.unwrap().extract())?;
@ -344,7 +363,9 @@ impl CallHooks<PythonRuntime> for SyntheticHooks {
}
Ok(())
}
}
impl ProviderHooks<PythonSurface> for SyntheticHooks {
async fn before_provider_request<'a>(
&'a self,
_: &'a dyn Awaiter,
@ -358,7 +379,23 @@ impl CallHooks<PythonRuntime> for SyntheticHooks {
}))
}
async fn transform_response<'a>(
async fn after_provider_response<'a>(
&'a self,
_: &'a dyn Awaiter,
raw: &RawResponse,
) -> PyResult<()> {
self.log.push(format!("response:{}", raw.body));
Ok(())
}
async fn result_ready<'a>(&'a self, _: &'a dyn Awaiter, _: &ExecutionFacts) -> PyResult<()> {
self.log.push("cache_hit");
Ok(())
}
}
impl ResponseHooks<PythonSurface> for SyntheticHooks {
async fn transform<'a>(
&'a self,
_: &'a dyn Awaiter,
response: Py<PyAny>,
@ -374,49 +411,56 @@ impl CallHooks<PythonRuntime> for SyntheticHooks {
}
}
async fn on_event<'a>(
&'a self,
_: &'a dyn Awaiter,
event: PythonCallEvent<'_>,
) -> PyResult<()> {
let terminal = matches!(
event,
PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. }
);
if terminal && matches!(self.script, HookScript::CancelTerminal) {
return Err(CancelledError::new_err("callback cancelled"));
}
if terminal && matches!(self.script, HookScript::FailTerminal) {
return Err(PyValueError::new_err("callback failed"));
}
self.log.push(Python::attach(|py| match event {
PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "cache_hit".into(),
PythonCallEvent::Started { .. } => "started".into(),
PythonCallEvent::Cancelled { .. } => "cancelled".into(),
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
format!("response:{}", raw.body)
}
PythonCallEvent::Succeeded { response, .. } => {
format!("succeeded:{}", response.bind(py))
}
PythonCallEvent::Failed { origin, error, .. } => {
format!("failed:{origin:?}:{}", error.bind(py))
}
}));
Ok(())
}
fn on_stream_open(&self, _: &Py<PyAny>) -> PyResult<()> {
fn stream_opened(&self, _: &Py<PyAny>) -> PyResult<()> {
self.log.push("opened");
Ok(())
}
fn on_stream_chunk(&self, _: &Py<PyAny>) -> PyResult<()> {
fn stream_chunk(&self, _: &Py<PyAny>) -> PyResult<()> {
self.log.push("delivered");
Ok(())
}
}
impl OutcomeHooks<PythonSurface> for SyntheticHooks {
async fn started<'a>(&'a self, _: &'a dyn Awaiter, _: f64) -> Result<(), Cancelled<PyErr>> {
self.log.push("started");
Ok(())
}
async fn succeeded<'a>(
&'a self,
_: &'a dyn Awaiter,
_: Timing,
response: &Py<PyAny>,
) -> Result<(), Cancelled<PyErr>> {
self.outcome(Python::attach(|py| {
format!("succeeded:{}", response.bind(py))
}))
}
async fn failed<'a>(
&'a self,
_: &'a dyn Awaiter,
_: Timing,
origin: FailureOrigin,
error: &Py<PyBaseException>,
) -> Result<(), Cancelled<PyErr>> {
self.outcome(Python::attach(|py| {
format!("failed:{origin:?}:{}", error.bind(py))
}))
}
async fn cancelled<'a>(
&'a self,
_: &'a dyn Awaiter,
_: Timing,
) -> Result<(), Cancelled<PyErr>> {
self.log.push("cancelled");
Ok(())
}
}
impl PythonOwned for SyntheticHooks {
fn close(&mut self, _: Python<'_>) {
self.log.push("adapter.close");
@ -495,7 +539,7 @@ pub fn run_scripted(
)
}
pub fn run_composed<L: PythonCallHooks + 'static>(
pub fn run_composed<L: PythonGroup + 'static>(
py: Python<'_>,
machine: MachineFactory,
binding: SyntheticBinding,

View file

@ -4,7 +4,7 @@
| --- | --- | --- | --- |
| Input and output conversion | `Protocol::{Request, Response, StreamHead, Chunk, Error}` | Typed input. `ResponseEncoder` and `StreamEncoder` produce HTTP values | `PythonBinding` decodes prepared arguments, encodes public values and maps native errors |
| Host services | `Protocol::HostCall`, `HostServices::call` | `host-native::services::HostCallHandler` answers typed calls. `()` handles protocols without host calls | `PythonHostCalls` invokes retained Python objects. It may share an owner with the binding |
| Active hooks | `Interceptors::{before_provider_request, after_provider_response}` and `hooks::CallHooks<Runtime>` | Request interception and fallible execution callbacks | `PythonCallHooks` also prepares arguments, transforms public responses and receives stream callbacks |
| Active hooks | `Interceptors::{before_provider_request, after_provider_response}` and the `hooks` facets (`ArgumentHooks`, `ProviderHooks`, `ResponseHooks`, `OutcomeHooks`) over a `Surface` | Request interception and fallible execution callbacks | `PythonGroup` implements the four facets on `PythonSurface`, mounted through `Hooks` |
| Passive observation | `observation::ObservationSender` | Queued execution and lifecycle snapshots, retained by the response body | Public Python callbacks remain active hooks with their existing failure policy |
| Runtime driving | `Machine`, `HostRequest::{HostCall, Intercept, Stream}` | Body polling controls demand | Native polling and inline caller-task Python awaits control progress |
@ -16,7 +16,7 @@ Each driver owns terminal dispatch. Hooks can change or fail execution. Passive
`interceptors.rs` owns `Interceptors` and its request/response payload types. `lifecycle.rs` owns `CallObserver`, `CallEvent`, `ExecutionEvent`, timing, failure origin, and the observation wrappers. Event payloads are generic so a runtime can retain its own response, exception and raw-response references without introducing a language dependency. `snapshot()` projects them into the owned observation contract without retaining runtime objects. Pass interceptors and observers separately at direct route and HTTP entrypoints. Routes publish execution events independently of interception
`hooks.rs` owns the async call-stage interface and its runtime-associated types. `hooks::HookChain` composes stages in order, using `DynCallHooks` for boxed storage and `CallHooks` for callback authors. A terminal notification reaches every group, and `ChainRuntime::terminal_notification_failed` decides whether a group's error is reported or stops the chain. A runtime supplies its context through `HookRuntime` and its stored hook type through `ChainRuntime`. Python ownership and cancellation guards stay in the Python adapter
`hooks.rs` owns the call's facets, one trait per payload, each with its own failure policy in its return type: `ArgumentHooks` and `ProviderHooks` fold a value through the groups in mount order and stop on the first error, `ResponseHooks` unwinds in reverse mount order, `OutcomeHooks` reaches every group in reverse and only a `Cancelled` group stops the rest. What a group does with an ordinary outcome error is the group's business. `Surface` names a host's public values (`Task`, `Arguments`, `Response`, `Chunk`, `Error`, `Failure`) and how a group is stored (`Group`). `hooks::Hooks<S>` composes mounted groups and implements every facet; `AnyGroup` is the object-safe mirror a group is boxed through, never the trait a group author implements. Python ownership and cancellation guards stay in the Python adapter
`protocol.rs` owns `Protocol` and suspension messages, including `InterceptRequest` and `StreamDelivery`. `call.rs` owns route outputs and their adaptation into a hosted machine. Rust service handling belongs in `host-native::services`. Coroutine channel handles stay in `machine/context.rs`

View file

@ -1,73 +1,141 @@
mod adapter;
mod chain;
//! What a call lets its host hook into, one trait per payload. A host picks the facets
//! its public surface has and composes any number of groups through [`Hooks`].
pub use adapter::DynCallHooks;
pub use chain::{ChainRuntime, HookChain, HookGroupAdapter};
mod chain;
mod group;
use std::future::{Future, ready};
pub use chain::Hooks;
pub use group::AnyGroup;
use crate::{
interceptors::{RawResponse, RequestContext, WireRequest},
lifecycle::{CallEvent, Timing},
interceptors::{ExecutionFacts, RawResponse, RequestContext, WireRequest},
lifecycle::{FailureOrigin, Timing},
};
pub trait HookRuntime {
type Context<'a>: Send;
/// A host's public values as its hooks see them.
pub trait Surface: Sized + 'static {
/// The caller's task a hook awaits on, or `()` for a host that never suspends.
type Task<'a>: Copy + Send;
type Arguments: Send;
type Response: Send + Sync;
type Chunk;
type Error: Send + Sync;
/// What a failed outcome carries: the host's retained form of its error.
type Failure: Sync;
/// How [`Hooks`] stores a mounted group.
type Group: ?Sized + AnyGroup<Self> + Send + Sync;
}
pub type RuntimeCallEvent<'a, R> =
CallEvent<&'a <R as HookRuntime>::Response, &'a <R as HookRuntime>::Failure, &'a RawResponse>;
/// A hook stopped the call: nothing after it runs and the error is raised as it is.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Cancelled<E>(pub E);
pub trait CallHooks<R: HookRuntime>: Sync {
fn prepare_arguments<'a>(
/// The caller's arguments before the route decodes them. Groups run in mount order and
/// each receives the previous group's output; the first error stops the chain.
pub trait ArgumentHooks<S: Surface>: Send + Sync {
fn prepare<'a>(
&'a self,
_runtime: R::Context<'a>,
arguments: R::Arguments,
_task: S::Task<'a>,
arguments: S::Arguments,
_started_at: f64,
) -> impl Future<Output = Result<R::Arguments, R::Error>> + Send {
) -> impl Future<Output = Result<S::Arguments, S::Error>> + Send {
ready(Ok(arguments))
}
fn arguments_prepared(&self, _arguments: &R::Arguments) -> Result<(), R::Error> {
Ok(())
}
fn before_provider_request<'a>(
&'a self,
_runtime: R::Context<'a>,
wire: Box<WireRequest>,
_context: &RequestContext,
) -> impl Future<Output = Result<Box<WireRequest>, R::Error>> + Send {
ready(Ok(wire))
}
fn transform_response<'a>(
&'a self,
_runtime: R::Context<'a>,
response: R::Response,
_timing: Timing,
) -> impl Future<Output = Result<R::Response, R::Error>> + Send {
ready(Ok(response))
}
fn on_event<'a>(
&'a self,
_runtime: R::Context<'a>,
_event: RuntimeCallEvent<'_, R>,
) -> impl Future<Output = Result<(), R::Error>> + Send {
ready(Ok(()))
}
fn on_stream_open(&self, _head: &R::Response) -> Result<(), R::Error> {
Ok(())
}
fn on_stream_chunk(&self, _chunk: &R::Chunk) -> Result<(), R::Error> {
fn prepared(&self, _arguments: &S::Arguments) -> Result<(), S::Error> {
Ok(())
}
}
/// The provider exchange: the wire request about to leave, the raw response that came
/// back, and the facts of the execution. Groups run in mount order; the first error
/// interrupts the call.
pub trait ProviderHooks<S: Surface>: Send + Sync {
fn before_provider_request<'a>(
&'a self,
_task: S::Task<'a>,
wire: Box<WireRequest>,
_context: &RequestContext,
) -> impl Future<Output = Result<Box<WireRequest>, S::Error>> + Send {
ready(Ok(wire))
}
fn after_provider_response<'a>(
&'a self,
_task: S::Task<'a>,
_raw: &RawResponse,
) -> impl Future<Output = Result<(), S::Error>> + Send {
ready(Ok(()))
}
fn result_ready<'a>(
&'a self,
_task: S::Task<'a>,
_facts: &ExecutionFacts,
) -> impl Future<Output = Result<(), S::Error>> + Send {
ready(Ok(()))
}
}
/// The public response and stream. Transforms unwind in reverse mount order, so the first
/// group mounted sees the response last; the first error stops the chain.
pub trait ResponseHooks<S: Surface>: Send + Sync {
fn transform<'a>(
&'a self,
_task: S::Task<'a>,
response: S::Response,
_timing: Timing,
) -> impl Future<Output = Result<S::Response, S::Error>> + Send {
ready(Ok(response))
}
fn stream_opened(&self, _head: &S::Response) -> Result<(), S::Error> {
Ok(())
}
fn stream_chunk(&self, _chunk: &S::Chunk) -> Result<(), S::Error> {
Ok(())
}
}
/// How the call ended. Every group hears the outcome, terminal ones in reverse mount
/// order, and only a [`Cancelled`] group stops the rest. What a group does with its own
/// errors is its business.
pub trait OutcomeHooks<S: Surface>: Send + Sync {
fn started<'a>(
&'a self,
_task: S::Task<'a>,
_start_time: f64,
) -> impl Future<Output = Result<(), Cancelled<S::Error>>> + Send {
ready(Ok(()))
}
fn succeeded<'a>(
&'a self,
_task: S::Task<'a>,
_timing: Timing,
_response: &S::Response,
) -> impl Future<Output = Result<(), Cancelled<S::Error>>> + Send {
ready(Ok(()))
}
fn failed<'a>(
&'a self,
_task: S::Task<'a>,
_timing: Timing,
_origin: FailureOrigin,
_error: &S::Failure,
) -> impl Future<Output = Result<(), Cancelled<S::Error>>> + Send {
ready(Ok(()))
}
fn cancelled<'a>(
&'a self,
_task: S::Task<'a>,
_timing: Timing,
) -> impl Future<Output = Result<(), Cancelled<S::Error>>> + Send {
ready(Ok(()))
}
}

View file

@ -1,91 +0,0 @@
use futures_util::future::BoxFuture;
use super::{CallHooks, HookRuntime, RuntimeCallEvent};
use crate::{
interceptors::{RequestContext, WireRequest},
lifecycle::Timing,
};
pub trait DynCallHooks<R: HookRuntime + 'static>: Sync {
fn prepare_arguments<'a>(
&'a self,
runtime: R::Context<'a>,
arguments: R::Arguments,
started_at: f64,
) -> BoxFuture<'a, Result<R::Arguments, R::Error>>;
fn arguments_prepared(&self, arguments: &R::Arguments) -> Result<(), R::Error>;
fn before_provider_request<'a: 'b, 'b>(
&'a self,
runtime: R::Context<'a>,
wire: Box<WireRequest>,
context: &'b RequestContext,
) -> BoxFuture<'b, Result<Box<WireRequest>, R::Error>>;
fn transform_response<'a>(
&'a self,
runtime: R::Context<'a>,
response: R::Response,
timing: Timing,
) -> BoxFuture<'a, Result<R::Response, R::Error>>;
fn on_event<'a: 'b, 'b>(
&'a self,
runtime: R::Context<'a>,
event: RuntimeCallEvent<'b, R>,
) -> BoxFuture<'b, Result<(), R::Error>>;
fn on_stream_open(&self, head: &R::Response) -> Result<(), R::Error>;
fn on_stream_chunk(&self, chunk: &R::Chunk) -> Result<(), R::Error>;
}
impl<R: HookRuntime + 'static, H: CallHooks<R>> DynCallHooks<R> for H {
fn prepare_arguments<'a>(
&'a self,
runtime: R::Context<'a>,
arguments: R::Arguments,
started_at: f64,
) -> BoxFuture<'a, Result<R::Arguments, R::Error>> {
Box::pin(CallHooks::prepare_arguments(
self, runtime, arguments, started_at,
))
}
fn arguments_prepared(&self, arguments: &R::Arguments) -> Result<(), R::Error> {
CallHooks::arguments_prepared(self, arguments)
}
fn before_provider_request<'a: 'b, 'b>(
&'a self,
runtime: R::Context<'a>,
wire: Box<WireRequest>,
context: &'b RequestContext,
) -> BoxFuture<'b, Result<Box<WireRequest>, R::Error>> {
Box::pin(CallHooks::before_provider_request(
self, runtime, wire, context,
))
}
fn transform_response<'a>(
&'a self,
runtime: R::Context<'a>,
response: R::Response,
timing: Timing,
) -> BoxFuture<'a, Result<R::Response, R::Error>> {
Box::pin(CallHooks::transform_response(
self, runtime, response, timing,
))
}
fn on_event<'a: 'b, 'b>(
&'a self,
runtime: R::Context<'a>,
event: RuntimeCallEvent<'b, R>,
) -> BoxFuture<'b, Result<(), R::Error>> {
Box::pin(CallHooks::on_event(self, runtime, event))
}
fn on_stream_open(&self, head: &R::Response) -> Result<(), R::Error> {
CallHooks::on_stream_open(self, head)
}
fn on_stream_chunk(&self, chunk: &R::Chunk) -> Result<(), R::Error> {
CallHooks::on_stream_chunk(self, chunk)
}
}

View file

@ -1,133 +1,175 @@
use super::{CallHooks, DynCallHooks, HookRuntime, RuntimeCallEvent};
use super::{
AnyGroup, ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks, Surface,
};
use crate::{
interceptors::{RequestContext, WireRequest},
lifecycle::{CallEvent, Timing},
interceptors::{ExecutionFacts, RawResponse, RequestContext, WireRequest},
lifecycle::{FailureOrigin, Timing},
};
pub trait ChainRuntime: HookRuntime + Sized + 'static {
type Group: DynCallHooks<Self> + ?Sized;
/// `Ok` lets the next group receive the terminal event, `Err` stops the chain.
fn terminal_notification_failed(error: Self::Error) -> Result<(), Self::Error>;
/// The groups mounted on a call, run as one group of each facet. Request-side stages run
/// in mount order, response-side and terminal stages unwind in reverse, like layers.
pub struct Hooks<S: Surface> {
groups: Vec<Box<S::Group>>,
}
pub trait HookGroupAdapter<H>: ChainRuntime {
fn adapt_group(hooks: H) -> Box<Self::Group>;
}
pub struct HookChain<R: ChainRuntime> {
hooks: Vec<Box<R::Group>>,
}
impl<R: ChainRuntime> HookChain<R> {
pub fn groups(&self) -> impl Iterator<Item = &R::Group> {
self.hooks.iter().map(Box::as_ref)
}
pub fn groups_mut(&mut self) -> impl Iterator<Item = &mut R::Group> {
self.hooks.iter_mut().map(Box::as_mut)
}
impl<S: Surface> Hooks<S> {
pub fn new() -> Self {
Self::default()
}
pub fn with<H>(mut self, hooks: H) -> Self
where
R: HookGroupAdapter<H>,
{
self.hooks.push(R::adapt_group(hooks));
pub fn with(mut self, group: Box<S::Group>) -> Self {
self.groups.push(group);
self
}
pub fn groups(&self) -> impl Iterator<Item = &S::Group> {
self.groups.iter().map(Box::as_ref)
}
pub fn groups_mut(&mut self) -> impl Iterator<Item = &mut S::Group> {
self.groups.iter_mut().map(Box::as_mut)
}
}
impl<R: ChainRuntime> CallHooks<R> for HookChain<R>
where
for<'a> R::Context<'a>: Copy,
{
async fn prepare_arguments<'a>(
impl<S: Surface> Default for Hooks<S> {
fn default() -> Self {
Self { groups: Vec::new() }
}
}
impl<S: Surface> ArgumentHooks<S> for Hooks<S> {
async fn prepare<'a>(
&'a self,
runtime: R::Context<'a>,
arguments: R::Arguments,
task: S::Task<'a>,
arguments: S::Arguments,
started_at: f64,
) -> Result<R::Arguments, R::Error> {
) -> Result<S::Arguments, S::Error> {
let mut arguments = arguments;
for hooks in &self.hooks {
arguments = hooks
.prepare_arguments(runtime, arguments, started_at)
.await?;
for group in &self.groups {
arguments = group.prepare(task, arguments, started_at).await?;
}
Ok(arguments)
}
fn arguments_prepared(&self, arguments: &R::Arguments) -> Result<(), R::Error> {
self.hooks
fn prepared(&self, arguments: &S::Arguments) -> Result<(), S::Error> {
self.groups
.iter()
.try_for_each(|hooks| hooks.arguments_prepared(arguments))
.try_for_each(|group| group.prepared(arguments))
}
}
impl<S: Surface> ProviderHooks<S> for Hooks<S> {
async fn before_provider_request<'a>(
&'a self,
runtime: R::Context<'a>,
task: S::Task<'a>,
wire: Box<WireRequest>,
context: &RequestContext,
) -> Result<Box<WireRequest>, R::Error> {
) -> Result<Box<WireRequest>, S::Error> {
let mut wire = wire;
for hooks in &self.hooks {
wire = hooks
.before_provider_request(runtime, wire, context)
.await?;
for group in &self.groups {
wire = group.before_provider_request(task, wire, context).await?;
}
Ok(wire)
}
async fn transform_response<'a>(
async fn after_provider_response<'a>(
&'a self,
runtime: R::Context<'a>,
response: R::Response,
timing: Timing,
) -> Result<R::Response, R::Error> {
let mut response = response;
for hooks in &self.hooks {
response = hooks.transform_response(runtime, response, timing).await?;
}
Ok(response)
}
async fn on_event<'a>(
&'a self,
runtime: R::Context<'a>,
event: RuntimeCallEvent<'_, R>,
) -> Result<(), R::Error> {
let terminal = matches!(
event,
CallEvent::Succeeded { .. } | CallEvent::Failed { .. }
);
for hooks in &self.hooks {
match hooks.on_event(runtime, event.clone()).await {
Err(error) if terminal => R::terminal_notification_failed(error)?,
result => result?,
}
task: S::Task<'a>,
raw: &RawResponse,
) -> Result<(), S::Error> {
for group in &self.groups {
group.after_provider_response(task, raw).await?;
}
Ok(())
}
fn on_stream_open(&self, head: &R::Response) -> Result<(), R::Error> {
self.hooks
.iter()
.try_for_each(|hooks| hooks.on_stream_open(head))
}
fn on_stream_chunk(&self, chunk: &R::Chunk) -> Result<(), R::Error> {
self.hooks
.iter()
.try_for_each(|hooks| hooks.on_stream_chunk(chunk))
async fn result_ready<'a>(
&'a self,
task: S::Task<'a>,
facts: &ExecutionFacts,
) -> Result<(), S::Error> {
for group in &self.groups {
group.result_ready(task, facts).await?;
}
Ok(())
}
}
impl<R: ChainRuntime> Default for HookChain<R> {
fn default() -> Self {
Self { hooks: Vec::new() }
impl<S: Surface> ResponseHooks<S> for Hooks<S> {
async fn transform<'a>(
&'a self,
task: S::Task<'a>,
response: S::Response,
timing: Timing,
) -> Result<S::Response, S::Error> {
let mut response = response;
for group in self.groups.iter().rev() {
response = group.transform(task, response, timing).await?;
}
Ok(response)
}
fn stream_opened(&self, head: &S::Response) -> Result<(), S::Error> {
self.groups
.iter()
.rev()
.try_for_each(|group| group.stream_opened(head))
}
fn stream_chunk(&self, chunk: &S::Chunk) -> Result<(), S::Error> {
self.groups
.iter()
.rev()
.try_for_each(|group| group.stream_chunk(chunk))
}
}
impl<S: Surface> OutcomeHooks<S> for Hooks<S> {
async fn started<'a>(
&'a self,
task: S::Task<'a>,
start_time: f64,
) -> Result<(), Cancelled<S::Error>> {
for group in &self.groups {
group.started(task, start_time).await?;
}
Ok(())
}
async fn succeeded<'a>(
&'a self,
task: S::Task<'a>,
timing: Timing,
response: &S::Response,
) -> Result<(), Cancelled<S::Error>> {
for group in self.groups.iter().rev() {
group.succeeded(task, timing, response).await?;
}
Ok(())
}
async fn failed<'a>(
&'a self,
task: S::Task<'a>,
timing: Timing,
origin: FailureOrigin,
error: &S::Failure,
) -> Result<(), Cancelled<S::Error>> {
for group in self.groups.iter().rev() {
group.failed(task, timing, origin, error).await?;
}
Ok(())
}
async fn cancelled<'a>(
&'a self,
task: S::Task<'a>,
timing: Timing,
) -> Result<(), Cancelled<S::Error>> {
for group in self.groups.iter().rev() {
group.cancelled(task, timing).await?;
}
Ok(())
}
}
@ -138,36 +180,21 @@ mod tests {
use rstest::{fixture, rstest};
use super::*;
use crate::lifecycle::FailureOrigin;
struct TestRuntime;
struct TestSurface;
impl HookRuntime for TestRuntime {
type Context<'a> = &'a Mutex<Vec<String>>;
impl Surface for TestSurface {
type Task<'a> = &'a Mutex<Vec<String>>;
type Arguments = String;
type Response = String;
type Chunk = u32;
type Error = &'static str;
type Failure = &'static str;
type Group = dyn AnyGroup<Self>;
}
impl ChainRuntime for TestRuntime {
type Group = dyn DynCallHooks<Self>;
fn terminal_notification_failed(error: &'static str) -> Result<(), &'static str> {
match error {
"cancelled" => Err(error),
_ => Ok(()),
}
}
}
impl<H: CallHooks<TestRuntime> + 'static> HookGroupAdapter<H> for TestRuntime {
fn adapt_group(hooks: H) -> Box<Self::Group> {
Box::new(hooks)
}
}
/// A group that tags every value it transforms with its name, logs every stage it
/// runs, and fails one stage when told to.
struct Group {
name: &'static str,
suspensions: u8,
@ -175,6 +202,10 @@ mod tests {
}
impl Group {
fn mounted(self) -> Box<dyn AnyGroup<TestSurface>> {
Box::new(self)
}
async fn value(
&self,
trace: &Mutex<Vec<String>>,
@ -192,8 +223,8 @@ mod tests {
}
}
impl CallHooks<TestRuntime> for Group {
async fn prepare_arguments<'a>(
impl ArgumentHooks<TestSurface> for Group {
async fn prepare<'a>(
&'a self,
trace: &'a Mutex<Vec<String>>,
arguments: String,
@ -201,7 +232,9 @@ mod tests {
) -> Result<String, &'static str> {
self.value(trace, "prepare", arguments).await
}
}
impl ProviderHooks<TestSurface> for Group {
async fn before_provider_request<'a>(
&'a self,
trace: &'a Mutex<Vec<String>>,
@ -211,8 +244,10 @@ mod tests {
let url = self.value(trace, "wire", wire.url.clone()).await?;
Ok(Box::new(WireRequest { url, ..*wire }))
}
}
async fn transform_response<'a>(
impl ResponseHooks<TestSurface> for Group {
async fn transform<'a>(
&'a self,
trace: &'a Mutex<Vec<String>>,
response: String,
@ -220,18 +255,43 @@ mod tests {
) -> Result<String, &'static str> {
self.value(trace, "response", response).await
}
}
async fn on_event<'a>(
impl OutcomeHooks<TestSurface> for Group {
async fn started<'a>(
&'a self,
trace: &'a Mutex<Vec<String>>,
event: RuntimeCallEvent<'_, TestRuntime>,
) -> Result<(), &'static str> {
let label = match event {
CallEvent::Succeeded { response, .. } => response.as_str(),
CallEvent::Failed { error, .. } => *error,
_ => "event",
};
self.value(trace, label, String::new()).await.map(|_| ())
_: f64,
) -> Result<(), Cancelled<&'static str>> {
self.value(trace, "started", String::new())
.await
.map(drop)
.map_err(Cancelled)
}
async fn succeeded<'a>(
&'a self,
trace: &'a Mutex<Vec<String>>,
_: Timing,
response: &String,
) -> Result<(), Cancelled<&'static str>> {
self.value(trace, response, String::new())
.await
.map(drop)
.map_err(Cancelled)
}
async fn failed<'a>(
&'a self,
trace: &'a Mutex<Vec<String>>,
_: Timing,
_: FailureOrigin,
error: &&'static str,
) -> Result<(), Cancelled<&'static str>> {
self.value(trace, error, String::new())
.await
.map(drop)
.map_err(Cancelled)
}
}
@ -259,147 +319,128 @@ mod tests {
}
}
fn chain(
suspensions: u8,
rejecting: Option<(&'static str, &'static str)>,
) -> Hooks<TestSurface> {
["a", "b", "c"]
.into_iter()
.fold(Hooks::new(), |hooks, name| {
hooks.with(
Group {
name,
suspensions,
reject: rejecting
.filter(|(group, _)| *group == name)
.map(|(_, error)| error),
}
.mounted(),
)
})
}
#[rstest]
#[tokio::test]
async fn groups_compose_without_python(
async fn request_stages_run_in_mount_order_and_response_stages_unwind(
trace: Mutex<Vec<String>>,
timing: Timing,
context: RequestContext,
#[values(0, 1, 2)] suspensions: u8,
) {
let hooks = HookChain::<TestRuntime>::new()
.with(Group {
name: "a",
suspensions,
reject: None,
})
.with(Group {
name: "b",
suspensions,
reject: None,
});
let arguments = CallHooks::prepare_arguments(&hooks, &trace, "input".into(), 0.0)
let hooks = chain(suspensions, None);
OutcomeHooks::started(&hooks, &trace, 0.0).await.unwrap();
let arguments = ArgumentHooks::prepare(&hooks, &trace, "input".into(), 0.0)
.await
.unwrap();
assert_eq!(arguments, "inputab");
assert_eq!(arguments, "inputabc");
let wire = Box::new(WireRequest {
url: "url".into(),
headers: Vec::new(),
body: serde_json::Value::Null,
});
let rewritten = CallHooks::before_provider_request(&hooks, &trace, wire, &context)
let rewritten = ProviderHooks::before_provider_request(&hooks, &trace, wire, &context)
.await
.unwrap();
assert_eq!(rewritten.url, "urlab");
let response = CallHooks::transform_response(&hooks, &trace, "output".into(), timing)
assert_eq!(rewritten.url, "urlabc");
let response = ResponseHooks::transform(&hooks, &trace, "output".into(), timing)
.await
.unwrap();
assert_eq!(response, "outputcba");
OutcomeHooks::succeeded(&hooks, &trace, timing, &response)
.await
.unwrap();
assert_eq!(response, "outputab");
CallHooks::on_event(
&hooks,
&trace,
CallEvent::Succeeded {
timing,
response: &response,
},
)
.await
.unwrap();
assert_eq!(
*trace.lock().unwrap(),
[
"started:a",
"started:b",
"started:c",
"prepare:a",
"prepare:b",
"prepare:c",
"wire:a",
"wire:b",
"response:a",
"wire:c",
"response:c",
"response:b",
"outputab:a",
"outputab:b",
"response:a",
"outputcba:c",
"outputcba:b",
"outputcba:a",
]
);
}
fn rejecting_chain(suspensions: u8, error: &'static str) -> HookChain<TestRuntime> {
["a", "b", "c"]
.into_iter()
.fold(HookChain::new(), |hooks, name| {
hooks.with(Group {
name,
suspensions,
reject: (name == "b").then_some(error),
})
})
}
#[rstest]
#[tokio::test]
async fn transformation_failure_stops_later_groups(
async fn a_transformation_error_stops_later_groups(
trace: Mutex<Vec<String>>,
#[values(0, 2)] suspensions: u8,
) {
let hooks = rejecting_chain(suspensions, "rejected");
let result = CallHooks::prepare_arguments(&hooks, &trace, String::new(), 0.0).await;
let hooks = chain(suspensions, Some(("b", "rejected")));
let result = ArgumentHooks::prepare(&hooks, &trace, String::new(), 0.0).await;
assert_eq!(result, Err("rejected"));
assert_eq!(*trace.lock().unwrap(), ["prepare:a", "prepare:b"]);
}
#[rstest]
#[case::success_reaches_every_group(true, "rejected", Ok(()), &["selected:a", "selected:b", "selected:c"])]
#[case::failure_reaches_every_group(false, "rejected", Ok(()), &["failed:a", "failed:b", "failed:c"])]
#[case::cancellation_stops_success(true, "cancelled", Err("cancelled"), &["selected:a", "selected:b"])]
#[case::cancellation_stops_failure(false, "cancelled", Err("cancelled"), &["failed:a", "failed:b"])]
#[case::success(true)]
#[case::failure(false)]
#[tokio::test]
async fn terminal_notification_errors_follow_the_runtime(
async fn a_cancelled_outcome_stops_the_groups_mounted_before_it(
trace: Mutex<Vec<String>>,
timing: Timing,
#[values(0, 2)] suspensions: u8,
#[case] succeeded: bool,
#[case] error: &'static str,
#[case] expected: Result<(), &'static str>,
#[case] reached: &[&str],
) {
let hooks = rejecting_chain(suspensions, error);
let hooks = chain(suspensions, Some(("b", "cancelled")));
let response = "selected".to_owned();
let event = if succeeded {
CallEvent::Succeeded {
timing,
response: &response,
}
let result = if succeeded {
OutcomeHooks::succeeded(&hooks, &trace, timing, &response).await
} else {
CallEvent::Failed {
timing,
origin: FailureOrigin::Host,
error: &"failed",
}
OutcomeHooks::failed(&hooks, &trace, timing, FailureOrigin::Host, &"failed").await
};
assert_eq!(result, Err(Cancelled("cancelled")));
let reached = if succeeded {
["selected:c", "selected:b"]
} else {
["failed:c", "failed:b"]
};
let result = CallHooks::on_event(&hooks, &trace, event).await;
assert_eq!(result, expected);
assert_eq!(*trace.lock().unwrap(), reached);
}
#[rstest]
#[tokio::test]
async fn non_terminal_notification_error_stops_later_groups(trace: Mutex<Vec<String>>) {
let hooks = rejecting_chain(0, "rejected");
let result =
CallHooks::on_event(&hooks, &trace, CallEvent::Started { start_time: 0.0 }).await;
assert_eq!(result, Err("rejected"));
assert_eq!(*trace.lock().unwrap(), ["event:a", "event:b"]);
}
#[rstest]
#[tokio::test]
async fn empty_chain_preserves_values(trace: Mutex<Vec<String>>, timing: Timing) {
let hooks = HookChain::<TestRuntime>::new();
async fn an_empty_chain_preserves_values(trace: Mutex<Vec<String>>, timing: Timing) {
let hooks = Hooks::<TestSurface>::new();
assert_eq!(
CallHooks::prepare_arguments(&hooks, &trace, "input".into(), 0.0)
ArgumentHooks::prepare(&hooks, &trace, "input".into(), 0.0)
.await
.unwrap(),
"input"
);
assert_eq!(
CallHooks::transform_response(&hooks, &trace, "output".into(), timing)
ResponseHooks::transform(&hooks, &trace, "output".into(), timing)
.await
.unwrap(),
"output"

View file

@ -0,0 +1,164 @@
use futures_util::future::BoxFuture;
use super::{ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks, Surface};
use crate::{
interceptors::{ExecutionFacts, RawResponse, RequestContext, WireRequest},
lifecycle::{FailureOrigin, Timing},
};
/// Every facet behind one object-safe trait, so [`super::Hooks`] can store groups of
/// different types. Implemented for anything that implements the four facets.
pub trait AnyGroup<S: Surface>: Send + Sync {
fn prepare<'a>(
&'a self,
task: S::Task<'a>,
arguments: S::Arguments,
started_at: f64,
) -> BoxFuture<'a, Result<S::Arguments, S::Error>>;
fn prepared(&self, arguments: &S::Arguments) -> Result<(), S::Error>;
fn before_provider_request<'a: 'b, 'b>(
&'a self,
task: S::Task<'a>,
wire: Box<WireRequest>,
context: &'b RequestContext,
) -> BoxFuture<'b, Result<Box<WireRequest>, S::Error>>;
fn after_provider_response<'a: 'b, 'b>(
&'a self,
task: S::Task<'a>,
raw: &'b RawResponse,
) -> BoxFuture<'b, Result<(), S::Error>>;
fn result_ready<'a: 'b, 'b>(
&'a self,
task: S::Task<'a>,
facts: &'b ExecutionFacts,
) -> BoxFuture<'b, Result<(), S::Error>>;
fn transform<'a>(
&'a self,
task: S::Task<'a>,
response: S::Response,
timing: Timing,
) -> BoxFuture<'a, Result<S::Response, S::Error>>;
fn stream_opened(&self, head: &S::Response) -> Result<(), S::Error>;
fn stream_chunk(&self, chunk: &S::Chunk) -> Result<(), S::Error>;
fn started<'a>(
&'a self,
task: S::Task<'a>,
start_time: f64,
) -> BoxFuture<'a, Result<(), Cancelled<S::Error>>>;
fn succeeded<'a: 'b, 'b>(
&'a self,
task: S::Task<'a>,
timing: Timing,
response: &'b S::Response,
) -> BoxFuture<'b, Result<(), Cancelled<S::Error>>>;
fn failed<'a: 'b, 'b>(
&'a self,
task: S::Task<'a>,
timing: Timing,
origin: FailureOrigin,
error: &'b S::Failure,
) -> BoxFuture<'b, Result<(), Cancelled<S::Error>>>;
fn cancelled<'a>(
&'a self,
task: S::Task<'a>,
timing: Timing,
) -> BoxFuture<'a, Result<(), Cancelled<S::Error>>>;
}
impl<S, H> AnyGroup<S> for H
where
S: Surface,
H: ArgumentHooks<S> + ProviderHooks<S> + ResponseHooks<S> + OutcomeHooks<S>,
{
fn prepare<'a>(
&'a self,
task: S::Task<'a>,
arguments: S::Arguments,
started_at: f64,
) -> BoxFuture<'a, Result<S::Arguments, S::Error>> {
Box::pin(ArgumentHooks::prepare(self, task, arguments, started_at))
}
fn prepared(&self, arguments: &S::Arguments) -> Result<(), S::Error> {
ArgumentHooks::prepared(self, arguments)
}
fn before_provider_request<'a: 'b, 'b>(
&'a self,
task: S::Task<'a>,
wire: Box<WireRequest>,
context: &'b RequestContext,
) -> BoxFuture<'b, Result<Box<WireRequest>, S::Error>> {
Box::pin(ProviderHooks::before_provider_request(
self, task, wire, context,
))
}
fn after_provider_response<'a: 'b, 'b>(
&'a self,
task: S::Task<'a>,
raw: &'b RawResponse,
) -> BoxFuture<'b, Result<(), S::Error>> {
Box::pin(ProviderHooks::after_provider_response(self, task, raw))
}
fn result_ready<'a: 'b, 'b>(
&'a self,
task: S::Task<'a>,
facts: &'b ExecutionFacts,
) -> BoxFuture<'b, Result<(), S::Error>> {
Box::pin(ProviderHooks::result_ready(self, task, facts))
}
fn transform<'a>(
&'a self,
task: S::Task<'a>,
response: S::Response,
timing: Timing,
) -> BoxFuture<'a, Result<S::Response, S::Error>> {
Box::pin(ResponseHooks::transform(self, task, response, timing))
}
fn stream_opened(&self, head: &S::Response) -> Result<(), S::Error> {
ResponseHooks::stream_opened(self, head)
}
fn stream_chunk(&self, chunk: &S::Chunk) -> Result<(), S::Error> {
ResponseHooks::stream_chunk(self, chunk)
}
fn started<'a>(
&'a self,
task: S::Task<'a>,
start_time: f64,
) -> BoxFuture<'a, Result<(), Cancelled<S::Error>>> {
Box::pin(OutcomeHooks::started(self, task, start_time))
}
fn succeeded<'a: 'b, 'b>(
&'a self,
task: S::Task<'a>,
timing: Timing,
response: &'b S::Response,
) -> BoxFuture<'b, Result<(), Cancelled<S::Error>>> {
Box::pin(OutcomeHooks::succeeded(self, task, timing, response))
}
fn failed<'a: 'b, 'b>(
&'a self,
task: S::Task<'a>,
timing: Timing,
origin: FailureOrigin,
error: &'b S::Failure,
) -> BoxFuture<'b, Result<(), Cancelled<S::Error>>> {
Box::pin(OutcomeHooks::failed(self, task, timing, origin, error))
}
fn cancelled<'a>(
&'a self,
task: S::Task<'a>,
timing: Timing,
) -> BoxFuture<'a, Result<(), Cancelled<S::Error>>> {
Box::pin(OutcomeHooks::cancelled(self, task, timing))
}
}

View file

@ -1,5 +1,5 @@
use litellm_host::hooks::CallHooks;
use litellm_host_python::{PythonOwned, PythonRuntime};
use litellm_host::hooks::{ArgumentHooks, OutcomeHooks, ProviderHooks, ResponseHooks};
use litellm_host_python::{Awaiter, PythonOwned, PythonSurface};
use pyo3::{
gc::{PyTraverseError, PyVisit},
prelude::*,
@ -33,22 +33,33 @@ impl PythonPreflight {
#[cfg(test)]
pub(crate) const PYTHON_CONTRACT: &str = include_str!("../preflight_contract.json");
/// The SDK's request policy as one argument group: inherited credentials are written into
/// the keyword view and a call over its limits is rejected before any provider work.
pub(crate) struct SdkPolicy;
impl CallHooks<PythonRuntime> for SdkPolicy {
fn arguments_prepared(&self, arguments: &Py<PyDict>) -> PyResult<()> {
impl ArgumentHooks<PythonSurface> for SdkPolicy {
async fn prepare<'a>(
&'a self,
_: &'a dyn Awaiter,
arguments: Py<PyDict>,
_: f64,
) -> PyResult<Py<PyDict>> {
Python::attach(|py| {
inherit_credentials(py, arguments.bind(py), || {
Ok(PythonPreflight::CredentialList
.call(py, ())?
.cast_into::<PyList>()?)
})?;
PythonPreflight::CheckLimits.call(py, (arguments,))?;
Ok(())
PythonPreflight::CheckLimits.call(py, (&arguments,))?;
Ok(arguments)
})
}
}
impl ProviderHooks<PythonSurface> for SdkPolicy {}
impl ResponseHooks<PythonSurface> for SdkPolicy {}
impl OutcomeHooks<PythonSurface> for SdkPolicy {}
impl PythonOwned for SdkPolicy {
fn close(&mut self, _: Python<'_>) {}
@ -108,7 +119,11 @@ fn inherit_credentials<'py>(
#[cfg(test)]
mod tests {
use std::collections::BTreeSet;
use std::pin::pin;
use std::sync::Mutex;
use std::task::{Context, Poll, Waker};
use futures_util::future::BoxFuture;
use super::*;
use strum::VariantArray;
@ -427,9 +442,7 @@ preflight.credential_list = lambda: [Credential(), Credential()]
arguments = {'litellm_credential_name': 'missing'}
",
);
SdkPolicy
.arguments_prepared(&argument_dict(&locals).unbind())
.unwrap();
prepared(&argument_dict(&locals).unbind()).unwrap();
py.run(
c"
assert arguments == {'litellm_credential_name': 'missing'}, arguments
@ -465,9 +478,7 @@ preflight.check_limits = check_limits
arguments = {'litellm_credential_name': 'ocr-test'}
",
);
let error = SdkPolicy
.arguments_prepared(&argument_dict(&locals).unbind())
.unwrap_err();
let error = prepared(&argument_dict(&locals).unbind()).unwrap_err();
assert!(
error
.value(py)
@ -485,6 +496,25 @@ assert arguments['api_key'] == 'inherited'
});
}
/// Runs the policy over `arguments`. It never awaits, so the task is never asked.
fn prepared(arguments: &Py<PyDict>) -> PyResult<Py<PyDict>> {
struct NeverAwaited;
impl Awaiter for NeverAwaited {
fn awaited(&self, _: Py<PyAny>) -> BoxFuture<'_, PyResult<Py<PyAny>>> {
unreachable!("the SDK policy never awaits")
}
}
let arguments = Python::attach(|py| arguments.clone_ref(py));
let mut prepare = pin!(SdkPolicy.prepare(&NeverAwaited, arguments, 0.0));
match prepare
.as_mut()
.poll(&mut Context::from_waker(Waker::noop()))
{
Poll::Ready(result) => result,
Poll::Pending => unreachable!("the SDK policy never suspends"),
}
}
fn argument_dict<'py>(locals: &Bound<'py, PyDict>) -> Bound<'py, PyDict> {
locals
.get_item("arguments")

View file

@ -11,10 +11,11 @@ pub(crate) mod traces;
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::{
CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
CallOptions, Hooks, PythonBinding, PythonGroup, PythonHostCalls, PythonSurface,
effective_py_args,
};
use litellm_host_python::{Error as HostError, present};
use parking_lot::Mutex;
use pyo3::{
gc::{PyTraverseError, PyVisit},
@ -119,7 +120,7 @@ fn call_hooks(
operation: LoggingOperation,
call: &NativeCall<'_>,
asynchronous: bool,
) -> PyResult<(Py<PyDict>, impl PythonCallHooks + use<>)> {
) -> PyResult<(Py<PyDict>, impl PythonGroup + use<>)> {
let call = PublicCall::capture(&call.base, &call.args, &call.kwargs)?;
let arguments = call.arguments(py);
Ok((
@ -144,7 +145,7 @@ fn run_public_call<H, M>(
+ Sync
+ 'static,
host: H,
hooks: impl PythonCallHooks + 'static,
hooks: impl PythonGroup + 'static,
asynchronous: bool,
) -> PyResult<Py<PyAny>>
where
@ -158,9 +159,9 @@ where
start(py, arguments, request).map(crate::logger::LoggedMachine::new)
},
host,
HookChain::new()
.with(hooks)
.with(crate::preflight::SdkPolicy),
Hooks::new()
.with(PythonSurface::group(hooks))
.with(PythonSurface::group(crate::preflight::SdkPolicy)),
arguments,
CallOptions::new(asynchronous, execution_steps),
)