lc_agents/hooks/
mod.rs

1// lc-agents/src/hooks/mod.rs
2//! Agent Hook/Middleware system for composable lifecycle interception.
3//!
4//! Hooks allow injecting custom behavior at key points in the agent execution
5//! loop: before/after LLM calls, before/after tool calls, on stream tokens,
6//! and on errors.
7//!
8//! # Example
9//!
10//! ```rust,ignore
11//! use lc_agents::hooks::{AgentHook, ApprovalHook, ContentFilterHook};
12//! use lc_agents::AgentExecutor;
13//!
14//! let executor = AgentExecutor::new(agent, tools)
15//!     .hook(ApprovalHook::new())           // Require approval before tool calls
16//!     .hook(ContentFilterHook::new(words)); // Filter sensitive words from stream
17//! ```
18
19mod approval;
20mod content_filter;
21mod injection;
22mod logging;
23mod rate_limit;
24
25pub use approval::ApprovalHook;
26pub use content_filter::ContentFilterHook;
27pub use injection::PromptInjectionHook;
28pub use logging::LoggingHook;
29pub use rate_limit::TokenBudgetHook;
30
31use async_trait::async_trait;
32use lc_schema::Message;
33use serde_json::Value;
34use std::collections::HashMap;
35
36/// Error type for hook operations.
37#[derive(Debug, thiserror::Error)]
38#[non_exhaustive]
39pub enum HookError {
40    /// The hook rejected the operation.
41    #[error("Hook rejected: {0}")]
42    Rejected(String),
43
44    /// The hook encountered an error.
45    #[error("Hook error: {0}")]
46    Other(String),
47}
48
49/// Action to take for a completion (LLM call).
50#[derive(Debug, Clone)]
51pub enum CompletionAction {
52    /// Allow the completion to proceed.
53    Continue,
54    /// Modify the messages before the LLM call.
55    Modify {
56        /// The replacement messages to send to the LLM.
57        messages: Vec<Message>,
58    },
59    /// Reject the LLM call entirely.
60    Reject {
61        /// The reason for the rejection.
62        reason: String,
63    },
64}
65
66/// Action to take for a tool call.
67#[derive(Debug, Clone)]
68pub enum ToolCallAction {
69    /// Allow the tool call to proceed.
70    Continue,
71    /// Modify the tool call parameters.
72    Modify {
73        /// The tool name.
74        name: String,
75        /// The modified tool call arguments.
76        arguments: Value,
77    },
78    /// Reject the tool call.
79    Reject {
80        /// The reason for the rejection.
81        reason: String,
82    },
83    /// Skip this tool call (don't execute, don't error).
84    Skip,
85}
86
87/// Action to take for a stream chunk.
88#[derive(Debug, Clone)]
89pub enum StreamAction {
90    /// Forward the token to the stream.
91    Forward(String),
92    /// Filter (drop) this token.
93    Filter,
94    /// Replace the token with different content.
95    Replace(String),
96}
97
98/// Action to take on error.
99#[derive(Debug, Clone)]
100pub enum ErrorAction {
101    /// Propagate the error normally.
102    Propagate,
103    /// Retry the operation.
104    Retry,
105    /// Ignore the error and continue.
106    Ignore,
107}
108
109/// Context for a completion (LLM call) hook.
110#[derive(Debug, Clone)]
111pub struct CompletionContext {
112    /// The messages being sent to the LLM.
113    pub messages: Vec<Message>,
114    /// The model being used.
115    pub model: String,
116    /// Additional metadata.
117    pub metadata: HashMap<String, Value>,
118}
119
120/// Result context after a completion (LLM call).
121#[derive(Debug, Clone)]
122pub struct CompletionResult {
123    /// The response message from the LLM.
124    pub message: Message,
125    /// Token usage if available.
126    pub tokens_used: Option<lc_core::language_models::TokenUsage>,
127}
128
129/// Context for a tool call hook.
130#[derive(Debug, Clone)]
131pub struct ToolCallContext {
132    /// The tool name.
133    pub name: String,
134    /// The tool arguments.
135    pub arguments: Value,
136    /// The tool call ID (for function calling style).
137    pub tool_id: String,
138}
139
140/// Result context after a tool call.
141#[derive(Debug, Clone)]
142pub struct ToolResultContext {
143    /// The tool name.
144    pub name: String,
145    /// The tool result.
146    pub result: String,
147    /// The tool call ID.
148    pub tool_id: String,
149}
150
151/// Trait for agent lifecycle hooks.
152///
153/// Implement this trait to inject custom behavior at key points in the
154/// agent execution loop. All methods have default no-op implementations,
155/// so you only need to override the ones you care about.
156#[async_trait]
157pub trait AgentHook: Send + Sync {
158    /// Called before an LLM completion. Can modify messages or reject the call.
159    fn on_before_completion(&self, _ctx: &mut CompletionContext) -> CompletionAction {
160        CompletionAction::Continue
161    }
162
163    /// Called after an LLM completion. Can modify the response.
164    fn on_after_completion(&self, _ctx: &mut CompletionResult) -> Result<(), HookError> {
165        Ok(())
166    }
167
168    /// Called before a tool call. Can approve, reject, modify, or skip.
169    fn on_before_tool_call(&self, _ctx: &mut ToolCallContext) -> ToolCallAction {
170        ToolCallAction::Continue
171    }
172
173    /// Called after a tool call. Can modify the result.
174    fn on_after_tool_call(&self, _ctx: &mut ToolResultContext) -> Result<(), HookError> {
175        Ok(())
176    }
177
178    /// Called for each streaming token. Can filter, replace, or forward.
179    fn on_stream_chunk(&self, chunk: &str) -> StreamAction {
180        StreamAction::Forward(chunk.to_string())
181    }
182
183    /// Called when the agent starts execution.
184    fn on_agent_start(&self, _input: &str) -> Result<(), HookError> {
185        Ok(())
186    }
187
188    /// Called when the agent finishes execution.
189    fn on_agent_end(&self, _output: &str) -> Result<(), HookError> {
190        Ok(())
191    }
192
193    /// Called when an error occurs. Can retry, ignore, or propagate.
194    fn on_error(&self, _error: &HookError) -> ErrorAction {
195        ErrorAction::Propagate
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    #[test]
204    fn test_completion_action_default_continue() {
205        let action = CompletionAction::Continue;
206        assert!(matches!(action, CompletionAction::Continue));
207    }
208
209    #[test]
210    fn test_tool_call_action_variants() {
211        let continue_action = ToolCallAction::Continue;
212        let modify_action = ToolCallAction::Modify {
213            name: "calc".to_string(),
214            arguments: serde_json::json!({"x": 1}),
215        };
216        let reject_action = ToolCallAction::Reject {
217            reason: "not allowed".to_string(),
218        };
219        let skip_action = ToolCallAction::Skip;
220
221        assert!(matches!(continue_action, ToolCallAction::Continue));
222        assert!(matches!(modify_action, ToolCallAction::Modify { .. }));
223        assert!(matches!(reject_action, ToolCallAction::Reject { .. }));
224        assert!(matches!(skip_action, ToolCallAction::Skip));
225    }
226
227    #[test]
228    fn test_stream_action_variants() {
229        let forward = StreamAction::Forward("hello".to_string());
230        let filter = StreamAction::Filter;
231        let replace = StreamAction::Replace("[REDACTED]".to_string());
232
233        assert!(matches!(forward, StreamAction::Forward(_)));
234        assert!(matches!(filter, StreamAction::Filter));
235        assert!(matches!(replace, StreamAction::Replace(_)));
236    }
237
238    #[test]
239    fn test_error_action_variants() {
240        assert!(matches!(ErrorAction::Propagate, ErrorAction::Propagate));
241        assert!(matches!(ErrorAction::Retry, ErrorAction::Retry));
242        assert!(matches!(ErrorAction::Ignore, ErrorAction::Ignore));
243    }
244
245    #[test]
246    fn test_hook_error_display() {
247        let rejected = HookError::Rejected("not allowed".to_string());
248        assert_eq!(format!("{}", rejected), "Hook rejected: not allowed");
249
250        let other = HookError::Other("something broke".to_string());
251        assert_eq!(format!("{}", other), "Hook error: something broke");
252    }
253
254    #[test]
255    fn test_completion_context_default() {
256        let ctx = CompletionContext {
257            messages: vec![],
258            model: "gpt-4".to_string(),
259            metadata: HashMap::new(),
260        };
261        assert_eq!(ctx.model, "gpt-4");
262        assert!(ctx.messages.is_empty());
263    }
264
265    #[test]
266    fn test_tool_call_context() {
267        let ctx = ToolCallContext {
268            name: "calculator".to_string(),
269            arguments: serde_json::json!({"expr": "2+2"}),
270            tool_id: "call_123".to_string(),
271        };
272        assert_eq!(ctx.name, "calculator");
273        assert_eq!(ctx.tool_id, "call_123");
274    }
275
276    #[test]
277    fn test_tool_result_context() {
278        let ctx = ToolResultContext {
279            name: "calculator".to_string(),
280            result: "4".to_string(),
281            tool_id: "call_123".to_string(),
282        };
283        assert_eq!(ctx.result, "4");
284    }
285
286    #[test]
287    fn test_completion_result() {
288        let result = CompletionResult {
289            message: lc_schema::Message::ai("Hello!"),
290            tokens_used: None,
291        };
292        assert_eq!(result.message.content, "Hello!");
293    }
294
295    /// A custom hook that tracks all hook calls for testing.
296    struct TrackingHook {
297        before_completion_called: std::sync::atomic::AtomicBool,
298        after_completion_called: std::sync::atomic::AtomicBool,
299        before_tool_called: std::sync::atomic::AtomicBool,
300        after_tool_called: std::sync::atomic::AtomicBool,
301        agent_start_called: std::sync::atomic::AtomicBool,
302        agent_end_called: std::sync::atomic::AtomicBool,
303        error_called: std::sync::atomic::AtomicBool,
304    }
305
306    impl TrackingHook {
307        fn new() -> Self {
308            Self {
309                before_completion_called: std::sync::atomic::AtomicBool::new(false),
310                after_completion_called: std::sync::atomic::AtomicBool::new(false),
311                before_tool_called: std::sync::atomic::AtomicBool::new(false),
312                after_tool_called: std::sync::atomic::AtomicBool::new(false),
313                agent_start_called: std::sync::atomic::AtomicBool::new(false),
314                agent_end_called: std::sync::atomic::AtomicBool::new(false),
315                error_called: std::sync::atomic::AtomicBool::new(false),
316            }
317        }
318    }
319
320    #[async_trait]
321    impl AgentHook for TrackingHook {
322        fn on_before_completion(&self, _ctx: &mut CompletionContext) -> CompletionAction {
323            self.before_completion_called
324                .store(true, std::sync::atomic::Ordering::SeqCst);
325            CompletionAction::Continue
326        }
327
328        fn on_after_completion(&self, _ctx: &mut CompletionResult) -> Result<(), HookError> {
329            self.after_completion_called
330                .store(true, std::sync::atomic::Ordering::SeqCst);
331            Ok(())
332        }
333
334        fn on_before_tool_call(&self, _ctx: &mut ToolCallContext) -> ToolCallAction {
335            self.before_tool_called
336                .store(true, std::sync::atomic::Ordering::SeqCst);
337            ToolCallAction::Continue
338        }
339
340        fn on_after_tool_call(&self, _ctx: &mut ToolResultContext) -> Result<(), HookError> {
341            self.after_tool_called
342                .store(true, std::sync::atomic::Ordering::SeqCst);
343            Ok(())
344        }
345
346        fn on_agent_start(&self, _input: &str) -> Result<(), HookError> {
347            self.agent_start_called
348                .store(true, std::sync::atomic::Ordering::SeqCst);
349            Ok(())
350        }
351
352        fn on_agent_end(&self, _output: &str) -> Result<(), HookError> {
353            self.agent_end_called
354                .store(true, std::sync::atomic::Ordering::SeqCst);
355            Ok(())
356        }
357
358        fn on_error(&self, _error: &HookError) -> ErrorAction {
359            self.error_called
360                .store(true, std::sync::atomic::Ordering::SeqCst);
361            ErrorAction::Propagate
362        }
363    }
364
365    #[test]
366    fn test_custom_hook_tracking() {
367        let hook = TrackingHook::new();
368
369        // Simulate hook calls
370        let mut ctx = CompletionContext {
371            messages: vec![],
372            model: "gpt-4".to_string(),
373            metadata: HashMap::new(),
374        };
375        hook.on_before_completion(&mut ctx);
376        assert!(hook
377            .before_completion_called
378            .load(std::sync::atomic::Ordering::SeqCst));
379
380        hook.on_agent_start("test input").unwrap();
381        assert!(hook
382            .agent_start_called
383            .load(std::sync::atomic::Ordering::SeqCst));
384
385        hook.on_agent_end("test output").unwrap();
386        assert!(hook
387            .agent_end_called
388            .load(std::sync::atomic::Ordering::SeqCst));
389    }
390
391    #[test]
392    fn test_completion_action_reject() {
393        let action = CompletionAction::Reject {
394            reason: "blocked".to_string(),
395        };
396        if let CompletionAction::Reject { reason } = action {
397            assert_eq!(reason, "blocked");
398        } else {
399            panic!("Expected Reject");
400        }
401    }
402
403    #[test]
404    fn test_completion_action_modify() {
405        let action = CompletionAction::Modify {
406            messages: vec![lc_schema::Message::system("test")],
407        };
408        if let CompletionAction::Modify { messages } = action {
409            assert_eq!(messages.len(), 1);
410        } else {
411            panic!("Expected Modify");
412        }
413    }
414}
lc_agents/hooks/mod.rs

lc_agents/hooks/
mod.rs
