1use serde::{Deserialize, Serialize};
16use std::collections::HashMap;
17use std::time::{SystemTime, UNIX_EPOCH};
18
19#[derive(Debug, thiserror::Error)]
20pub enum HookError {
21 #[error("Hook '{0}' execution failed: {1}")]
22 ExecutionFailed(String, String),
23 #[error("Hook '{0}' not registered")]
24 NotRegistered(String),
25 #[error("Memory store disabled by configuration")]
26 StoreDisabled,
27 #[error("Write operations disabled by configuration")]
28 WriteDisabled,
29 #[error("Serialization error: {0}")]
30 Serialization(#[from] serde_json::Error),
31}
32
33pub type HookResult<T> = std::result::Result<T, HookError>;
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
38pub enum HookEvent {
39 SessionStart,
40 BeforePrompt,
41 BeforeToolCall,
42 AfterToolCall,
43 AfterResponse,
44 Error,
45}
46
47impl HookEvent {
48 pub fn label(&self) -> &'static str {
49 match self {
50 HookEvent::SessionStart => "on_session_start",
51 HookEvent::BeforePrompt => "before_prompt",
52 HookEvent::BeforeToolCall => "before_tool_call",
53 HookEvent::AfterToolCall => "after_tool_call",
54 HookEvent::AfterResponse => "after_response",
55 HookEvent::Error => "on_error",
56 }
57 }
58
59 pub fn all() -> [HookEvent; 6] {
60 [
61 HookEvent::SessionStart,
62 HookEvent::BeforePrompt,
63 HookEvent::BeforeToolCall,
64 HookEvent::AfterToolCall,
65 HookEvent::AfterResponse,
66 HookEvent::Error,
67 ]
68 }
69}
70
71#[derive(Debug, Clone, Serialize, Deserialize)]
74pub struct HookContext {
75 pub chain_id: String,
76 pub event: HookEvent,
77 pub timestamp_ms: u64,
78 pub query: Option<String>,
79 pub response: Option<String>,
80 pub tool_name: Option<String>,
81 pub tool_args: Option<serde_json::Value>,
82 pub tool_result: Option<String>,
83 pub error_message: Option<String>,
84 pub error_code: Option<String>,
85 pub project_context: Option<String>,
86 pub constraints: Vec<String>,
87 pub memory_contract: Option<String>,
88 pub retrieved_memories: Vec<MemoryBlockRef>,
89 pub memory_candidates: Vec<MemoryCandidate>,
90 pub metadata: HashMap<String, String>,
91}
92
93impl HookContext {
94 pub fn new(event: HookEvent) -> Self {
95 let now = SystemTime::now()
96 .duration_since(UNIX_EPOCH)
97 .unwrap_or_default()
98 .as_millis() as u64;
99 Self {
100 chain_id: uuid::Uuid::new_v4().to_string(),
101 event,
102 timestamp_ms: now,
103 query: None,
104 response: None,
105 tool_name: None,
106 tool_args: None,
107 tool_result: None,
108 error_message: None,
109 error_code: None,
110 project_context: None,
111 constraints: Vec::new(),
112 memory_contract: None,
113 retrieved_memories: Vec::new(),
114 memory_candidates: Vec::new(),
115 metadata: HashMap::new(),
116 }
117 }
118
119 pub fn with_query(mut self, query: &str) -> Self {
120 self.query = Some(query.to_string());
121 self
122 }
123
124 pub fn with_response(mut self, response: &str) -> Self {
125 self.response = Some(response.to_string());
126 self
127 }
128
129 pub fn with_tool(mut self, name: &str, args: serde_json::Value) -> Self {
130 self.tool_name = Some(name.to_string());
131 self.tool_args = Some(args);
132 self
133 }
134
135 pub fn with_error(mut self, message: &str, code: &str) -> Self {
136 self.error_message = Some(message.to_string());
137 self.error_code = Some(code.to_string());
138 self
139 }
140}
141
142#[derive(Debug, Clone, Serialize, Deserialize)]
145pub struct MemoryBlockRef {
146 pub id: String,
147 pub layer: String,
148 pub depth: u8,
149 pub importance: u8,
150 pub snippet: String,
151 pub distance: f32,
152}
153
154#[derive(Debug, Clone, Serialize, Deserialize)]
157pub struct MemoryCandidate {
158 pub text: String,
159 pub layer: String,
160 pub importance: u8,
161 pub source_event: HookEvent,
162 pub source_tool: Option<String>,
163 pub confidence: f32,
164 pub is_error: bool,
165 pub is_task: bool,
166 pub association_id: Option<String>,
167}
168
169#[derive(Debug, Clone, Serialize, Deserialize)]
172pub struct HookConfig {
173 pub enabled_hooks: HashMap<HookEvent, bool>,
174 pub write: WriteConfig,
175 pub read_only: bool,
176 pub max_candidates: usize,
177 pub min_importance: u8,
178}
179
180#[derive(Debug, Clone, Serialize, Deserialize, Default)]
181pub struct WriteConfig {
182 pub enabled: bool,
183 pub require_confirmation: bool,
184 pub filter_secrets: bool,
185 pub filter_api_keys: bool,
186}
187
188impl Default for HookConfig {
189 fn default() -> Self {
190 let mut enabled = HashMap::new();
191 enabled.insert(HookEvent::SessionStart, true);
192 enabled.insert(HookEvent::BeforePrompt, true);
193 enabled.insert(HookEvent::BeforeToolCall, true);
194 enabled.insert(HookEvent::AfterToolCall, true);
195 enabled.insert(HookEvent::AfterResponse, false);
196 enabled.insert(HookEvent::Error, true);
197 Self {
198 enabled_hooks: enabled,
199 write: WriteConfig {
200 enabled: false,
201 require_confirmation: true,
202 filter_secrets: true,
203 filter_api_keys: true,
204 },
205 read_only: false,
206 max_candidates: 5,
207 min_importance: 3,
208 }
209 }
210}
211
212impl HookConfig {
213 pub fn read_only() -> Self {
214 let mut enabled_hooks = HashMap::new();
215 enabled_hooks.insert(HookEvent::SessionStart, true);
216 enabled_hooks.insert(HookEvent::BeforePrompt, true);
217 enabled_hooks.insert(HookEvent::AfterResponse, false);
218 Self {
219 enabled_hooks,
220 write: WriteConfig {
221 enabled: false,
222 ..Default::default()
223 },
224 read_only: true,
225 ..Default::default()
226 }
227 }
228
229 pub fn full() -> Self {
230 let mut config = Self::default();
231 config.write.enabled = true;
232 config.enabled_hooks.insert(HookEvent::AfterResponse, true);
233 config
234 }
235
236 pub fn is_enabled(&self, event: &HookEvent) -> bool {
237 *self.enabled_hooks.get(event).unwrap_or(&false)
238 }
239
240 pub fn can_write(&self) -> bool {
241 !self.read_only && self.write.enabled
242 }
243}
244
245pub trait HookHandler: Send + Sync {
248 fn event(&self) -> HookEvent;
249 fn execute(&self, ctx: HookContext) -> HookContext;
250}
251
252pub fn default_on_session_start(mut ctx: HookContext) -> HookContext {
255 ctx.memory_contract = Some(include_str!("contract.txt").to_string());
256 ctx.constraints
257 .push("Never store secrets or API keys".to_string());
258 ctx.constraints
259 .push("Memory is authoritative - never invent memories".to_string());
260 ctx
261}
262
263pub fn default_before_prompt(mut ctx: HookContext) -> HookContext {
264 if ctx.memory_contract.is_none() {
265 ctx.memory_contract = Some(include_str!("contract.txt").to_string());
266 }
267 ctx
268}
269
270pub fn default_before_tool_call(mut ctx: HookContext) -> HookContext {
271 if let Some(_project) = &ctx.project_context {
272 ctx.metadata
273 .insert("project_context_active".to_string(), "true".to_string());
274 }
275 ctx
276}
277
278pub fn default_after_tool_call(mut ctx: HookContext) -> HookContext {
279 if let Some(result) = &ctx.tool_result {
280 if !result.is_empty() && result.len() > 20 {
281 let tool = ctx.tool_name.clone().unwrap_or_default();
282 ctx.memory_candidates.push(MemoryCandidate {
283 text: format!(
284 "[Tool: {}] {}",
285 tool,
286 &result.chars().take(500).collect::<String>()
287 ),
288 layer: "session".to_string(),
289 importance: 4,
290 source_event: HookEvent::AfterToolCall,
291 source_tool: ctx.tool_name.clone(),
292 confidence: 0.6,
293 is_error: false,
294 is_task: true,
295 association_id: None,
296 });
297 }
298 }
299 ctx
300}
301
302pub fn default_after_response(mut ctx: HookContext) -> HookContext {
303 if let (Some(query), Some(response)) = (&ctx.query, &ctx.response) {
304 if response.len() > 30 {
305 ctx.memory_candidates.push(MemoryCandidate {
306 text: format!(
307 "Q: {} | A: {}",
308 query,
309 &response.chars().take(300).collect::<String>()
310 ),
311 layer: "session".to_string(),
312 importance: 3,
313 source_event: HookEvent::AfterResponse,
314 source_tool: None,
315 confidence: 0.5,
316 is_error: false,
317 is_task: false,
318 association_id: None,
319 });
320 }
321 }
322 ctx
323}
324
325pub fn default_on_error(mut ctx: HookContext) -> HookContext {
326 if let (Some(msg), Some(code)) = (&ctx.error_message, &ctx.error_code) {
327 ctx.memory_candidates.push(MemoryCandidate {
328 text: format!("ERROR [{}]: {}", code, msg),
329 layer: "session".to_string(),
330 importance: 7,
331 source_event: HookEvent::Error,
332 source_tool: ctx.tool_name.clone(),
333 confidence: 0.9,
334 is_error: true,
335 is_task: true,
336 association_id: None,
337 });
338 }
339 ctx
340}
341
342pub struct HookManager {
345 config: HookConfig,
346 handlers: HashMap<HookEvent, Vec<Box<dyn HookHandler>>>,
347 use_defaults: bool,
348}
349
350impl Default for HookManager {
351 fn default() -> Self {
352 Self::new(HookConfig::default())
353 }
354}
355
356impl HookManager {
357 pub fn new(config: HookConfig) -> Self {
358 Self {
359 config,
360 handlers: HashMap::new(),
361 use_defaults: true,
362 }
363 }
364
365 pub fn read_only() -> Self {
366 Self::new(HookConfig::read_only())
367 }
368
369 pub fn register(&mut self, handler: Box<dyn HookHandler>) {
370 let event = handler.event();
371 self.handlers.entry(event).or_default().push(handler);
372 }
373
374 pub fn set_use_defaults(&mut self, use_defaults: bool) {
375 self.use_defaults = use_defaults;
376 }
377
378 pub fn config(&self) -> &HookConfig {
379 &self.config
380 }
381
382 pub fn set_config(&mut self, config: HookConfig) {
383 self.config = config;
384 }
385
386 pub fn execute(&self, event: HookEvent, ctx: HookContext) -> HookContext {
387 if !self.config.is_enabled(&event) {
388 tracing::debug!("Hook '{}' is disabled, skipping", event.label());
389 return ctx;
390 }
391
392 tracing::debug!(
393 "Executing hook '{}' (chain: {})",
394 event.label(),
395 ctx.chain_id
396 );
397 let mut current_ctx = ctx;
398
399 if self.use_defaults {
400 current_ctx = match event {
401 HookEvent::SessionStart => default_on_session_start(current_ctx),
402 HookEvent::BeforePrompt => default_before_prompt(current_ctx),
403 HookEvent::BeforeToolCall => default_before_tool_call(current_ctx),
404 HookEvent::AfterToolCall => default_after_tool_call(current_ctx),
405 HookEvent::AfterResponse => default_after_response(current_ctx),
406 HookEvent::Error => default_on_error(current_ctx),
407 };
408 }
409
410 if let Some(handlers) = self.handlers.get(&event) {
411 for handler in handlers {
412 current_ctx = handler.execute(current_ctx);
413 }
414 }
415
416 if !self.config.can_write() {
417 current_ctx.memory_candidates.clear();
418 } else {
419 if self.config.write.filter_secrets {
420 current_ctx
421 .memory_candidates
422 .retain(|c| !contains_secrets(&c.text));
423 }
424 if self.config.write.filter_api_keys {
425 current_ctx
426 .memory_candidates
427 .retain(|c| !contains_api_key(&c.text));
428 }
429 current_ctx
430 .memory_candidates
431 .retain(|c| c.importance >= self.config.min_importance);
432 current_ctx
433 .memory_candidates
434 .truncate(self.config.max_candidates);
435 }
436
437 current_ctx
438 }
439
440 pub fn execute_chain(&self, events: &[HookEvent], initial_ctx: HookContext) -> HookContext {
441 let mut ctx = initial_ctx;
442 for event in events {
443 ctx = self.execute(*event, ctx);
444 }
445 ctx
446 }
447
448 pub fn extract_candidates(&self, ctx: &HookContext) -> Vec<MemoryCandidate> {
449 ctx.memory_candidates.clone()
450 }
451
452 pub fn is_enabled(&self, event: &HookEvent) -> bool {
453 self.config.is_enabled(event)
454 }
455
456 pub fn can_write(&self) -> bool {
457 self.config.can_write()
458 }
459}
460
461fn contains_secrets(text: &str) -> bool {
464 let lower = text.to_lowercase();
465 let patterns = [
466 "password",
467 "passwd",
468 "pwd",
469 "secret",
470 "token",
471 "credential",
472 "auth_token",
473 "bearer",
474 "private_key",
475 "-----begin",
476 ];
477 patterns.iter().any(|p| lower.contains(p))
478}
479
480fn contains_api_key(text: &str) -> bool {
481 let lower = text.to_lowercase();
482 let patterns = ["api_key", "apikey", "api-key", "sk-", "pk-", "openai_key"];
483 patterns.iter().any(|p| lower.contains(p))
484}
485
486pub fn build_context_package(ctx: &HookContext) -> String {
489 let mut package = String::new();
490
491 if let Some(contract) = &ctx.memory_contract {
492 package.push_str(contract);
493 package.push_str("\n\n");
494 }
495
496 if let Some(_project) = &ctx.project_context {
497 package.push_str(&format!("## Project Context\n{}\n\n", _project));
498 }
499
500 if !ctx.constraints.is_empty() {
501 package.push_str("## Active Constraints\n");
502 for c in &ctx.constraints {
503 package.push_str(&format!("- {}\n", c));
504 }
505 package.push('\n');
506 }
507
508 if !ctx.retrieved_memories.is_empty() {
509 package.push_str(&format!(
510 "## Retrieved Memories ({} blocks)\n",
511 ctx.retrieved_memories.len()
512 ));
513 for m in &ctx.retrieved_memories {
514 package.push_str(&format!(
515 "- [{}] D{} imp={} dist={:.3}: {}\n",
516 m.layer, m.depth, m.importance, m.distance, m.snippet
517 ));
518 }
519 package.push('\n');
520 }
521
522 package
523}
524
525#[cfg(test)]
526mod tests {
527 use super::*;
528
529 #[test]
530 fn test_hook_event_labels() {
531 assert_eq!(HookEvent::SessionStart.label(), "on_session_start");
532 assert_eq!(HookEvent::BeforePrompt.label(), "before_prompt");
533 assert_eq!(HookEvent::BeforeToolCall.label(), "before_tool_call");
534 assert_eq!(HookEvent::AfterToolCall.label(), "after_tool_call");
535 assert_eq!(HookEvent::AfterResponse.label(), "after_response");
536 assert_eq!(HookEvent::Error.label(), "on_error");
537 }
538
539 #[test]
540 fn test_hook_config_default() {
541 let config = HookConfig::default();
542 assert!(config.is_enabled(&HookEvent::SessionStart));
543 assert!(config.is_enabled(&HookEvent::BeforePrompt));
544 assert!(!config.is_enabled(&HookEvent::AfterResponse));
545 assert!(!config.can_write());
546 }
547
548 #[test]
549 fn test_hook_config_read_only() {
550 let config = HookConfig::read_only();
551 assert!(config.read_only);
552 assert!(!config.can_write());
553 }
554
555 #[test]
556 fn test_hook_config_full() {
557 let config = HookConfig::full();
558 assert!(config.can_write());
559 assert!(config.is_enabled(&HookEvent::AfterResponse));
560 }
561
562 #[test]
563 fn test_hook_context_creation() {
564 let ctx = HookContext::new(HookEvent::SessionStart);
565 assert_eq!(ctx.event, HookEvent::SessionStart);
566 assert!(!ctx.chain_id.is_empty());
567 assert!(ctx.timestamp_ms > 0);
568 }
569
570 #[test]
571 fn test_hook_context_builder() {
572 let ctx = HookContext::new(HookEvent::BeforePrompt)
573 .with_query("test query")
574 .with_response("test response");
575 assert_eq!(ctx.query, Some("test query".to_string()));
576 assert_eq!(ctx.response, Some("test response".to_string()));
577 }
578
579 #[test]
580 fn test_default_on_session_start() {
581 let ctx = HookContext::new(HookEvent::SessionStart);
582 let result = default_on_session_start(ctx);
583 assert!(result.memory_contract.is_some());
584 assert!(!result.constraints.is_empty());
585 }
586
587 #[test]
588 fn test_default_after_tool_call() {
589 let ctx = HookContext::new(HookEvent::AfterToolCall)
590 .with_tool("test_tool", serde_json::json!({}));
591 let mut ctx_with_result = ctx.clone();
592 ctx_with_result.tool_result =
593 Some("This is a significant tool result with useful information.".to_string());
594 let result = default_after_tool_call(ctx_with_result);
595 assert!(!result.memory_candidates.is_empty());
596 assert_eq!(
597 result.memory_candidates[0].source_tool,
598 Some("test_tool".to_string())
599 );
600 }
601
602 #[test]
603 fn test_default_on_error() {
604 let ctx = HookContext::new(HookEvent::Error).with_error("Connection timeout", "E1001");
605 let result = default_on_error(ctx);
606 assert!(!result.memory_candidates.is_empty());
607 assert!(result.memory_candidates[0].is_error);
608 assert_eq!(result.memory_candidates[0].importance, 7);
609 }
610
611 #[test]
612 fn test_hook_manager_execute_disabled() {
613 let mut config = HookConfig::default();
614 config.enabled_hooks.insert(HookEvent::AfterResponse, false);
615 let manager = HookManager::new(config);
616 let ctx = HookContext::new(HookEvent::AfterResponse);
617 let result = manager.execute(HookEvent::AfterResponse, ctx);
618 assert!(result.memory_candidates.is_empty());
619 }
620
621 #[test]
622 fn test_hook_manager_execute_enabled() {
623 let config = HookConfig::full();
624 let manager = HookManager::new(config);
625 let ctx = HookContext::new(HookEvent::SessionStart);
626 let result = manager.execute(HookEvent::SessionStart, ctx);
627 assert!(result.memory_contract.is_some());
628 }
629
630 #[test]
631 fn test_security_filters() {
632 assert!(contains_secrets("my password is 1234"));
633 assert!(contains_api_key("api_key=sk-test123"));
634 assert!(!contains_secrets("hello world"));
635 assert!(!contains_api_key("normal text"));
636 }
637
638 #[test]
639 fn test_write_disabled_filters_candidates() {
640 let config = HookConfig::default();
641 let manager = HookManager::new(config);
642 let mut ctx = HookContext::new(HookEvent::AfterToolCall);
643 ctx.tool_result = Some("Important result".to_string());
644 ctx = default_after_tool_call(ctx);
645 let result = manager.execute(HookEvent::AfterToolCall, ctx);
646 assert!(result.memory_candidates.is_empty());
647 }
648
649 #[test]
650 fn test_build_context_package() {
651 let mut ctx = HookContext::new(HookEvent::BeforePrompt);
652 ctx.memory_contract = Some("Memory Contract".to_string());
653 ctx.project_context = Some("Project X".to_string());
654 ctx.constraints.push("Test constraint".to_string());
655 ctx.retrieved_memories.push(MemoryBlockRef {
656 id: "1".to_string(),
657 layer: "long_term".to_string(),
658 depth: 3,
659 importance: 8,
660 snippet: "Important memory".to_string(),
661 distance: 0.1,
662 });
663
664 let package = build_context_package(&ctx);
665 assert!(package.contains("Memory Contract"));
666 assert!(package.contains("Project X"));
667 assert!(package.contains("Test constraint"));
668 assert!(package.contains("Important memory"));
669 }
670
671 #[test]
672 fn test_execute_chain() {
673 let config = HookConfig::full();
674 let manager = HookManager::new(config);
675 let ctx = HookContext::new(HookEvent::SessionStart);
676 let events = [HookEvent::SessionStart, HookEvent::BeforePrompt];
677 let result = manager.execute_chain(&events, ctx);
678 assert!(result.memory_contract.is_some());
679 }
680
681 #[test]
682 fn test_custom_handler() {
683 struct TestHandler;
684 impl HookHandler for TestHandler {
685 fn event(&self) -> HookEvent {
686 HookEvent::BeforePrompt
687 }
688 fn execute(&self, mut ctx: HookContext) -> HookContext {
689 ctx.metadata.insert("custom".to_string(), "ran".to_string());
690 ctx
691 }
692 }
693
694 let mut manager = HookManager::new(HookConfig::full());
695 manager.register(Box::new(TestHandler));
696 let ctx = HookContext::new(HookEvent::BeforePrompt);
697 let result = manager.execute(HookEvent::BeforePrompt, ctx);
698 assert_eq!(result.metadata.get("custom").unwrap(), "ran");
699 }
700}