Skip to main content

microscope_hooks/
lib.rs

1//! Microscope Hooks - thin orchestration layer for LLM lifecycle events.
2//!
3//! The model is only the motor.
4//! Hooks are the nervous system.
5//! Microscope Memory is the memory.
6//!
7//! Six lifecycle hooks:
8//! - on_session_start  - load identity, project context, constraints
9//! - before_prompt     - inspect request, search memory, build context, inject contract
10//! - before_tool_call  - retrieve task memory, check constraints, add project context
11//! - after_tool_call   - summarize result, detect facts, prepare memory candidate
12//! - after_response   - extract durable memory, assign layer/importance, write if enabled
13//! - on_error         - store error trace, associate with task, expose recovery context
14
15use 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// ── HookEvent ────────────────────────────────────────────────────────────────
36
37#[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// ── HookContext ──────────────────────────────────────────────────────────────
72
73#[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// ── MemoryBlockRef ──────────────────────────────────────────────────────────
143
144#[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// ── MemoryCandidate ──────────────────────────────────────────────────────────
155
156#[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// ── HookConfig ────────────────────────────────────────────────────────────────
170
171#[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
245// ── HookHandler trait ────────────────────────────────────────────────────────
246
247pub trait HookHandler: Send + Sync {
248    fn event(&self) -> HookEvent;
249    fn execute(&self, ctx: HookContext) -> HookContext;
250}
251
252// ── Default Handlers ─────────────────────────────────────────────────────────
253
254pub 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
342// ── HookManager ───────────────────────────────────────────────────────────────
343
344pub 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
461// ── Security Filters ─────────────────────────────────────────────────────────
462
463fn 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
486// ── Helper: Build context package ────────────────────────────────────────────
487
488pub 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}