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(
78 &mut self,
79 tool_name: &str,
80 status: ExecutionStatus,
81 duration_ms: u64,
82 was_cached: bool,
83 ) {
84 let tool_id = self.get_or_intern_tool_id(tool_name);
86
87 self.stats.total_executions += 1;
89 self.stats.total_duration_ms += duration_ms;
90 match status {
91 ExecutionStatus::Success => self.stats.successful += 1,
92 ExecutionStatus::Failed => self.stats.failed += 1,
93 ExecutionStatus::TimedOut => self.stats.timed_out += 1,
94 ExecutionStatus::Cancelled => {}
95 }
96 if was_cached {
97 self.stats.cached_hits += 1;
98 }
99
100 let record = ExecutionRecord {
102 tool_id,
103 status,
104 duration_ms,
105 timestamp: Instant::now(),
106 was_cached,
107 };
108
109 self.history.push_back(record);
110 if self.history.len() > self.max_history {
111 self.history.pop_front();
112 }
113 }
114
115 fn get_or_intern_tool_id(&mut self, tool_name: &str) -> u32 {
117 if let Some(pos) = self.tool_ids.iter().position(|t| t == tool_name) {
118 pos as u32
119 } else {
120 self.tool_ids.push(tool_name.to_string());
121 self.tool_ids.len().saturating_sub(1) as u32
122 }
123 }
124
125 pub fn get_tool_name(&self, id: u32) -> Option<&str> {
127 self.tool_ids.get(id as usize).map(|s| s.as_str())
128 }
129
130 pub fn recent_executions(&self, n: usize) -> Vec<(String, ExecutionStatus, u64)> {
132 self.history
133 .iter()
134 .rev()
135 .take(n)
136 .filter_map(|rec| {
137 self.get_tool_name(rec.tool_id)
138 .map(|name| (name.to_string(), rec.status, rec.duration_ms))
139 })
140 .collect()
141 }
142
143 pub fn stats(&self) -> &ExecutionStats {
145 &self.stats
146 }
147
148 pub fn avg_duration_for_tool(&self, tool_name: &str) -> Option<u64> {
150 let (total, count) = self
151 .history
152 .iter()
153 .filter(|rec| self.get_tool_name(rec.tool_id) == Some(tool_name))
154 .fold((0u64, 0u64), |(sum, count), rec| {
155 (sum + rec.duration_ms, count + 1)
156 });
157
158 if count == 0 {
159 return None;
160 }
161
162 Some(total / count)
163 }
164
165 pub fn clear(&mut self) {
167 self.history.clear();
168 self.stats = ExecutionStats::default();
169 }
170}
171
172#[cfg(test)]
173mod tests {
174 use super::*;
175
176 #[test]
177 fn test_execution_tracker_bounds() {
178 let mut tracker = ExecutionTracker::new(5);
179
180 for i in 0..10 {
181 tracker.record("read_file", ExecutionStatus::Success, 100, i % 2 == 0);
182 }
183
184 assert_eq!(tracker.history.len(), 5);
185 assert_eq!(tracker.stats.total_executions, 10);
186 assert_eq!(tracker.stats.cached_hits, 5);
187 }
188
189 #[test]
190 fn test_tool_id_interning() {
191 let mut tracker = ExecutionTracker::new(10);
192
193 tracker.record("read_file", ExecutionStatus::Success, 100, false);
194 tracker.record("read_file", ExecutionStatus::Success, 50, true);
195 tracker.record("write_file", ExecutionStatus::Success, 200, false);
196
197 assert_eq!(tracker.tool_ids.len(), 2);
198 }
199
200 #[test]
201 fn test_avg_duration() {
202 let mut tracker = ExecutionTracker::new(10);
203
204 tracker.record("read_file", ExecutionStatus::Success, 100, false);
205 tracker.record("read_file", ExecutionStatus::Success, 200, false);
206 tracker.record("read_file", ExecutionStatus::Success, 300, false);
207
208 assert_eq!(tracker.avg_duration_for_tool("read_file"), Some(200));
209 assert_eq!(tracker.avg_duration_for_tool("write_file"), None);
210 }
211}