Skip to main content

ferrin_core/
hooks.rs

1//! Lifecycle hooks: asynchronous callbacks awaited by the core.
2
3use std::fmt;
4use std::future::Future;
5use std::panic::AssertUnwindSafe;
6use std::panic::catch_unwind;
7use std::sync::Arc;
8
9use ferrin_spec::BoxFuture;
10use futures_util::FutureExt;
11use futures_util::future::join_all;
12
13use crate::generate_text::StepResult;
14use crate::stream_text::StreamEvent;
15use crate::telemetry::AbortEvent;
16use crate::telemetry::EndEvent;
17use crate::telemetry::ModelCallEndEvent;
18use crate::telemetry::ModelCallStartEvent;
19use crate::telemetry::StartEvent;
20use crate::telemetry::StepStartEvent;
21use crate::telemetry::ToolExecutionEndEvent;
22use crate::telemetry::ToolExecutionStartEvent;
23
24/// An asynchronous callback for event type `E`.
25///
26/// Implemented for every `Fn(Arc<E>) -> impl Future<Output = ()>` closure, so
27/// builders accept `|event| async move { ... }` directly. The core awaits the
28/// returned future before continuing, concurrently with other hooks for the
29/// same event. Unwinding callback panics are isolated; the normal panic handler
30/// still runs, and `panic = "abort"` cannot be isolated.
31pub trait HookFn<E>: Send + Sync + 'static {
32    /// Handles one event.
33    fn call(&self, event: Arc<E>) -> BoxFuture<'static, ()>;
34}
35
36impl<E, F, Fut> HookFn<E> for F
37where
38    F: Fn(Arc<E>) -> Fut + Send + Sync + 'static,
39    Fut: Future<Output = ()> + Send + 'static,
40{
41    fn call(&self, event: Arc<E>) -> BoxFuture<'static, ()> {
42        Box::pin(self(event))
43    }
44}
45
46/// A list of hooks for one event type.
47pub type HookList<E> = Vec<Arc<dyn HookFn<E>>>;
48
49/// All lifecycle hooks of a call.
50#[derive(Clone, Default)]
51pub struct Hooks {
52    /// The call started.
53    pub on_start: HookList<StartEvent>,
54    /// A step started.
55    pub on_step_start: HookList<StepStartEvent>,
56    /// A model call is about to be made.
57    pub on_language_model_call_start: HookList<ModelCallStartEvent>,
58    /// A model call finished (tool calls parsed, not yet executed).
59    pub on_language_model_call_end: HookList<ModelCallEndEvent>,
60    /// A tool execution started.
61    pub on_tool_execution_start: HookList<ToolExecutionStartEvent>,
62    /// A tool execution finished.
63    pub on_tool_execution_end: HookList<ToolExecutionEndEvent>,
64    /// A step finished.
65    pub on_step_end: HookList<StepResult>,
66    /// The call finished.
67    pub on_end: HookList<EndEvent>,
68    /// Streaming: a transformed event or a provider error entering retry handling.
69    pub on_chunk: HookList<StreamEvent>,
70    /// Streaming: the call was aborted.
71    pub on_abort: HookList<AbortEvent>,
72}
73
74impl Hooks {
75    /// Appends the hooks of `other` after those of `self` (settings-level
76    /// callbacks are invoked before call-level callbacks; their futures run
77    /// concurrently).
78    #[must_use]
79    pub fn merged(mut self, other: Hooks) -> Hooks {
80        self.on_start.extend(other.on_start);
81        self.on_step_start.extend(other.on_step_start);
82        self.on_language_model_call_start
83            .extend(other.on_language_model_call_start);
84        self.on_language_model_call_end
85            .extend(other.on_language_model_call_end);
86        self.on_tool_execution_start
87            .extend(other.on_tool_execution_start);
88        self.on_tool_execution_end
89            .extend(other.on_tool_execution_end);
90        self.on_step_end.extend(other.on_step_end);
91        self.on_end.extend(other.on_end);
92        self.on_chunk.extend(other.on_chunk);
93        self.on_abort.extend(other.on_abort);
94        self
95    }
96
97    /// Invokes hooks in list order and awaits all returned futures concurrently.
98    ///
99    /// Completion order is unspecified. Unwinding panics from invoking or
100    /// polling a callback are ignored after the normal panic handler runs.
101    /// This cannot isolate panics when compiled with `panic = "abort"`.
102    pub async fn emit<E: 'static>(list: &[Arc<dyn HookFn<E>>], event: Arc<E>) {
103        let futures = list.iter().filter_map(|hook| {
104            catch_unwind(AssertUnwindSafe(|| hook.call(Arc::clone(&event))))
105                .ok()
106                .map(|future| AssertUnwindSafe(future).catch_unwind())
107        });
108        let _ = join_all(futures).await;
109    }
110}
111
112impl fmt::Debug for Hooks {
113    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
114        f.debug_struct("Hooks")
115            .field("on_start", &self.on_start.len())
116            .field("on_step_start", &self.on_step_start.len())
117            .field(
118                "on_language_model_call_start",
119                &self.on_language_model_call_start.len(),
120            )
121            .field(
122                "on_language_model_call_end",
123                &self.on_language_model_call_end.len(),
124            )
125            .field(
126                "on_tool_execution_start",
127                &self.on_tool_execution_start.len(),
128            )
129            .field("on_tool_execution_end", &self.on_tool_execution_end.len())
130            .field("on_step_end", &self.on_step_end.len())
131            .field("on_end", &self.on_end.len())
132            .field("on_chunk", &self.on_chunk.len())
133            .field("on_abort", &self.on_abort.len())
134            .finish()
135    }
136}