Skip to main content

vtcode_core/tools/
execution_tracker.rs

1//! Tool execution tracking with optimized storage
2//!
3//! Tracks tool executions for observability and optimization while
4//! minimizing memory overhead through strategic data retention.
5
6use serde::{Deserialize, Serialize};
7use std::collections::VecDeque;
8use std::time::Instant;
9
10/// Lightweight execution record (optimized for memory)
11#[derive(Debug, Clone)]
12pub struct ExecutionRecord {
13    /// Tool name (stored once, referenced by ID)
14    pub tool_id: u32,
15    /// Execution status
16    pub status: ExecutionStatus,
17    /// Execution duration in milliseconds
18    pub duration_ms: u64,
19    /// Timestamp of execution
20    pub timestamp: Instant,
21    /// Whether result was cached
22    pub was_cached: bool,
23}
24
25/// Serializable version of ExecutionRecord
26#[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/// Tool execution status
35#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
36pub enum ExecutionStatus {
37    Success,
38    Failed,
39    TimedOut,
40    Cancelled,
41}
42
43/// Tool execution tracker with bounded memory usage
44pub struct ExecutionTracker {
45    /// Mapping of tool names to IDs (deduplicate strings)
46    tool_ids: Vec<String>,
47    /// Bounded execution history
48    history: VecDeque<ExecutionRecord>,
49    /// Maximum history size
50    max_history: usize,
51    /// Execution stats
52    stats: ExecutionStats,
53}
54
55/// Execution statistics
56#[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    /// Record a tool execution
77    pub fn record(
78        &mut self,
79        tool_name: &str,
80        status: ExecutionStatus,
81        duration_ms: u64,
82        was_cached: bool,
83    ) {
84        // Intern tool name to reduce memory
85        let tool_id = self.get_or_intern_tool_id(tool_name);
86
87        // Update stats
88        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        // Add record with bounded history
101        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    /// Get or create a tool ID (interns string)
116    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    /// Get tool name from ID
126    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    /// Get recent executions (bounded to prevent large allocations)
131    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    /// Get execution statistics
144    pub fn stats(&self) -> &ExecutionStats {
145        &self.stats
146    }
147
148    /// Get average execution time for a tool
149    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    /// Clear history (useful for session boundaries)
166    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}