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(&mut self, tool_name: &str, status: ExecutionStatus, duration_ms: u64, was_cached: bool) {
78        // Intern tool name to reduce memory
79        let tool_id = self.get_or_intern_tool_id(tool_name);
80
81        // Update stats
82        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        // Add record with bounded history
95        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    /// Get or create a tool ID (interns string)
110    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    /// Get tool name from ID
120    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    /// Get recent executions (bounded to prevent large allocations)
125    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    /// Get execution statistics
138    pub fn stats(&self) -> &ExecutionStats {
139        &self.stats
140    }
141
142    /// Get average execution time for a tool
143    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    /// Clear history (useful for session boundaries)
158    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}