vtcode_core/core/
telemetry.rs

1use anyhow::{Context, Result};
2use hashbrown::HashMap;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::{Arc, Mutex, TryLockError};
5use std::time::{Duration, Instant};
6use vtcode_config::constants::tools;
7
8/// Aggregates telemetry data for the agent session.
9#[derive(Debug, Clone, Default)]
10pub struct TelemetryManager {
11    stats: Arc<Mutex<TelemetryStats>>,
12    start_time: Option<Instant>,
13    dropped_metric_updates: Arc<AtomicU64>,
14}
15
16#[derive(Debug, Clone, Default)]
17pub struct ModelUsageStats {
18    /// Cumulative API call duration for this model.
19    pub api_time: Duration,
20    /// Total prompt tokens consumed.
21    pub prompt_tokens: u64,
22    /// Total completion tokens generated.
23    pub completion_tokens: u64,
24    /// Total cached prompt tokens reused.
25    pub cached_prompt_tokens: u64,
26    /// Total tokens read from the prompt cache.
27    pub cache_read_tokens: u64,
28    /// Total tokens written into the prompt cache.
29    pub cache_creation_tokens: u64,
30}
31
32#[derive(Debug, Clone, Default)]
33pub struct TelemetryStats {
34    /// Total agent turns in the session.
35    pub total_turns: usize,
36    /// Total tool invocations across all tools.
37    pub total_tool_calls: usize,
38    /// Total tokens consumed (placeholder for future token tracking).
39    pub total_tokens: usize,
40    /// Per-tool invocation counts.
41    pub tool_counts: HashMap<String, usize>,
42    /// Per-tool error counts.
43    pub tool_errors: HashMap<String, usize>,
44    /// Total elapsed session duration.
45    pub session_duration: Duration,
46    /// Cumulative time spent waiting for LLM API responses.
47    pub api_time_spent: Duration,
48    /// Per-model usage breakdown.
49    pub model_usage: HashMap<String, ModelUsageStats>,
50    /// Number of metric updates dropped due to lock contention.
51    pub dropped_metric_updates: u64,
52}
53
54impl TelemetryManager {
55    /// Create a new telemetry manager and start the session timer.
56    pub fn new() -> Self {
57        Self {
58            stats: Arc::new(Mutex::new(TelemetryStats::default())),
59            start_time: Some(Instant::now()),
60            dropped_metric_updates: Arc::new(AtomicU64::new(0)),
61        }
62    }
63
64    /// Record that a new agent turn has started.
65    pub fn record_turn(&self) {
66        self.with_stats_mut_non_blocking(|stats| {
67            stats.total_turns += 1;
68            if let Some(start) = self.start_time {
69                stats.session_duration = start.elapsed();
70            }
71        });
72    }
73
74    /// Record a tool invocation, incrementing error count if it failed.
75    pub fn record_tool_usage(&self, tool: &str, success: bool) {
76        let tool = public_tool_telemetry_label(tool);
77        self.with_stats_mut_non_blocking(|stats| {
78            stats.total_tool_calls += 1;
79            if let Some(count) = stats.tool_counts.get_mut(&tool) {
80                *count += 1;
81            } else {
82                stats.tool_counts.insert(tool.clone(), 1);
83            }
84            if !success {
85                if let Some(count) = stats.tool_errors.get_mut(&tool) {
86                    *count += 1;
87                } else {
88                    stats.tool_errors.insert(tool, 1);
89                }
90            }
91        });
92    }
93
94    /// Record an LLM API request with its duration and optional token usage.
95    pub fn record_llm_request(&self, model: &str, duration: Duration, usage: Option<&crate::llm::provider::Usage>) {
96        self.with_stats_mut_non_blocking(|stats| {
97            stats.api_time_spent = stats.api_time_spent.saturating_add(duration);
98            let model_stats = if let Some(existing) = stats.model_usage.get_mut(model) {
99                existing
100            } else {
101                stats.model_usage.entry(model.to_owned()).or_default()
102            };
103            model_stats.api_time = model_stats.api_time.saturating_add(duration);
104
105            if let Some(usage) = usage {
106                model_stats.prompt_tokens = model_stats.prompt_tokens.saturating_add(usage.prompt_tokens as u64);
107                model_stats.completion_tokens =
108                    model_stats.completion_tokens.saturating_add(usage.completion_tokens as u64);
109                model_stats.cached_prompt_tokens = model_stats
110                    .cached_prompt_tokens
111                    .saturating_add(usage.cached_prompt_tokens.unwrap_or(0) as u64);
112                model_stats.cache_read_tokens = model_stats
113                    .cache_read_tokens
114                    .saturating_add(usage.cache_read_tokens_or_fallback() as u64);
115                model_stats.cache_creation_tokens = model_stats
116                    .cache_creation_tokens
117                    .saturating_add(usage.cache_creation_tokens_or_zero() as u64);
118            }
119        });
120    }
121
122    fn with_stats_mut_non_blocking<F>(&self, update: F)
123    where
124        F: FnOnce(&mut TelemetryStats),
125    {
126        match self.stats.try_lock() {
127            Ok(mut stats) => update(&mut stats),
128            Err(TryLockError::WouldBlock) | Err(TryLockError::Poisoned(_)) => {
129                self.dropped_metric_updates.fetch_add(1, Ordering::Relaxed);
130            }
131        }
132    }
133
134    /// Take a snapshot of the current telemetry statistics.
135    pub fn get_snapshot(&self) -> Result<TelemetryStats> {
136        let stats = self
137            .stats
138            .lock()
139            .map_err(|err| anyhow::anyhow!("telemetry stats lock poisoned: {err}"))
140            .context("Failed to read telemetry snapshot")?;
141        let mut snapshot = stats.clone();
142        snapshot.dropped_metric_updates = self.dropped_metric_updates.load(Ordering::Relaxed);
143        Ok(snapshot)
144    }
145}
146
147fn public_tool_telemetry_label(tool: &str) -> String {
148    match tool {
149        tools::UNIFIED_EXEC => tools::EXEC_COMMAND.to_string(),
150        tools::UNIFIED_FILE => "file_operation".to_string(),
151        _ => tool.to_owned(),
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::TelemetryManager;
158    use std::time::Duration;
159    use vtcode_config::constants::tools;
160
161    #[test]
162    fn records_llm_request_usage_per_model() {
163        let telemetry = TelemetryManager::new();
164        telemetry.record_llm_request(
165            "gpt-5",
166            Duration::from_secs(30),
167            Some(&crate::llm::provider::Usage {
168                prompt_tokens: 100,
169                completion_tokens: 200,
170                total_tokens: 300,
171                cached_prompt_tokens: Some(50),
172                cache_creation_tokens: None,
173                cache_read_tokens: None,
174                iterations: None,
175            }),
176        );
177        telemetry.record_llm_request("gpt-5", Duration::from_secs(10), None);
178
179        let snapshot = telemetry.get_snapshot().expect("snapshot");
180        assert_eq!(snapshot.api_time_spent, Duration::from_secs(40));
181        let model = snapshot.model_usage.get("gpt-5").expect("model usage");
182        assert_eq!(model.api_time, Duration::from_secs(40));
183        assert_eq!(model.prompt_tokens, 100);
184        assert_eq!(model.completion_tokens, 200);
185        assert_eq!(model.cached_prompt_tokens, 50);
186        assert_eq!(model.cache_read_tokens, 50);
187        assert_eq!(model.cache_creation_tokens, 0);
188        assert_eq!(snapshot.dropped_metric_updates, 0);
189    }
190
191    #[test]
192    fn records_cache_read_fallback_and_creation_tokens() {
193        let telemetry = TelemetryManager::new();
194        telemetry.record_llm_request(
195            "gpt-5",
196            Duration::from_secs(5),
197            Some(&crate::llm::provider::Usage {
198                prompt_tokens: 500,
199                completion_tokens: 100,
200                total_tokens: 600,
201                cached_prompt_tokens: Some(320),
202                cache_creation_tokens: Some(80),
203                cache_read_tokens: None,
204                iterations: None,
205            }),
206        );
207
208        let snapshot = telemetry.get_snapshot().expect("snapshot");
209        let model = snapshot.model_usage.get("gpt-5").expect("model usage");
210        assert_eq!(model.cache_read_tokens, 320);
211        assert_eq!(model.cache_creation_tokens, 80);
212    }
213
214    #[test]
215    fn code_search_telemetry_label_remains_canonical() {
216        let telemetry = TelemetryManager::new();
217        telemetry.record_tool_usage(tools::UNIFIED_EXEC, true);
218        telemetry.record_tool_usage(tools::CODE_SEARCH, true);
219        telemetry.record_tool_usage(tools::UNIFIED_FILE, true);
220
221        let snapshot = telemetry.get_snapshot().expect("snapshot");
222        assert_eq!(snapshot.total_tool_calls, 3);
223        assert_eq!(snapshot.tool_counts.get("exec_command"), Some(&1));
224        assert_eq!(snapshot.tool_counts.get("code_search"), Some(&1));
225        assert_eq!(snapshot.tool_counts.get("file_operation"), Some(&1));
226        assert!(snapshot.tool_errors.get("code_search").is_none());
227        assert!(!snapshot.tool_counts.keys().any(|label| label.contains("unified_")));
228        assert!(!snapshot.tool_errors.keys().any(|label| label.contains("unified_")));
229    }
230}
vtcode_core/core/telemetry.rs

vtcode_core/core/
telemetry.rs
