use std::sync::Arc;
use serde_json::Value;
use crate::agent::AgentOutcome;
use std::collections::HashMap;
use std::time::Instant;
#[derive(Debug, Clone)]
pub enum HookAction {
Continue,
Skip,
Error(String),
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum HookEvent<'a> {
SessionStart {
goal: &'a str,
},
PreToolCall {
name: &'a str,
args: &'a Value,
},
PostToolCall {
name: &'a str,
args: &'a Value,
result: &'a str,
duration_ms: u64,
},
PreCompact {
transcript_len: usize,
},
PostCompact {
removed: usize,
summary_chars: usize,
},
SessionEnd {
outcome: &'a AgentOutcome,
},
}
pub trait Hook: Send + Sync {
fn on_event(&self, event: HookEvent) -> HookAction;
}
#[derive(Default)]
pub struct HookRegistry {
hooks: Vec<Arc<dyn Hook>>,
}
impl HookRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register(&mut self, hook: Arc<dyn Hook>) {
self.hooks.push(hook);
}
pub fn dispatch(&self, event: HookEvent) -> HookAction {
let is_pre_tool = matches!(event, HookEvent::PreToolCall { .. });
for hook in &self.hooks {
match hook.on_event(event.clone()) {
HookAction::Continue => continue,
HookAction::Skip if is_pre_tool => return HookAction::Skip,
HookAction::Error(msg) if is_pre_tool => return HookAction::Error(msg),
_ => continue,
}
}
HookAction::Continue
}
pub fn is_empty(&self) -> bool {
self.hooks.is_empty()
}
pub fn len(&self) -> usize {
self.hooks.len()
}
}
pub struct ToolTimingHook {
start_times: std::sync::Mutex<HashMap<String, Instant>>,
}
impl ToolTimingHook {
pub fn new() -> Self {
Self {
start_times: std::sync::Mutex::new(HashMap::new()),
}
}
}
impl Default for ToolTimingHook {
fn default() -> Self {
Self::new()
}
}
impl Hook for ToolTimingHook {
fn on_event(&self, event: HookEvent) -> HookAction {
match event {
HookEvent::PreToolCall { name, .. } => {
let mut map = self.start_times.lock().unwrap();
map.insert(name.to_string(), Instant::now());
HookAction::Continue
}
HookEvent::PostToolCall {
name, duration_ms, ..
} => {
eprintln!("[hook] {name} took {duration_ms}ms");
HookAction::Continue
}
_ => HookAction::Continue,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
struct SkipHook;
impl Hook for SkipHook {
fn on_event(&self, event: HookEvent) -> HookAction {
match event {
HookEvent::PreToolCall { .. } => HookAction::Skip,
_ => HookAction::Continue,
}
}
}
struct ErrorHook;
impl Hook for ErrorHook {
fn on_event(&self, event: HookEvent) -> HookAction {
match event {
HookEvent::PreToolCall { .. } => HookAction::Error("nope".into()),
_ => HookAction::Continue,
}
}
}
#[test]
fn empty_registry_returns_continue() {
let reg = HookRegistry::new();
let action = reg.dispatch(HookEvent::SessionStart { goal: "test" });
assert!(matches!(action, HookAction::Continue));
}
#[test]
fn session_start_fires_with_correct_goal() {
let captured = Arc::new(std::sync::Mutex::new(String::new()));
let c = captured.clone();
struct GoalCapture(Arc<std::sync::Mutex<String>>);
impl Hook for GoalCapture {
fn on_event(&self, event: HookEvent) -> HookAction {
if let HookEvent::SessionStart { goal } = event {
*self.0.lock().unwrap() = goal.to_string();
}
HookAction::Continue
}
}
let mut reg = HookRegistry::new();
reg.register(Arc::new(GoalCapture(c)));
reg.dispatch(HookEvent::SessionStart { goal: "my goal" });
assert_eq!(*captured.lock().unwrap(), "my goal");
}
#[test]
fn pre_tool_call_skip_prevents_execution() {
let mut reg = HookRegistry::new();
reg.register(Arc::new(SkipHook));
let action = reg.dispatch(HookEvent::PreToolCall {
name: "write_file",
args: &serde_json::json!({"path": "foo.txt"}),
});
assert!(matches!(action, HookAction::Skip));
}
#[test]
fn pre_tool_call_error_returns_message() {
let mut reg = HookRegistry::new();
reg.register(Arc::new(ErrorHook));
let action = reg.dispatch(HookEvent::PreToolCall {
name: "write_file",
args: &serde_json::json!({"path": "foo.txt"}),
});
assert!(matches!(action, HookAction::Error(ref msg) if msg == "nope"));
}
#[test]
fn skip_and_error_on_non_pre_tool_are_continue() {
let mut reg = HookRegistry::new();
reg.register(Arc::new(SkipHook));
reg.register(Arc::new(ErrorHook));
let action = reg.dispatch(HookEvent::SessionStart { goal: "test" });
assert!(matches!(action, HookAction::Continue));
let action = reg.dispatch(HookEvent::PostToolCall {
name: "read_file",
args: &serde_json::json!({"path": "foo.txt"}),
result: "ok",
duration_ms: 5,
});
assert!(matches!(action, HookAction::Continue));
}
#[test]
fn multiple_hooks_fire_in_order() {
let order = Arc::new(std::sync::Mutex::new(Vec::new()));
let o1 = order.clone();
let o2 = order.clone();
struct OrdHook(usize, Arc<std::sync::Mutex<Vec<usize>>>);
impl Hook for OrdHook {
fn on_event(&self, _event: HookEvent) -> HookAction {
self.1.lock().unwrap().push(self.0);
HookAction::Continue
}
}
let mut reg = HookRegistry::new();
reg.register(Arc::new(OrdHook(1, o1)));
reg.register(Arc::new(OrdHook(2, o2)));
reg.dispatch(HookEvent::SessionStart { goal: "test" });
assert_eq!(*order.lock().unwrap(), vec![1, 2]);
}
#[test]
fn first_skip_short_circuits_remaining_hooks() {
let count = Arc::new(AtomicUsize::new(0));
let c1 = count.clone();
let c2 = count.clone();
struct FirstSkip(Arc<AtomicUsize>);
impl Hook for FirstSkip {
fn on_event(&self, event: HookEvent) -> HookAction {
match event {
HookEvent::PreToolCall { .. } => HookAction::Skip,
_ => {
self.0.fetch_add(1, Ordering::SeqCst);
HookAction::Continue
}
}
}
}
struct SecondCounter(Arc<AtomicUsize>);
impl Hook for SecondCounter {
fn on_event(&self, _event: HookEvent) -> HookAction {
self.0.fetch_add(1, Ordering::SeqCst);
HookAction::Continue
}
}
let mut reg = HookRegistry::new();
reg.register(Arc::new(FirstSkip(c1)));
reg.register(Arc::new(SecondCounter(c2.clone())));
let action = reg.dispatch(HookEvent::PreToolCall {
name: "write_file",
args: &serde_json::json!({"path": "foo.txt"}),
});
assert!(matches!(action, HookAction::Skip));
assert_eq!(c2.load(Ordering::SeqCst), 0);
}
#[test]
fn post_tool_call_receives_correct_fields() {
let captured = Arc::new(std::sync::Mutex::new(None::<(String, String, u64)>));
let c = captured.clone();
struct CaptureHook(Arc<std::sync::Mutex<Option<(String, String, u64)>>>);
impl Hook for CaptureHook {
fn on_event(&self, event: HookEvent) -> HookAction {
if let HookEvent::PostToolCall {
name,
result,
duration_ms,
..
} = event
{
*self.0.lock().unwrap() =
Some((name.to_string(), result.to_string(), duration_ms));
}
HookAction::Continue
}
}
let mut reg = HookRegistry::new();
reg.register(Arc::new(CaptureHook(c)));
reg.dispatch(HookEvent::PostToolCall {
name: "read_file",
args: &serde_json::json!({"path": "foo.txt"}),
result: "file contents",
duration_ms: 42,
});
let captured = captured.lock().unwrap().clone().unwrap();
assert_eq!(captured.0, "read_file");
assert_eq!(captured.1, "file contents");
assert_eq!(captured.2, 42);
}
#[test]
fn session_end_receives_outcome() {
let captured = Arc::new(std::sync::Mutex::new(None));
let c = captured.clone();
struct CaptureOutcome(Arc<std::sync::Mutex<Option<AgentOutcome>>>);
impl Hook for CaptureOutcome {
fn on_event(&self, event: HookEvent) -> HookAction {
if let HookEvent::SessionEnd { outcome } = event {
*self.0.lock().unwrap() = Some(outcome.clone());
}
HookAction::Continue
}
}
let outcome = AgentOutcome {
final_message: Some("done".into()),
transcript: vec![],
steps: 3,
finish: crate::agent::FinishReason::NoMoreToolCalls,
total_usage: crate::llm::TokenUsage::default(),
total_llm_latency_ms: 100,
};
let mut reg = HookRegistry::new();
reg.register(Arc::new(CaptureOutcome(c)));
reg.dispatch(HookEvent::SessionEnd { outcome: &outcome });
let captured = captured.lock().unwrap().take().unwrap();
assert_eq!(captured.final_message.as_deref(), Some("done"));
assert_eq!(captured.steps, 3);
}
#[test]
fn pre_compact_receives_transcript_len() {
let captured = Arc::new(std::sync::Mutex::new(0usize));
let c = captured.clone();
struct CaptureLen(Arc<std::sync::Mutex<usize>>);
impl Hook for CaptureLen {
fn on_event(&self, event: HookEvent) -> HookAction {
if let HookEvent::PreCompact { transcript_len } = event {
*self.0.lock().unwrap() = transcript_len;
}
HookAction::Continue
}
}
let mut reg = HookRegistry::new();
reg.register(Arc::new(CaptureLen(c)));
reg.dispatch(HookEvent::PreCompact {
transcript_len: 5000,
});
assert_eq!(*captured.lock().unwrap(), 5000);
}
#[test]
fn post_compact_receives_removed_and_summary_chars() {
let captured = Arc::new(std::sync::Mutex::new((0usize, 0usize)));
let c = captured.clone();
struct CaptureCompact(Arc<std::sync::Mutex<(usize, usize)>>);
impl Hook for CaptureCompact {
fn on_event(&self, event: HookEvent) -> HookAction {
if let HookEvent::PostCompact {
removed,
summary_chars,
} = event
{
*self.0.lock().unwrap() = (removed, summary_chars);
}
HookAction::Continue
}
}
let mut reg = HookRegistry::new();
reg.register(Arc::new(CaptureCompact(c)));
reg.dispatch(HookEvent::PostCompact {
removed: 10,
summary_chars: 200,
});
assert_eq!(captured.lock().unwrap().0, 10);
assert_eq!(captured.lock().unwrap().1, 200);
}
#[test]
fn hook_event_is_non_exhaustive() {
let _ = HookEvent::SessionStart { goal: "test" };
}
#[test]
fn tool_timing_hook_prints_to_stderr_on_post_tool_call() {
let hook = ToolTimingHook::new();
let action = hook.on_event(HookEvent::PostToolCall {
name: "read_file",
args: &serde_json::json!({"path": "foo.txt"}),
result: "ok",
duration_ms: 42,
});
assert!(matches!(action, HookAction::Continue));
}
#[test]
fn tool_timing_hook_returns_continue_for_non_tool_events() {
let hook = ToolTimingHook::new();
let action = hook.on_event(HookEvent::SessionStart { goal: "test" });
assert!(matches!(action, HookAction::Continue));
let action = hook.on_event(HookEvent::PreCompact {
transcript_len: 100,
});
assert!(matches!(action, HookAction::Continue));
let action = hook.on_event(HookEvent::PostCompact {
removed: 5,
summary_chars: 50,
});
assert!(matches!(action, HookAction::Continue));
let outcome = AgentOutcome {
final_message: Some("done".into()),
transcript: vec![],
steps: 1,
finish: crate::agent::FinishReason::NoMoreToolCalls,
total_usage: crate::llm::TokenUsage::default(),
total_llm_latency_ms: 0,
};
let action = hook.on_event(HookEvent::SessionEnd { outcome: &outcome });
assert!(matches!(action, HookAction::Continue));
}
#[test]
fn tool_timing_hook_pre_tool_call_returns_continue() {
let hook = ToolTimingHook::new();
let action = hook.on_event(HookEvent::PreToolCall {
name: "write_file",
args: &serde_json::json!({"path": "foo.txt"}),
});
assert!(matches!(action, HookAction::Continue));
}
}