1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243//! Plugin lifecycle hooks — intercept points at key moments in agent execution.
//!
//! Provides a callback-based hook system (not dynamic loading) for safe extensibility.
//! Four hook types:
//! - `BeforeToolCall`: Fires before tool execution. Can block the call by returning Err.
//! - `AfterToolCall`: Fires after tool execution. Observe-only.
//! - `BeforePromptBuild`: Fires before system prompt construction. Observe-only.
//! - `AgentLoopEnd`: Fires after the agent loop completes. Observe-only.
use dashmap::DashMap;
use openfang_types::agent::HookEvent;
use std::sync::Arc;
/// Context passed to hook handlers.
pub struct HookContext<'a> {
/// Agent display name.
pub agent_name: &'a str,
/// Agent ID string.
pub agent_id: &'a str,
/// Which hook event triggered this call.
pub event: HookEvent,
/// Event-specific payload (tool name, input, result, etc.).
pub data: serde_json::Value,
}
/// Hook handler trait. Implementations must be thread-safe.
pub trait HookHandler: Send + Sync {
/// Called when the hook fires.
///
/// For `BeforeToolCall`: returning `Err(reason)` blocks the tool call.
/// For all other events: return value is ignored (observe-only).
fn on_event(&self, ctx: &HookContext) -> Result<(), String>;
}
/// Registry of hook handlers, keyed by event type.
///
/// Thread-safe via `DashMap`. Handlers fire in registration order.
pub struct HookRegistry {
handlers: DashMap<HookEvent, Vec<Arc<dyn HookHandler>>>,
}
impl HookRegistry {
/// Create an empty hook registry.
pub fn new() -> Self {
Self {
handlers: DashMap::new(),
}
}
/// Register a handler for a specific event type.
pub fn register(&self, event: HookEvent, handler: Arc<dyn HookHandler>) {
self.handlers.entry(event).or_default().push(handler);
}
/// Fire all handlers for an event. Returns Err if any handler blocks.
///
/// For `BeforeToolCall`, the first Err stops execution and returns the reason.
/// For other events, errors are logged but don't propagate.
pub fn fire(&self, ctx: &HookContext) -> Result<(), String> {
if let Some(handlers) = self.handlers.get(&ctx.event) {
for handler in handlers.iter() {
if let Err(reason) = handler.on_event(ctx) {
if ctx.event == HookEvent::BeforeToolCall {
return Err(reason);
}
// For non-blocking hooks, log and continue
tracing::warn!(
event = ?ctx.event,
agent = ctx.agent_name,
error = %reason,
"Hook handler returned error (non-blocking)"
);
}
}
}
Ok(())
}
/// Check if any handlers are registered for a given event.
pub fn has_handlers(&self, event: HookEvent) -> bool {
self.handlers
.get(&event)
.map(|v| !v.is_empty())
.unwrap_or(false)
}
}
impl Default for HookRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
/// A test handler that always succeeds.
struct OkHandler;
impl HookHandler for OkHandler {
fn on_event(&self, _ctx: &HookContext) -> Result<(), String> {
Ok(())
}
}
/// A test handler that always blocks.
struct BlockHandler {
reason: String,
}
impl HookHandler for BlockHandler {
fn on_event(&self, _ctx: &HookContext) -> Result<(), String> {
Err(self.reason.clone())
}
}
/// A test handler that records calls.
struct RecordHandler {
calls: std::sync::Mutex<Vec<String>>,
}
impl RecordHandler {
fn new() -> Self {
Self {
calls: std::sync::Mutex::new(Vec::new()),
}
}
fn call_count(&self) -> usize {
self.calls.lock().unwrap().len()
}
}
impl HookHandler for RecordHandler {
fn on_event(&self, ctx: &HookContext) -> Result<(), String> {
self.calls.lock().unwrap().push(format!("{:?}", ctx.event));
Ok(())
}
}
fn make_ctx(event: HookEvent) -> HookContext<'static> {
HookContext {
agent_name: "test-agent",
agent_id: "abc-123",
event,
data: serde_json::json!({}),
}
}
#[test]
fn test_empty_registry_is_noop() {
let registry = HookRegistry::new();
let ctx = make_ctx(HookEvent::BeforeToolCall);
assert!(registry.fire(&ctx).is_ok());
}
#[test]
fn test_before_tool_call_can_block() {
let registry = HookRegistry::new();
registry.register(
HookEvent::BeforeToolCall,
Arc::new(BlockHandler {
reason: "Not allowed".to_string(),
}),
);
let ctx = make_ctx(HookEvent::BeforeToolCall);
let result = registry.fire(&ctx);
assert!(result.is_err());
assert_eq!(result.unwrap_err(), "Not allowed");
}
#[test]
fn test_after_tool_call_receives_result() {
let recorder = Arc::new(RecordHandler::new());
let registry = HookRegistry::new();
registry.register(HookEvent::AfterToolCall, recorder.clone());
let ctx = HookContext {
agent_name: "test-agent",
agent_id: "abc-123",
event: HookEvent::AfterToolCall,
data: serde_json::json!({"tool_name": "file_read", "result": "ok"}),
};
assert!(registry.fire(&ctx).is_ok());
assert_eq!(recorder.call_count(), 1);
}
#[test]
fn test_multiple_handlers_all_fire() {
let r1 = Arc::new(RecordHandler::new());
let r2 = Arc::new(RecordHandler::new());
let registry = HookRegistry::new();
registry.register(HookEvent::AgentLoopEnd, r1.clone());
registry.register(HookEvent::AgentLoopEnd, r2.clone());
let ctx = make_ctx(HookEvent::AgentLoopEnd);
assert!(registry.fire(&ctx).is_ok());
assert_eq!(r1.call_count(), 1);
assert_eq!(r2.call_count(), 1);
}
#[test]
fn test_hook_errors_dont_crash_non_blocking() {
let registry = HookRegistry::new();
// Register a blocking handler for a non-blocking event
registry.register(
HookEvent::AfterToolCall,
Arc::new(BlockHandler {
reason: "oops".to_string(),
}),
);
let ctx = make_ctx(HookEvent::AfterToolCall);
// AfterToolCall is non-blocking, so error should be swallowed
assert!(registry.fire(&ctx).is_ok());
}
#[test]
fn test_all_four_events_fire() {
let recorder = Arc::new(RecordHandler::new());
let registry = HookRegistry::new();
registry.register(HookEvent::BeforeToolCall, recorder.clone());
registry.register(HookEvent::AfterToolCall, recorder.clone());
registry.register(HookEvent::BeforePromptBuild, recorder.clone());
registry.register(HookEvent::AgentLoopEnd, recorder.clone());
for event in [
HookEvent::BeforeToolCall,
HookEvent::AfterToolCall,
HookEvent::BeforePromptBuild,
HookEvent::AgentLoopEnd,
] {
let ctx = make_ctx(event);
let _ = registry.fire(&ctx);
}
assert_eq!(recorder.call_count(), 4);
}
#[test]
fn test_has_handlers() {
let registry = HookRegistry::new();
assert!(!registry.has_handlers(HookEvent::BeforeToolCall));
registry.register(HookEvent::BeforeToolCall, Arc::new(OkHandler));
assert!(registry.has_handlers(HookEvent::BeforeToolCall));
assert!(!registry.has_handlers(HookEvent::AfterToolCall));
}
}