1use serde::{Deserialize, Serialize};
7use std::collections::VecDeque;
8use std::time::Instant;
9
10#[derive(Debug, Clone)]
12pub struct ExecutionRecord {
13 pub tool_id: u32,
15 pub status: ExecutionStatus,
17 pub duration_ms: u64,
19 pub timestamp: Instant,
21 pub was_cached: bool,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct ExecutionRecordSnapshot {
28 pub tool_id: u32,
29 pub status: ExecutionStatus,
30 pub duration_ms: u64,
31 pub was_cached: bool,
32}
33
34#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
36pub enum ExecutionStatus {
37 Success,
38 Failed,
39 TimedOut,
40 Cancelled,
41}
42
43pub struct ExecutionTracker {
45 tool_ids: Vec<String>,
47 history: VecDeque<ExecutionRecord>,
49 max_history: usize,
51 stats: ExecutionStats,
53}
54
55#[derive(Debug, Clone, Default)]
57pub struct ExecutionStats {
58 pub total_executions: u64,
59 pub successful: u64,
60 pub failed: u64,
61 pub timed_out: u64,
62 pub cached_hits: u64,
63 pub total_duration_ms: u64,
64}
65
66impl ExecutionTracker {
67 pub fn new(max_history: usize) -> Self {
68 Self {
69 tool_ids: Vec::new(),
70 history: VecDeque::with_capacity(max_history),
71 max_history,
72 stats: ExecutionStats::default(),
73 }
74 }
75
76 pub fn record(&mut self, tool_name: &str, status: ExecutionStatus, duration_ms: u64, was_cached: bool) {
78 let tool_id = self.get_or_intern_tool_id(tool_name);
80
81 self.stats.total_executions += 1;
83 self.stats.total_duration_ms += duration_ms;
84 match status {
85 ExecutionStatus::Success => self.stats.successful += 1,
86 ExecutionStatus::Failed => self.stats.failed += 1,
87 ExecutionStatus::TimedOut => self.stats.timed_out += 1,
88 ExecutionStatus::Cancelled => {}
89 }
90 if was_cached {
91 self.stats.cached_hits += 1;
92 }
93
94 let record = ExecutionRecord {
96 tool_id,
97 status,
98 duration_ms,
99 timestamp: Instant::now(),
100 was_cached,
101 };
102
103 self.history.push_back(record);
104 if self.history.len() > self.max_history {
105 self.history.pop_front();
106 }
107 }
108
109 fn get_or_intern_tool_id(&mut self, tool_name: &str) -> u32 {
111 if let Some(pos) = self.tool_ids.iter().position(|t| t == tool_name) {
112 pos as u32
113 } else {
114 self.tool_ids.push(tool_name.to_string());
115 self.tool_ids.len().saturating_sub(1) as u32
116 }
117 }
118
119 pub fn get_tool_name(&self, id: u32) -> Option<&str> {
121 self.tool_ids.get(id as usize).map(|s| s.as_str())
122 }
123
124 pub fn recent_executions(&self, n: usize) -> Vec<(String, ExecutionStatus, u64)> {
126 self.history
127 .iter()
128 .rev()
129 .take(n)
130 .filter_map(|rec| {
131 self.get_tool_name(rec.tool_id)
132 .map(|name| (name.to_string(), rec.status, rec.duration_ms))
133 })
134 .collect()
135 }
136
137 pub fn stats(&self) -> &ExecutionStats {
139 &self.stats
140 }
141
142 pub fn avg_duration_for_tool(&self, tool_name: &str) -> Option<u64> {
144 let (total, count) = self
145 .history
146 .iter()
147 .filter(|rec| self.get_tool_name(rec.tool_id) == Some(tool_name))
148 .fold((0u64, 0u64), |(sum, count), rec| (sum + rec.duration_ms, count + 1));
149
150 if count == 0 {
151 return None;
152 }
153
154 Some(total / count)
155 }
156
157 pub fn clear(&mut self) {
159 self.history.clear();
160 self.stats = ExecutionStats::default();
161 }
162}
163
164#[cfg(test)]
165mod tests {
166 use super::*;
167
168 #[test]
169 fn test_execution_tracker_bounds() {
170 let mut tracker = ExecutionTracker::new(5);
171
172 for i in 0..10 {
173 tracker.record("read_file", ExecutionStatus::Success, 100, i % 2 == 0);
174 }
175
176 assert_eq!(tracker.history.len(), 5);
177 assert_eq!(tracker.stats.total_executions, 10);
178 assert_eq!(tracker.stats.cached_hits, 5);
179 }
180
181 #[test]
182 fn test_tool_id_interning() {
183 let mut tracker = ExecutionTracker::new(10);
184
185 tracker.record("read_file", ExecutionStatus::Success, 100, false);
186 tracker.record("read_file", ExecutionStatus::Success, 50, true);
187 tracker.record("write_file", ExecutionStatus::Success, 200, false);
188
189 assert_eq!(tracker.tool_ids.len(), 2);
190 }
191
192 #[test]
193 fn test_avg_duration() {
194 let mut tracker = ExecutionTracker::new(10);
195
196 tracker.record("read_file", ExecutionStatus::Success, 100, false);
197 tracker.record("read_file", ExecutionStatus::Success, 200, false);
198 tracker.record("read_file", ExecutionStatus::Success, 300, false);
199
200 assert_eq!(tracker.avg_duration_for_tool("read_file"), Some(200));
201 assert_eq!(tracker.avg_duration_for_tool("write_file"), None);
202 }
203}