axon/
hooks.rs

1//! Execution hooks — pre/post step callbacks for instrumentation.
2//!
3//! Provides timing, token tracking, and metrics collection at the
4//! step and unit level. The `HookManager` accumulates events during
5//! execution and produces a summary at completion.
6//!
7//! Hook events:
8//!   UnitStart   — fired when an execution unit begins
9//!   StepStart   — fired before each step executes
10//!   StepEnd     — fired after each step completes
11//!   UnitEnd     — fired when an execution unit completes
12//!
13//! Metrics tracked:
14//!   - Per-step wall-clock duration (milliseconds)
15//!   - Per-step input/output token counts
16//!   - Per-unit aggregated timing and tokens
17//!   - Anchor breach count per step
18//!   - Total execution summary
19
20use std::time::Instant;
21
22/// A recorded step timing with associated metrics.
23#[derive(Debug, Clone)]
24pub struct StepMetrics {
25    pub unit_name: String,
26    pub step_name: String,
27    pub step_type: String,
28    pub duration_ms: u64,
29    pub input_tokens: u64,
30    pub output_tokens: u64,
31    pub anchor_breaches: u32,
32    pub chain_activations: u32,
33    pub was_retried: bool,
34}
35
36/// A recorded unit timing with aggregated metrics.
37#[derive(Debug, Clone)]
38pub struct UnitMetrics {
39    pub unit_name: String,
40    pub persona_name: String,
41    pub duration_ms: u64,
42    pub total_steps: usize,
43    pub total_input_tokens: u64,
44    pub total_output_tokens: u64,
45    pub total_anchor_breaches: u32,
46    pub total_chain_activations: u32,
47}
48
49/// Hook manager — accumulates execution metrics.
50#[derive(Debug)]
51pub struct HookManager {
52    step_metrics: Vec<StepMetrics>,
53    unit_metrics: Vec<UnitMetrics>,
54    // In-flight tracking
55    current_unit_start: Option<Instant>,
56    current_unit_name: String,
57    current_persona: String,
58    current_step_start: Option<Instant>,
59    current_step_name: String,
60    current_step_type: String,
61}
62
63impl HookManager {
64    /// Create a new hook manager.
65    pub fn new() -> Self {
66        HookManager {
67            step_metrics: Vec::new(),
68            unit_metrics: Vec::new(),
69            current_unit_start: None,
70            current_unit_name: String::new(),
71            current_persona: String::new(),
72            current_step_start: None,
73            current_step_name: String::new(),
74            current_step_type: String::new(),
75        }
76    }
77
78    /// Signal the start of an execution unit.
79    pub fn on_unit_start(&mut self, unit_name: &str, persona_name: &str) {
80        self.current_unit_start = Some(Instant::now());
81        self.current_unit_name = unit_name.to_string();
82        self.current_persona = persona_name.to_string();
83    }
84
85    /// Signal the end of an execution unit.
86    pub fn on_unit_end(&mut self) {
87        let duration_ms = self
88            .current_unit_start
89            .map(|s| s.elapsed().as_millis() as u64)
90            .unwrap_or(0);
91
92        // Aggregate step metrics for this unit
93        let unit_steps: Vec<&StepMetrics> = self
94            .step_metrics
95            .iter()
96            .filter(|s| s.unit_name == self.current_unit_name)
97            .collect();
98
99        self.unit_metrics.push(UnitMetrics {
100            unit_name: self.current_unit_name.clone(),
101            persona_name: self.current_persona.clone(),
102            duration_ms,
103            total_steps: unit_steps.len(),
104            total_input_tokens: unit_steps.iter().map(|s| s.input_tokens).sum(),
105            total_output_tokens: unit_steps.iter().map(|s| s.output_tokens).sum(),
106            total_anchor_breaches: unit_steps.iter().map(|s| s.anchor_breaches).sum(),
107            total_chain_activations: unit_steps.iter().map(|s| s.chain_activations).sum(),
108        });
109
110        self.current_unit_start = None;
111    }
112
113    /// Signal the start of a step.
114    pub fn on_step_start(&mut self, step_name: &str, step_type: &str) {
115        self.current_step_start = Some(Instant::now());
116        self.current_step_name = step_name.to_string();
117        self.current_step_type = step_type.to_string();
118    }
119
120    /// Signal the end of a step with metrics.
121    pub fn on_step_end(
122        &mut self,
123        input_tokens: u64,
124        output_tokens: u64,
125        anchor_breaches: u32,
126        chain_activations: u32,
127        was_retried: bool,
128    ) {
129        let duration_ms = self
130            .current_step_start
131            .map(|s| s.elapsed().as_millis() as u64)
132            .unwrap_or(0);
133
134        self.step_metrics.push(StepMetrics {
135            unit_name: self.current_unit_name.clone(),
136            step_name: self.current_step_name.clone(),
137            step_type: self.current_step_type.clone(),
138            duration_ms,
139            input_tokens,
140            output_tokens,
141            anchor_breaches,
142            chain_activations,
143            was_retried,
144        });
145
146        self.current_step_start = None;
147    }
148
149    /// Get all step metrics.
150    pub fn step_metrics(&self) -> &[StepMetrics] {
151        &self.step_metrics
152    }
153
154    /// Get all unit metrics.
155    pub fn unit_metrics(&self) -> &[UnitMetrics] {
156        &self.unit_metrics
157    }
158
159    /// Total execution time across all units.
160    pub fn total_duration_ms(&self) -> u64 {
161        self.unit_metrics.iter().map(|u| u.duration_ms).sum()
162    }
163
164    /// Total input tokens across all steps.
165    pub fn total_input_tokens(&self) -> u64 {
166        self.step_metrics.iter().map(|s| s.input_tokens).sum()
167    }
168
169    /// Total output tokens across all steps.
170    pub fn total_output_tokens(&self) -> u64 {
171        self.step_metrics.iter().map(|s| s.output_tokens).sum()
172    }
173
174    /// Total steps executed.
175    pub fn total_steps(&self) -> usize {
176        self.step_metrics.len()
177    }
178
179    /// Number of steps that were retried.
180    pub fn retried_steps(&self) -> usize {
181        self.step_metrics.iter().filter(|s| s.was_retried).count()
182    }
183
184    /// Slowest step by duration.
185    pub fn slowest_step(&self) -> Option<&StepMetrics> {
186        self.step_metrics.iter().max_by_key(|s| s.duration_ms)
187    }
188
189    /// Most expensive step by total tokens.
190    pub fn most_expensive_step(&self) -> Option<&StepMetrics> {
191        self.step_metrics
192            .iter()
193            .max_by_key(|s| s.input_tokens + s.output_tokens)
194    }
195
196    /// Average step duration in milliseconds.
197    pub fn avg_step_duration_ms(&self) -> u64 {
198        if self.step_metrics.is_empty() {
199            return 0;
200        }
201        let total: u64 = self.step_metrics.iter().map(|s| s.duration_ms).sum();
202        total / self.step_metrics.len() as u64
203    }
204}
205
206// ── Tests ──────────────────────────────────────────────────────────────────
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use std::thread;
212    use std::time::Duration;
213
214    #[test]
215    fn new_hook_manager_is_empty() {
216        let hm = HookManager::new();
217        assert_eq!(hm.total_steps(), 0);
218        assert_eq!(hm.total_duration_ms(), 0);
219        assert_eq!(hm.total_input_tokens(), 0);
220        assert_eq!(hm.total_output_tokens(), 0);
221        assert_eq!(hm.retried_steps(), 0);
222        assert!(hm.slowest_step().is_none());
223        assert!(hm.most_expensive_step().is_none());
224        assert_eq!(hm.avg_step_duration_ms(), 0);
225    }
226
227    #[test]
228    fn step_lifecycle() {
229        let mut hm = HookManager::new();
230        hm.on_unit_start("Flow1", "Expert");
231        hm.on_step_start("Analyze", "step");
232        // Simulate some work
233        thread::sleep(Duration::from_millis(5));
234        hm.on_step_end(100, 50, 0, 0, false);
235        hm.on_unit_end();
236
237        assert_eq!(hm.total_steps(), 1);
238        let s = &hm.step_metrics()[0];
239        assert_eq!(s.unit_name, "Flow1");
240        assert_eq!(s.step_name, "Analyze");
241        assert_eq!(s.step_type, "step");
242        assert_eq!(s.input_tokens, 100);
243        assert_eq!(s.output_tokens, 50);
244        assert!(s.duration_ms >= 4); // At least ~5ms
245        assert!(!s.was_retried);
246    }
247
248    #[test]
249    fn unit_aggregates_steps() {
250        let mut hm = HookManager::new();
251        hm.on_unit_start("Flow1", "Expert");
252
253        hm.on_step_start("Step1", "step");
254        hm.on_step_end(100, 50, 1, 0, false);
255
256        hm.on_step_start("Step2", "step");
257        hm.on_step_end(200, 100, 0, 1, true);
258
259        hm.on_unit_end();
260
261        let u = &hm.unit_metrics()[0];
262        assert_eq!(u.unit_name, "Flow1");
263        assert_eq!(u.total_steps, 2);
264        assert_eq!(u.total_input_tokens, 300);
265        assert_eq!(u.total_output_tokens, 150);
266        assert_eq!(u.total_anchor_breaches, 1);
267        assert_eq!(u.total_chain_activations, 1);
268    }
269
270    #[test]
271    fn multiple_units() {
272        let mut hm = HookManager::new();
273
274        hm.on_unit_start("Flow1", "P1");
275        hm.on_step_start("S1", "step");
276        hm.on_step_end(10, 5, 0, 0, false);
277        hm.on_unit_end();
278
279        hm.on_unit_start("Flow2", "P2");
280        hm.on_step_start("S2", "step");
281        hm.on_step_end(20, 10, 0, 0, false);
282        hm.on_unit_end();
283
284        assert_eq!(hm.unit_metrics().len(), 2);
285        assert_eq!(hm.total_steps(), 2);
286        assert_eq!(hm.total_input_tokens(), 30);
287        assert_eq!(hm.total_output_tokens(), 15);
288    }
289
290    #[test]
291    fn retried_steps_count() {
292        let mut hm = HookManager::new();
293        hm.on_unit_start("F", "P");
294        hm.on_step_start("S1", "step");
295        hm.on_step_end(10, 5, 0, 0, false);
296        hm.on_step_start("S2", "step");
297        hm.on_step_end(20, 10, 2, 0, true);
298        hm.on_step_start("S3", "step");
299        hm.on_step_end(15, 8, 0, 0, false);
300        hm.on_unit_end();
301
302        assert_eq!(hm.retried_steps(), 1);
303    }
304
305    #[test]
306    fn slowest_step() {
307        let mut hm = HookManager::new();
308        hm.on_unit_start("F", "P");
309
310        hm.on_step_start("Fast", "step");
311        hm.on_step_end(10, 5, 0, 0, false);
312
313        hm.on_step_start("Slow", "step");
314        thread::sleep(Duration::from_millis(10));
315        hm.on_step_end(10, 5, 0, 0, false);
316
317        hm.on_unit_end();
318
319        let slowest = hm.slowest_step().unwrap();
320        assert_eq!(slowest.step_name, "Slow");
321    }
322
323    #[test]
324    fn most_expensive_step() {
325        let mut hm = HookManager::new();
326        hm.on_unit_start("F", "P");
327
328        hm.on_step_start("Cheap", "step");
329        hm.on_step_end(10, 5, 0, 0, false);
330
331        hm.on_step_start("Expensive", "step");
332        hm.on_step_end(1000, 500, 0, 0, false);
333
334        hm.on_unit_end();
335
336        let expensive = hm.most_expensive_step().unwrap();
337        assert_eq!(expensive.step_name, "Expensive");
338        assert_eq!(expensive.input_tokens + expensive.output_tokens, 1500);
339    }
340
341    #[test]
342    fn avg_step_duration() {
343        let mut hm = HookManager::new();
344        hm.on_unit_start("F", "P");
345
346        // Manually create metrics to avoid timing flakiness
347        hm.step_metrics.push(StepMetrics {
348            unit_name: "F".into(),
349            step_name: "S1".into(),
350            step_type: "step".into(),
351            duration_ms: 100,
352            input_tokens: 0,
353            output_tokens: 0,
354            anchor_breaches: 0,
355            chain_activations: 0,
356            was_retried: false,
357        });
358        hm.step_metrics.push(StepMetrics {
359            unit_name: "F".into(),
360            step_name: "S2".into(),
361            step_type: "step".into(),
362            duration_ms: 200,
363            input_tokens: 0,
364            output_tokens: 0,
365            anchor_breaches: 0,
366            chain_activations: 0,
367            was_retried: false,
368        });
369
370        assert_eq!(hm.avg_step_duration_ms(), 150);
371    }
372
373    #[test]
374    fn step_with_anchor_breaches_and_chains() {
375        let mut hm = HookManager::new();
376        hm.on_unit_start("F", "P");
377        hm.on_step_start("S1", "step");
378        hm.on_step_end(100, 50, 3, 2, true);
379        hm.on_unit_end();
380
381        let s = &hm.step_metrics()[0];
382        assert_eq!(s.anchor_breaches, 3);
383        assert_eq!(s.chain_activations, 2);
384        assert!(s.was_retried);
385    }
386}
axon/hooks.rs

axon/
hooks.rs
