Skip to main content

agentzero_core/
metrics.rs

1use crate::types::MetricsSink;
2use serde_json::{json, Map, Value};
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex};
5
6#[derive(Debug, Default, Clone)]
7pub struct RuntimeMetrics {
8    inner: Arc<Mutex<RuntimeMetricsInner>>,
9}
10
11#[derive(Debug, Default)]
12struct RuntimeMetricsInner {
13    counters: HashMap<&'static str, u64>,
14    histograms: HashMap<&'static str, HistogramState>,
15}
16
17#[derive(Debug, Default, Clone)]
18struct HistogramState {
19    count: u64,
20    sum: f64,
21    min: f64,
22    max: f64,
23}
24
25#[derive(Debug, Clone, PartialEq)]
26pub struct HistogramSnapshot {
27    pub count: u64,
28    pub avg: f64,
29    pub min: f64,
30    pub max: f64,
31}
32
33#[derive(Debug, Clone, PartialEq)]
34pub struct RuntimeMetricsSnapshot {
35    pub counters: HashMap<String, u64>,
36    pub histograms: HashMap<String, HistogramSnapshot>,
37}
38
39impl RuntimeMetrics {
40    pub fn new() -> Self {
41        Self::default()
42    }
43
44    pub fn snapshot(&self) -> RuntimeMetricsSnapshot {
45        let inner = self.inner.lock().expect("metrics lock poisoned");
46        let counters = inner
47            .counters
48            .iter()
49            .map(|(k, v)| ((*k).to_string(), *v))
50            .collect::<HashMap<_, _>>();
51
52        let histograms = inner
53            .histograms
54            .iter()
55            .map(|(k, state)| {
56                let avg = if state.count == 0 {
57                    0.0
58                } else {
59                    state.sum / state.count as f64
60                };
61                (
62                    (*k).to_string(),
63                    HistogramSnapshot {
64                        count: state.count,
65                        avg,
66                        min: state.min,
67                        max: state.max,
68                    },
69                )
70            })
71            .collect::<HashMap<_, _>>();
72
73        RuntimeMetricsSnapshot {
74            counters,
75            histograms,
76        }
77    }
78
79    pub fn export_json(&self) -> Value {
80        let snapshot = self.snapshot();
81        let counters = snapshot
82            .counters
83            .iter()
84            .map(|(k, v)| (k.clone(), json!(v)))
85            .collect::<Map<_, _>>();
86        let histograms = snapshot
87            .histograms
88            .iter()
89            .map(|(k, h)| {
90                (
91                    k.clone(),
92                    json!({
93                        "count": h.count,
94                        "avg": h.avg,
95                        "min": h.min,
96                        "max": h.max
97                    }),
98                )
99            })
100            .collect::<Map<_, _>>();
101
102        json!({
103            "counters": counters,
104            "histograms": histograms,
105        })
106    }
107}
108
109impl MetricsSink for RuntimeMetrics {
110    fn increment_counter(&self, name: &'static str, value: u64) {
111        let mut inner = self.inner.lock().expect("metrics lock poisoned");
112        *inner.counters.entry(name).or_insert(0) += value;
113    }
114
115    fn observe_histogram(&self, name: &'static str, value: f64) {
116        let mut inner = self.inner.lock().expect("metrics lock poisoned");
117        let entry = inner.histograms.entry(name).or_default();
118        if entry.count == 0 {
119            entry.min = value;
120            entry.max = value;
121        } else {
122            entry.min = entry.min.min(value);
123            entry.max = entry.max.max(value);
124        }
125        entry.count += 1;
126        entry.sum += value;
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::RuntimeMetrics;
133    use crate::types::MetricsSink;
134
135    #[test]
136    fn runtime_metrics_collects_counters_and_histograms() {
137        let metrics = RuntimeMetrics::new();
138        metrics.increment_counter("requests_total", 1);
139        metrics.increment_counter("requests_total", 2);
140        metrics.observe_histogram("provider_latency_ms", 10.0);
141        metrics.observe_histogram("provider_latency_ms", 30.0);
142
143        let snapshot = metrics.snapshot();
144        assert_eq!(snapshot.counters.get("requests_total").copied(), Some(3));
145        let hist = snapshot
146            .histograms
147            .get("provider_latency_ms")
148            .expect("provider histogram should exist");
149        assert_eq!(hist.count, 2);
150        assert_eq!(hist.min, 10.0);
151        assert_eq!(hist.max, 30.0);
152        assert_eq!(hist.avg, 20.0);
153    }
154
155    #[test]
156    fn runtime_metrics_export_handles_empty_histograms() {
157        let metrics = RuntimeMetrics::new();
158        metrics.increment_counter("tool_errors_total", 0);
159
160        let exported = metrics.export_json();
161        assert_eq!(exported["counters"]["tool_errors_total"], 0);
162        assert!(exported["histograms"]
163            .as_object()
164            .expect("histograms should be an object")
165            .is_empty());
166    }
167}