Skip to main content

codetether_agent/telemetry/
mod.rs

1//! Telemetry and usage tracking
2//!
3//! Tracks token usage, costs, tool executions, file changes, and other metrics
4//! for monitoring agent performance and providing audit trails.
5
6use parking_lot::RwLock;
7use ratatui::style::Color;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::sync::Arc;
11use std::sync::atomic::{AtomicU64, Ordering};
12use std::time::{Duration, SystemTime};
13
14// ============================================================================
15// Tool Execution Tracking
16// ============================================================================
17
18/// Unique identifier for a tool execution
19pub type ToolExecId = u64;
20
21/// A single tool execution record
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct ToolExecution {
24    /// Unique execution ID
25    pub id: ToolExecId,
26    /// Tool name/identifier
27    pub tool_name: String,
28    /// Full input arguments (JSON)
29    pub input: serde_json::Value,
30    /// Output result
31    pub output: Option<String>,
32    /// Whether execution succeeded
33    pub success: bool,
34    /// Error message if failed
35    pub error: Option<String>,
36    /// Execution start time (Unix timestamp ms)
37    pub started_at: u64,
38    /// Duration in milliseconds
39    pub duration_ms: u64,
40    /// Files affected by this tool
41    pub files_affected: Vec<FileChange>,
42    /// Token usage for this execution (if applicable)
43    pub tokens: Option<TokenCounts>,
44    /// Parent execution ID (for nested/sub-agent calls)
45    pub parent_id: Option<ToolExecId>,
46    /// Session or agent ID
47    pub session_id: Option<String>,
48    /// Model used (if applicable)
49    pub model: Option<String>,
50    /// Additional metadata
51    pub metadata: HashMap<String, serde_json::Value>,
52}
53
54impl ToolExecution {
55    /// Create a new tool execution record (call this when starting execution)
56    pub fn start(tool_name: impl Into<String>, input: serde_json::Value) -> Self {
57        static COUNTER: AtomicU64 = AtomicU64::new(1);
58        let id = COUNTER.fetch_add(1, Ordering::SeqCst);
59
60        let started_at = SystemTime::now()
61            .duration_since(SystemTime::UNIX_EPOCH)
62            .map(|d| d.as_millis() as u64)
63            .unwrap_or(0);
64
65        Self {
66            id,
67            tool_name: tool_name.into(),
68            input,
69            output: None,
70            success: false,
71            error: None,
72            started_at,
73            duration_ms: 0,
74            files_affected: Vec::new(),
75            tokens: None,
76            parent_id: None,
77            session_id: None,
78            model: None,
79            metadata: HashMap::new(),
80        }
81    }
82
83    /// Complete the execution with success
84    pub fn complete_success(mut self, output: String, duration: Duration) -> Self {
85        self.output = Some(output);
86        self.success = true;
87        self.duration_ms = duration.as_millis() as u64;
88        self
89    }
90
91    /// Complete the execution with an error
92    pub fn complete_error(mut self, error: String, duration: Duration) -> Self {
93        self.error = Some(error);
94        self.success = false;
95        self.duration_ms = duration.as_millis() as u64;
96        self
97    }
98
99    /// Add a file change to this execution
100    pub fn add_file_change(&mut self, change: FileChange) {
101        self.files_affected.push(change);
102    }
103
104    /// Set parent execution ID
105    pub fn with_parent(mut self, parent_id: ToolExecId) -> Self {
106        self.parent_id = Some(parent_id);
107        self
108    }
109
110    /// Set session ID
111    pub fn with_session(mut self, session_id: impl Into<String>) -> Self {
112        self.session_id = Some(session_id.into());
113        self
114    }
115
116    /// Set model
117    pub fn with_model(mut self, model: impl Into<String>) -> Self {
118        self.model = Some(model.into());
119        self
120    }
121}
122
123/// Type of file change
124#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
125#[serde(rename_all = "snake_case")]
126pub enum FileChangeType {
127    Create,
128    Modify,
129    Delete,
130    Rename,
131    Read,
132}
133
134/// A single file change record
135#[derive(Debug, Clone, Serialize, Deserialize)]
136pub struct FileChange {
137    /// Type of change
138    pub change_type: FileChangeType,
139    /// File path (relative to workspace)
140    pub path: String,
141    /// Old path (for renames)
142    pub old_path: Option<String>,
143    /// Lines affected (start, end) - 1-indexed
144    pub lines_affected: Option<(u32, u32)>,
145    /// Content before change (for modifications)
146    pub before: Option<String>,
147    /// Content after change (for modifications/creates)
148    pub after: Option<String>,
149    /// Unified diff (if available)
150    pub diff: Option<String>,
151    /// Size in bytes before
152    pub size_before: Option<u64>,
153    /// Size in bytes after
154    pub size_after: Option<u64>,
155    /// Timestamp (Unix ms)
156    pub timestamp: u64,
157}
158
159impl FileChange {
160    /// Create a file creation record
161    pub fn create(path: impl Into<String>, content: impl Into<String>) -> Self {
162        let content = content.into();
163        let lines = content.lines().count() as u32;
164        Self {
165            change_type: FileChangeType::Create,
166            path: path.into(),
167            old_path: None,
168            lines_affected: Some((1, lines)),
169            before: None,
170            after: Some(content.clone()),
171            diff: None,
172            size_before: None,
173            size_after: Some(content.len() as u64),
174            timestamp: current_timestamp_ms(),
175        }
176    }
177
178    /// Create a file modification record
179    pub fn modify(
180        path: impl Into<String>,
181        before: impl Into<String>,
182        after: impl Into<String>,
183        lines: Option<(u32, u32)>,
184    ) -> Self {
185        let before = before.into();
186        let after = after.into();
187        Self {
188            change_type: FileChangeType::Modify,
189            path: path.into(),
190            old_path: None,
191            lines_affected: lines,
192            before: Some(before.clone()),
193            after: Some(after.clone()),
194            diff: None,
195            size_before: Some(before.len() as u64),
196            size_after: Some(after.len() as u64),
197            timestamp: current_timestamp_ms(),
198        }
199    }
200
201    /// Create a file modification with diff
202    pub fn modify_with_diff(
203        path: impl Into<String>,
204        before: impl Into<String>,
205        after: impl Into<String>,
206        diff: impl Into<String>,
207        lines: Option<(u32, u32)>,
208    ) -> Self {
209        let mut change = Self::modify(path, before, after, lines);
210        change.diff = Some(diff.into());
211        change
212    }
213
214    /// Create a file deletion record
215    pub fn delete(path: impl Into<String>, content: impl Into<String>) -> Self {
216        let content = content.into();
217        Self {
218            change_type: FileChangeType::Delete,
219            path: path.into(),
220            old_path: None,
221            lines_affected: None,
222            before: Some(content.clone()),
223            after: None,
224            diff: None,
225            size_before: Some(content.len() as u64),
226            size_after: None,
227            timestamp: current_timestamp_ms(),
228        }
229    }
230
231    /// Create a file read record (for audit trail)
232    pub fn read(path: impl Into<String>, lines: Option<(u32, u32)>) -> Self {
233        Self {
234            change_type: FileChangeType::Read,
235            path: path.into(),
236            old_path: None,
237            lines_affected: lines,
238            before: None,
239            after: None,
240            diff: None,
241            size_before: None,
242            size_after: None,
243            timestamp: current_timestamp_ms(),
244        }
245    }
246
247    /// Create a file rename record
248    pub fn rename(old_path: impl Into<String>, new_path: impl Into<String>) -> Self {
249        Self {
250            change_type: FileChangeType::Rename,
251            path: new_path.into(),
252            old_path: Some(old_path.into()),
253            lines_affected: None,
254            before: None,
255            after: None,
256            diff: None,
257            size_before: None,
258            size_after: None,
259            timestamp: current_timestamp_ms(),
260        }
261    }
262
263    /// Get a short summary of the change
264    pub fn summary(&self) -> String {
265        match self.change_type {
266            FileChangeType::Create => format!("+ {}", self.path),
267            FileChangeType::Modify => {
268                if let Some((start, end)) = self.lines_affected {
269                    format!("M {} (L{}-{})", self.path, start, end)
270                } else {
271                    format!("M {}", self.path)
272                }
273            }
274            FileChangeType::Delete => format!("- {}", self.path),
275            FileChangeType::Rename => format!(
276                "R {} -> {}",
277                self.old_path.as_deref().unwrap_or("?"),
278                self.path
279            ),
280            FileChangeType::Read => {
281                if let Some((start, end)) = self.lines_affected {
282                    format!("R {} (L{}-{})", self.path, start, end)
283                } else {
284                    format!("R {}", self.path)
285                }
286            }
287        }
288    }
289}
290
291/// Thread-safe tool execution tracker
292#[derive(Debug)]
293pub struct ToolExecutionTracker {
294    /// All executions (limited to last N)
295    executions: RwLock<Vec<ToolExecution>>,
296    /// Executions by tool name
297    by_tool: RwLock<HashMap<String, Vec<ToolExecId>>>,
298    /// Executions by file path
299    by_file: RwLock<HashMap<String, Vec<ToolExecId>>>,
300    /// Max executions to retain
301    max_executions: usize,
302    /// Total execution count (even if some evicted)
303    total_count: AtomicU64,
304    /// Total duration across all executions
305    total_duration_ms: AtomicU64,
306    /// Error count
307    error_count: AtomicU64,
308}
309
310impl ToolExecutionTracker {
311    /// Create a new tracker with default capacity (1000)
312    pub fn new() -> Self {
313        Self::with_capacity(1000)
314    }
315
316    /// Create a tracker with specified capacity
317    pub fn with_capacity(max_executions: usize) -> Self {
318        Self {
319            executions: RwLock::new(Vec::with_capacity(max_executions)),
320            by_tool: RwLock::new(HashMap::new()),
321            by_file: RwLock::new(HashMap::new()),
322            max_executions,
323            total_count: AtomicU64::new(0),
324            total_duration_ms: AtomicU64::new(0),
325            error_count: AtomicU64::new(0),
326        }
327    }
328
329    /// Record a completed tool execution
330    pub fn record(&self, execution: ToolExecution) {
331        let exec_id = execution.id;
332        let tool_name = execution.tool_name.clone();
333        let files: Vec<String> = execution
334            .files_affected
335            .iter()
336            .map(|f| f.path.clone())
337            .collect();
338        let duration = execution.duration_ms;
339        let success = execution.success;
340
341        // Update atomics
342        self.total_count.fetch_add(1, Ordering::Relaxed);
343        self.total_duration_ms
344            .fetch_add(duration, Ordering::Relaxed);
345        if !success {
346            self.error_count.fetch_add(1, Ordering::Relaxed);
347        }
348
349        // Store execution
350        {
351            let mut execs = self.executions.write();
352            if execs.len() >= self.max_executions {
353                execs.remove(0);
354            }
355            execs.push(execution);
356        }
357
358        // Index by tool
359        {
360            let mut by_tool = self.by_tool.write();
361            by_tool.entry(tool_name).or_default().push(exec_id);
362        }
363
364        // Index by file
365        {
366            let mut by_file = self.by_file.write();
367            for file in files {
368                by_file.entry(file).or_default().push(exec_id);
369            }
370        }
371    }
372
373    /// Get execution by ID
374    pub fn get(&self, id: ToolExecId) -> Option<ToolExecution> {
375        let execs = self.executions.read();
376        execs.iter().find(|e| e.id == id).cloned()
377    }
378
379    /// Get all executions for a tool
380    pub fn get_by_tool(&self, tool_name: &str) -> Vec<ToolExecution> {
381        let ids = {
382            let by_tool = self.by_tool.read();
383            by_tool.get(tool_name).cloned().unwrap_or_default()
384        };
385
386        let execs = self.executions.read();
387        ids.iter()
388            .filter_map(|id| execs.iter().find(|e| e.id == *id).cloned())
389            .collect()
390    }
391
392    /// Get all executions that affected a file
393    pub fn get_by_file(&self, path: &str) -> Vec<ToolExecution> {
394        let ids = {
395            let by_file = self.by_file.read();
396            by_file.get(path).cloned().unwrap_or_default()
397        };
398
399        let execs = self.executions.read();
400        ids.iter()
401            .filter_map(|id| execs.iter().find(|e| e.id == *id).cloned())
402            .collect()
403    }
404
405    /// Get recent executions (last N)
406    pub fn recent(&self, count: usize) -> Vec<ToolExecution> {
407        let execs = self.executions.read();
408        execs.iter().rev().take(count).cloned().collect()
409    }
410
411    /// Get all file changes across all executions
412    pub fn all_file_changes(&self) -> Vec<(ToolExecId, FileChange)> {
413        let execs = self.executions.read();
414        execs
415            .iter()
416            .flat_map(|e| e.files_affected.iter().map(|f| (e.id, f.clone())))
417            .collect()
418    }
419
420    /// Get statistics
421    pub fn stats(&self) -> ToolExecutionStats {
422        let total = self.total_count.load(Ordering::Relaxed);
423        let errors = self.error_count.load(Ordering::Relaxed);
424        let duration = self.total_duration_ms.load(Ordering::Relaxed);
425
426        let execs = self.executions.read();
427        let by_tool = self.by_tool.read();
428
429        let tool_counts: HashMap<String, u64> = by_tool
430            .iter()
431            .map(|(k, v)| (k.clone(), v.len() as u64))
432            .collect();
433
434        let file_count = execs
435            .iter()
436            .flat_map(|e| e.files_affected.iter().map(|f| f.path.clone()))
437            .collect::<std::collections::HashSet<_>>()
438            .len();
439
440        ToolExecutionStats {
441            total_executions: total,
442            successful_executions: total.saturating_sub(errors),
443            failed_executions: errors,
444            total_duration_ms: duration,
445            avg_duration_ms: if total > 0 { duration / total } else { 0 },
446            executions_by_tool: tool_counts,
447            unique_files_affected: file_count as u64,
448        }
449    }
450}
451
452impl Default for ToolExecutionTracker {
453    fn default() -> Self {
454        Self::new()
455    }
456}
457
458/// Statistics about tool executions
459#[derive(Debug, Clone, Serialize, Deserialize)]
460pub struct ToolExecutionStats {
461    pub total_executions: u64,
462    pub successful_executions: u64,
463    pub failed_executions: u64,
464    pub total_duration_ms: u64,
465    pub avg_duration_ms: u64,
466    pub executions_by_tool: HashMap<String, u64>,
467    pub unique_files_affected: u64,
468}
469
470impl ToolExecutionStats {
471    /// Format as summary string
472    pub fn summary(&self) -> String {
473        let success_rate = if self.total_executions > 0 {
474            (self.successful_executions as f64 / self.total_executions as f64) * 100.0
475        } else {
476            100.0
477        };
478
479        format!(
480            "{} executions ({:.1}% success), {} files, avg {:.0}ms",
481            self.total_executions, success_rate, self.unique_files_affected, self.avg_duration_ms
482        )
483    }
484}
485
486/// Helper to get current timestamp in milliseconds
487fn current_timestamp_ms() -> u64 {
488    SystemTime::now()
489        .duration_since(SystemTime::UNIX_EPOCH)
490        .map(|d| d.as_millis() as u64)
491        .unwrap_or(0)
492}
493
494/// Global tool execution tracker
495pub static TOOL_EXECUTIONS: once_cell::sync::Lazy<ToolExecutionTracker> =
496    once_cell::sync::Lazy::new(ToolExecutionTracker::new);
497
498// ============================================================================
499// Token Usage Tracking (existing code below)
500// ============================================================================
501
502/// Token counts for a single request/operation
503#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
504pub struct TokenCounts {
505    /// Number of tokens in the input/prompt
506    pub input: u64,
507    /// Number of tokens in the output/completion
508    pub output: u64,
509}
510
511impl TokenCounts {
512    /// Create new token counts
513    pub fn new(input: u64, output: u64) -> Self {
514        Self { input, output }
515    }
516
517    /// Total tokens (input + output)
518    pub fn total(&self) -> u64 {
519        self.input.saturating_add(self.output)
520    }
521
522    /// Add another TokenCounts to this one
523    pub fn add(&mut self, other: &TokenCounts) {
524        self.input = self.input.saturating_add(other.input);
525        self.output = self.output.saturating_add(other.output);
526    }
527}
528
529impl std::fmt::Display for TokenCounts {
530    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
531        write!(
532            f,
533            "{} in / {} out ({} total)",
534            self.input,
535            self.output,
536            self.total()
537        )
538    }
539}
540
541/// Statistics for token usage
542#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
543pub struct TokenStats {
544    /// Average tokens per request
545    pub avg_input: f64,
546    pub avg_output: f64,
547    pub avg_total: f64,
548    /// Maximum tokens in a single request
549    pub max_input: u64,
550    pub max_output: u64,
551    pub max_total: u64,
552}
553
554/// Thread-safe token usage tracker for a single model or operation type
555#[derive(Debug)]
556pub struct TokenUsageTracker {
557    /// Atomic counters for thread-safe updates
558    total_input: AtomicU64,
559    total_output: AtomicU64,
560    total_tokens: AtomicU64,
561    request_count: AtomicU64,
562    max_input: AtomicU64,
563    max_output: AtomicU64,
564    max_total: AtomicU64,
565    /// Model or operation name
566    name: String,
567}
568
569impl TokenUsageTracker {
570    /// Create a new tracker for a named model/operation
571    pub fn new(name: impl Into<String>) -> Self {
572        Self {
573            total_input: AtomicU64::new(0),
574            total_output: AtomicU64::new(0),
575            total_tokens: AtomicU64::new(0),
576            request_count: AtomicU64::new(0),
577            max_input: AtomicU64::new(0),
578            max_output: AtomicU64::new(0),
579            max_total: AtomicU64::new(0),
580            name: name.into(),
581        }
582    }
583
584    /// Record token usage for a single request
585    pub fn record(&self, input: u64, output: u64) {
586        let counts = TokenCounts::new(input, output);
587
588        // Update atomics
589        self.total_input.fetch_add(input, Ordering::Relaxed);
590        self.total_output.fetch_add(output, Ordering::Relaxed);
591        self.total_tokens
592            .fetch_add(counts.total(), Ordering::Relaxed);
593        let _count = self.request_count.fetch_add(1, Ordering::Relaxed) + 1;
594
595        // Update max values
596        self.max_input.fetch_max(input, Ordering::Relaxed);
597        self.max_output.fetch_max(output, Ordering::Relaxed);
598        self.max_total.fetch_max(counts.total(), Ordering::Relaxed);
599    }
600
601    /// Record with TokenCounts struct
602    pub fn record_counts(&self, counts: &TokenCounts) {
603        self.record(counts.input, counts.output);
604    }
605
606    /// Get current totals (fast, atomic)
607    pub fn totals(&self) -> TokenCounts {
608        TokenCounts {
609            input: self.total_input.load(Ordering::Relaxed),
610            output: self.total_output.load(Ordering::Relaxed),
611        }
612    }
613
614    /// Get the request count
615    pub fn request_count(&self) -> u64 {
616        self.request_count.load(Ordering::Relaxed)
617    }
618
619    /// Get a snapshot of current state
620    pub fn snapshot(&self) -> TokenUsageSnapshot {
621        let totals = self.totals();
622        let request_count = self.request_count();
623        let max_input = self.max_input.load(Ordering::Relaxed);
624        let max_output = self.max_output.load(Ordering::Relaxed);
625        let max_total = self.max_total.load(Ordering::Relaxed);
626
627        let stats = TokenStats {
628            avg_input: if request_count > 0 {
629                totals.input as f64 / request_count as f64
630            } else {
631                0.0
632            },
633            avg_output: if request_count > 0 {
634                totals.output as f64 / request_count as f64
635            } else {
636                0.0
637            },
638            avg_total: if request_count > 0 {
639                totals.total() as f64 / request_count as f64
640            } else {
641                0.0
642            },
643            max_input,
644            max_output,
645            max_total,
646        };
647
648        TokenUsageSnapshot {
649            name: self.name.clone(),
650            totals,
651            request_count,
652            stats,
653        }
654    }
655}
656
657/// Immutable snapshot of token usage state
658#[derive(Debug, Clone)]
659pub struct TokenUsageSnapshot {
660    pub name: String,
661    pub totals: TokenCounts,
662    pub request_count: u64,
663    pub stats: TokenStats,
664}
665
666impl TokenUsageSnapshot {
667    /// Display summary
668    pub fn summary(&self) -> String {
669        format!(
670            "{}: {} tokens ({} requests)",
671            self.name,
672            self.totals.total(),
673            self.request_count
674        )
675    }
676
677    /// Display detailed stats
678    pub fn detailed(&self) -> String {
679        format!(
680            "{}: {} tokens ({} requests)\n  Avg: {:.1} in / {:.1} out\n  Max: {} in / {} out",
681            self.name,
682            self.totals.total(),
683            self.request_count,
684            self.stats.avg_input,
685            self.stats.avg_output,
686            self.stats.max_input,
687            self.stats.max_output
688        )
689    }
690}
691
692/// Tracks usage by model, by operation type, and provides aggregation
693#[derive(Debug)]
694pub struct TokenUsageRegistry {
695    /// Trackers by model name
696    by_model: RwLock<HashMap<String, Arc<TokenUsageTracker>>>,
697    /// Trackers by operation type
698    by_operation: RwLock<HashMap<String, Arc<TokenUsageTracker>>>,
699    /// Global tracker for all usage
700    global: Arc<TokenUsageTracker>,
701}
702
703impl TokenUsageRegistry {
704    /// Create new registry
705    pub fn new() -> Self {
706        Self {
707            by_model: RwLock::new(HashMap::new()),
708            by_operation: RwLock::new(HashMap::new()),
709            global: Arc::new(TokenUsageTracker::new("global")),
710        }
711    }
712
713    /// Record usage for a specific model
714    pub fn record_model_usage(&self, model: &str, input: u64, output: u64) {
715        // Get or create tracker for this model
716        let tracker = {
717            let mut models = self.by_model.write();
718            models
719                .entry(model.to_string())
720                .or_insert_with(|| Arc::new(TokenUsageTracker::new(model)))
721                .clone()
722        };
723
724        tracker.record(input, output);
725        self.global.record(input, output);
726    }
727
728    /// Record usage for a specific operation type
729    pub fn record_operation_usage(&self, operation: &str, input: u64, output: u64) {
730        let tracker = {
731            let mut operations = self.by_operation.write();
732            operations
733                .entry(operation.to_string())
734                .or_insert_with(|| Arc::new(TokenUsageTracker::new(operation)))
735                .clone()
736        };
737
738        tracker.record(input, output);
739    }
740
741    /// Get tracker for a specific model
742    pub fn get_model_tracker(&self, model: &str) -> Option<Arc<TokenUsageTracker>> {
743        let models = self.by_model.read();
744        models.get(model).cloned()
745    }
746
747    /// Get tracker for a specific operation
748    pub fn get_operation_tracker(&self, operation: &str) -> Option<Arc<TokenUsageTracker>> {
749        let operations = self.by_operation.read();
750        operations.get(operation).cloned()
751    }
752
753    /// Get global tracker
754    pub fn global_tracker(&self) -> Arc<TokenUsageTracker> {
755        self.global.clone()
756    }
757
758    /// Get all model snapshots
759    pub fn model_snapshots(&self) -> Vec<TokenUsageSnapshot> {
760        let models = self.by_model.read();
761        models.values().map(|tracker| tracker.snapshot()).collect()
762    }
763
764    /// Get all operation snapshots
765    pub fn operation_snapshots(&self) -> Vec<TokenUsageSnapshot> {
766        let operations = self.by_operation.read();
767        operations
768            .values()
769            .map(|tracker| tracker.snapshot())
770            .collect()
771    }
772
773    /// Get global snapshot
774    pub fn global_snapshot(&self) -> TokenUsageSnapshot {
775        self.global.snapshot()
776    }
777}
778
779impl Default for TokenUsageRegistry {
780    fn default() -> Self {
781        Self::new()
782    }
783}
784
785/// Cost calculation utilities
786#[derive(Debug, Clone, Copy, Default)]
787pub struct CostEstimate {
788    pub input_cost: f64,
789    pub output_cost: f64,
790    pub total_cost: f64,
791}
792
793impl CostEstimate {
794    /// Calculate cost from token counts and pricing
795    pub fn from_tokens(
796        counts: &TokenCounts,
797        input_cost_per_million: f64,
798        output_cost_per_million: f64,
799    ) -> Self {
800        let input_cost = (counts.input as f64 / 1_000_000.0) * input_cost_per_million;
801        let output_cost = (counts.output as f64 / 1_000_000.0) * output_cost_per_million;
802        let total_cost = input_cost + output_cost;
803
804        Self {
805            input_cost,
806            output_cost,
807            total_cost,
808        }
809    }
810
811    /// Format cost as currency
812    pub fn format_currency(&self) -> String {
813        format!("${:.4}", self.total_cost)
814    }
815
816    /// Format cost with appropriate precision
817    pub fn format_smart(&self) -> String {
818        if self.total_cost < 0.0001 {
819            "< $0.0001".to_string()
820        } else if self.total_cost < 0.01 {
821            format!("${:.4}", self.total_cost)
822        } else if self.total_cost < 1.0 {
823            format!("${:.3}", self.total_cost)
824        } else {
825            format!("${:.2}", self.total_cost)
826        }
827    }
828}
829
830/// Context limit tracking
831#[derive(Debug, Clone, Copy, Default)]
832pub struct ContextLimit {
833    pub current: u64,
834    pub limit: u64,
835    pub percentage: f64,
836}
837
838impl ContextLimit {
839    /// Create new context limit info
840    pub fn new(current: u64, limit: u64) -> Self {
841        let percentage = if limit > 0 {
842            (current as f64 / limit as f64) * 100.0
843        } else {
844            0.0
845        };
846        Self {
847            current,
848            limit,
849            percentage,
850        }
851    }
852
853    /// Get warning level based on percentage
854    pub fn warning_level(&self) -> &'static str {
855        match self.percentage {
856            p if p >= 100.0 => "CRITICAL",
857            p if p >= 90.0 => "HIGH",
858            p if p >= 75.0 => "MEDIUM",
859            p if p >= 50.0 => "LOW",
860            _ => "OK",
861        }
862    }
863
864    /// Get color for warning level
865    pub fn warning_color(&self) -> Color {
866        match self.warning_level() {
867            "CRITICAL" => Color::Red,
868            "HIGH" => Color::LightRed,
869            "MEDIUM" => Color::Yellow,
870            "LOW" => Color::LightYellow,
871            _ => Color::Green,
872        }
873    }
874}
875
876/// Global token usage registry
877pub static TOKEN_USAGE: once_cell::sync::Lazy<TokenUsageRegistry> =
878    once_cell::sync::Lazy::new(TokenUsageRegistry::new);
879
880// ============================================================================
881// Telemetry Persistence
882// ============================================================================
883
884/// Persistent telemetry data that can be saved/loaded from disk
885#[derive(Debug, Clone, Serialize, Deserialize, Default)]
886pub struct TelemetryData {
887    /// Tool executions (last N)
888    pub executions: Vec<ToolExecution>,
889    /// Summary stats
890    pub stats: TelemetryStats,
891    /// Last updated timestamp
892    pub last_updated: u64,
893}
894
895/// Aggregated statistics for persistence
896#[derive(Debug, Clone, Serialize, Deserialize, Default)]
897pub struct TelemetryStats {
898    pub total_executions: u64,
899    pub successful_executions: u64,
900    pub failed_executions: u64,
901    pub total_duration_ms: u64,
902    pub total_input_tokens: u64,
903    pub total_output_tokens: u64,
904    pub total_requests: u64,
905    pub executions_by_tool: HashMap<String, u64>,
906    pub files_modified: HashMap<String, u64>,
907}
908
909impl TelemetryData {
910    /// Get the default telemetry file path
911    pub fn default_path() -> std::path::PathBuf {
912        directories::ProjectDirs::from("com", "codetether", "codetether")
913            .map(|p| p.data_dir().join("telemetry.json"))
914            .unwrap_or_else(|| std::path::PathBuf::from(".codetether/telemetry.json"))
915    }
916
917    /// Load telemetry from disk
918    pub fn load() -> Self {
919        Self::load_from(&Self::default_path())
920    }
921
922    /// Load telemetry from a specific path
923    pub fn load_from(path: &std::path::Path) -> Self {
924        match std::fs::read_to_string(path) {
925            Ok(content) => serde_json::from_str(&content).unwrap_or_default(),
926            Err(_) => Self::default(),
927        }
928    }
929
930    /// Save telemetry to disk
931    pub fn save(&self) -> std::io::Result<()> {
932        self.save_to(&Self::default_path())
933    }
934
935    /// Save telemetry to a specific path
936    pub fn save_to(&self, path: &std::path::Path) -> std::io::Result<()> {
937        if let Some(parent) = path.parent() {
938            std::fs::create_dir_all(parent)?;
939        }
940        let content = serde_json::to_string_pretty(self)?;
941        std::fs::write(path, content)
942    }
943
944    /// Add an execution and update stats
945    pub fn record_execution(&mut self, exec: ToolExecution) {
946        // Update stats
947        self.stats.total_executions += 1;
948        if exec.success {
949            self.stats.successful_executions += 1;
950        } else {
951            self.stats.failed_executions += 1;
952        }
953        self.stats.total_duration_ms += exec.duration_ms;
954
955        // Track by tool
956        *self
957            .stats
958            .executions_by_tool
959            .entry(exec.tool_name.clone())
960            .or_insert(0) += 1;
961
962        // Track files modified
963        for file in &exec.files_affected {
964            if file.change_type != FileChangeType::Read {
965                *self
966                    .stats
967                    .files_modified
968                    .entry(file.path.clone())
969                    .or_insert(0) += 1;
970            }
971        }
972
973        // Track tokens if present
974        if let Some(tokens) = &exec.tokens {
975            self.stats.total_input_tokens += tokens.input;
976            self.stats.total_output_tokens += tokens.output;
977            self.stats.total_requests += 1;
978        }
979
980        // Keep only recent executions (limit to 500)
981        if self.executions.len() >= 500 {
982            self.executions.remove(0);
983        }
984        self.executions.push(exec);
985
986        self.last_updated = current_timestamp_ms();
987    }
988
989    /// Get recent executions
990    pub fn recent(&self, count: usize) -> Vec<&ToolExecution> {
991        self.executions.iter().rev().take(count).collect()
992    }
993
994    /// Get executions by tool name
995    pub fn by_tool(&self, tool_name: &str) -> Vec<&ToolExecution> {
996        self.executions
997            .iter()
998            .filter(|e| e.tool_name == tool_name)
999            .collect()
1000    }
1001
1002    /// Get executions that affected a file
1003    pub fn by_file(&self, path: &str) -> Vec<&ToolExecution> {
1004        self.executions
1005            .iter()
1006            .filter(|e| e.files_affected.iter().any(|f| f.path == path))
1007            .collect()
1008    }
1009
1010    /// Get all file changes
1011    pub fn all_file_changes(&self) -> Vec<(ToolExecId, &FileChange)> {
1012        self.executions
1013            .iter()
1014            .flat_map(|e| e.files_affected.iter().map(move |f| (e.id, f)))
1015            .collect()
1016    }
1017
1018    /// Format as summary
1019    pub fn summary(&self) -> String {
1020        let success_rate = if self.stats.total_executions > 0 {
1021            (self.stats.successful_executions as f64 / self.stats.total_executions as f64) * 100.0
1022        } else {
1023            100.0
1024        };
1025
1026        format!(
1027            "{} executions ({:.1}% success), {} files modified, {} tokens used",
1028            self.stats.total_executions,
1029            success_rate,
1030            self.stats.files_modified.len(),
1031            self.stats.total_input_tokens + self.stats.total_output_tokens
1032        )
1033    }
1034}
1035
1036/// Global persistent telemetry (loaded on first access, saved on record)
1037pub static PERSISTENT_TELEMETRY: once_cell::sync::Lazy<RwLock<TelemetryData>> =
1038    once_cell::sync::Lazy::new(|| RwLock::new(TelemetryData::load()));
1039
1040/// Record an execution to persistent telemetry
1041pub fn record_persistent(exec: ToolExecution) {
1042    let mut data = PERSISTENT_TELEMETRY.write();
1043    data.record_execution(exec);
1044    // Best effort save - don't fail the operation if save fails
1045    let _ = data.save();
1046}
1047
1048/// Get a snapshot of persistent telemetry
1049pub fn get_persistent_stats() -> TelemetryData {
1050    PERSISTENT_TELEMETRY.read().clone()
1051}
1052
1053// ============================================================================
1054// Swarm Telemetry (for executor.rs)
1055// ============================================================================
1056
1057/// Unique identifier for a swarm execution
1058pub type SwarmExecId = String;
1059
1060/// Unique identifier for a sub-agent execution
1061pub type SubAgentExecId = u64;
1062
1063/// Privacy level for telemetry data
1064#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1065pub enum PrivacyLevel {
1066    /// Full telemetry (all data)
1067    Full,
1068    /// Reduced telemetry (metadata only, no content)
1069    #[default]
1070    Reduced,
1071    /// Minimal telemetry (counts only)
1072    Minimal,
1073    /// No telemetry
1074    None,
1075}
1076
1077/// Metrics collected during swarm execution
1078#[derive(Debug, Clone, Default)]
1079pub struct SwarmTelemetryMetrics {
1080    /// Swarm execution ID
1081    pub swarm_id: SwarmExecId,
1082    /// Number of subtasks
1083    pub subtask_count: usize,
1084    /// Execution strategy used
1085    pub strategy: String,
1086    /// Whether execution succeeded
1087    pub success: bool,
1088    /// Total execution time
1089    pub total_duration_ms: u64,
1090    /// Stage-level metrics
1091    pub stage_metrics: Vec<StageTelemetry>,
1092}
1093
1094/// Per-stage telemetry
1095#[derive(Debug, Clone, Default)]
1096pub struct StageTelemetry {
1097    /// Stage number
1098    pub stage: usize,
1099    /// Number of subagents in stage
1100    pub subagent_count: usize,
1101    /// Successful completions
1102    pub completed: usize,
1103    /// Failures
1104    pub failed: usize,
1105    /// Stage duration
1106    pub duration_ms: u64,
1107}
1108
1109/// Telemetry collector for swarm operations
1110#[derive(Debug, Default)]
1111pub struct SwarmTelemetryCollector {
1112    /// Current swarm ID
1113    swarm_id: Option<SwarmExecId>,
1114    /// Start time
1115    start_time: Option<std::time::Instant>,
1116    /// Collected metrics
1117    metrics: Option<SwarmTelemetryMetrics>,
1118}
1119
1120impl SwarmTelemetryCollector {
1121    /// Create a new telemetry collector
1122    pub fn new() -> Self {
1123        Self::default()
1124    }
1125
1126    /// Start tracking a swarm execution
1127    pub fn start_swarm(&mut self, swarm_id: SwarmExecId, subtask_count: usize, strategy: &str) {
1128        self.swarm_id = Some(swarm_id.clone());
1129        self.start_time = Some(std::time::Instant::now());
1130        self.metrics = Some(SwarmTelemetryMetrics {
1131            swarm_id,
1132            subtask_count,
1133            strategy: strategy.to_string(),
1134            ..Default::default()
1135        });
1136    }
1137
1138    /// Record latency for a specific operation
1139    pub fn record_swarm_latency(&self, _operation: &str, _duration: std::time::Duration) {
1140        // Placeholder for latency tracking
1141    }
1142
1143    /// Complete swarm tracking and return metrics
1144    pub fn complete_swarm(&mut self, success: bool) -> SwarmTelemetryMetrics {
1145        let mut metrics = self.metrics.take().unwrap_or_default();
1146        metrics.success = success;
1147        if let Some(start) = self.start_time {
1148            metrics.total_duration_ms = start.elapsed().as_millis() as u64;
1149        }
1150        metrics
1151    }
1152}