1use async_trait::async_trait;
9use serde::{Deserialize, Serialize};
10use std::sync::atomic::{AtomicU64, Ordering};
11use std::time::Duration;
12
13pub type Args = serde_json::Value;
15
16pub type Output = serde_json::Value;
18
19#[derive(Debug, Default)]
24pub struct ToolStats {
25 pub call_count: AtomicU64,
27 pub success_count: AtomicU64,
29 pub p50_latency_ns: AtomicU64,
32 pub peak_latency_ns: AtomicU64,
35 pub cpu_time_ns: AtomicU64,
37 pub lmdb_pages_touched: AtomicU64,
39 pub last_used_unix: AtomicU64,
41 pub effectiveness: std::sync::atomic::AtomicU32,
43}
44
45impl ToolStats {
46 pub fn record_success(&self, latency: Duration, cpu_time: Duration) {
48 self.call_count.fetch_add(1, Ordering::Relaxed);
49 self.success_count.fetch_add(1, Ordering::Relaxed);
50 let latency_ns = latency.as_nanos() as u64;
51 self.p50_latency_ns
54 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |old| {
55 Some(old / 2 + latency_ns / 2)
56 })
57 .ok();
58 self.peak_latency_ns
61 .fetch_max(latency_ns, Ordering::Relaxed);
62 self.cpu_time_ns
63 .fetch_add(cpu_time.as_nanos() as u64, Ordering::Relaxed);
64 self.last_used_unix.store(
65 std::time::SystemTime::now()
66 .duration_since(std::time::UNIX_EPOCH)
67 .unwrap_or_default()
68 .as_secs(),
69 Ordering::Relaxed,
70 );
71 self.update_effectiveness();
73 }
74
75 pub fn record_failure(&self, latency: Duration) {
77 self.call_count.fetch_add(1, Ordering::Relaxed);
78 let latency_ns = latency.as_nanos() as u64;
79 self.p50_latency_ns
80 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |old| {
81 Some(old / 2 + latency_ns / 2)
82 })
83 .ok();
84 self.peak_latency_ns
85 .fetch_max(latency_ns, Ordering::Relaxed);
86 self.last_used_unix.store(
87 std::time::SystemTime::now()
88 .duration_since(std::time::UNIX_EPOCH)
89 .unwrap_or_default()
90 .as_secs(),
91 Ordering::Relaxed,
92 );
93 self.update_effectiveness();
95 }
96
97 pub fn success_rate(&self) -> f64 {
99 let calls = self.call_count.load(Ordering::Relaxed);
100 if calls == 0 {
101 return 1.0;
102 }
103 let successes = self.success_count.load(Ordering::Relaxed);
104 successes as f64 / calls as f64
105 }
106
107 pub fn effectiveness_f32(&self) -> f32 {
109 f32::from_bits(self.effectiveness.load(Ordering::Relaxed))
110 }
111
112 pub fn set_effectiveness(&self, score: f32) {
114 self.effectiveness.store(score.to_bits(), Ordering::Relaxed);
115 }
116
117 fn update_effectiveness(&self) {
123 let calls = self.call_count.load(Ordering::Relaxed);
124 if calls == 0 {
125 return;
126 }
127 let successes = self.success_count.load(Ordering::Relaxed);
128 let rate = (successes as f32) / (calls as f32);
129 self.effectiveness.store(rate.to_bits(), Ordering::Relaxed);
130 }
131
132 pub fn should_retire(&self, min_calls: u64, threshold: f32) -> bool {
134 let calls = self.call_count.load(Ordering::Relaxed);
135 if calls < min_calls {
136 return false;
137 }
138 self.effectiveness_f32() < threshold
139 }
140
141 pub fn is_hot(&self, threshold: u64) -> bool {
143 self.call_count.load(Ordering::Relaxed) > threshold
144 }
145
146 pub fn restore(&self, snap: &ToolStatsSnapshot) {
153 self.call_count.store(snap.call_count, Ordering::Relaxed);
154 self.success_count
155 .store(snap.success_count, Ordering::Relaxed);
156 self.p50_latency_ns
157 .store(snap.p50_latency_ns, Ordering::Relaxed);
158 self.peak_latency_ns
159 .store(snap.peak_latency_ns, Ordering::Relaxed);
160 self.cpu_time_ns.store(snap.cpu_time_ns, Ordering::Relaxed);
161 self.lmdb_pages_touched
162 .store(snap.lmdb_pages_touched, Ordering::Relaxed);
163 self.last_used_unix
164 .store(snap.last_used_unix, Ordering::Relaxed);
165 self.effectiveness
166 .store(snap.effectiveness.to_bits(), Ordering::Relaxed);
167 }
168
169 pub fn snapshot(&self) -> ToolStatsSnapshot {
171 ToolStatsSnapshot {
172 call_count: self.call_count.load(Ordering::Relaxed),
173 success_count: self.success_count.load(Ordering::Relaxed),
174 p50_latency_ns: self.p50_latency_ns.load(Ordering::Relaxed),
175 peak_latency_ns: self.peak_latency_ns.load(Ordering::Relaxed),
176 cpu_time_ns: self.cpu_time_ns.load(Ordering::Relaxed),
177 lmdb_pages_touched: self.lmdb_pages_touched.load(Ordering::Relaxed),
178 last_used_unix: self.last_used_unix.load(Ordering::Relaxed),
179 effectiveness: self.effectiveness_f32(),
180 }
181 }
182}
183
184#[derive(Debug, Clone, Default, Serialize, Deserialize)]
186pub struct ToolStatsSnapshot {
187 pub call_count: u64,
189 pub success_count: u64,
191 pub p50_latency_ns: u64,
193 pub peak_latency_ns: u64,
195 pub cpu_time_ns: u64,
197 pub lmdb_pages_touched: u64,
199 pub last_used_unix: u64,
201 pub effectiveness: f32,
203}
204
205#[async_trait]
211pub trait Tool: Send + Sync {
212 fn name(&self) -> &str;
214
215 fn gana(&self) -> crate::Gana;
217
218 fn effects(&self) -> &crate::EffectRow;
220
221 async fn call(&self, ctx: &mut crate::Context, args: Args) -> crate::Result<Output>;
223
224 fn stats(&self) -> &ToolStats;
226
227 fn description(&self) -> &str {
229 self.gana().description()
230 }
231
232 fn input_schema(&self) -> serde_json::Value {
237 serde_json::json!({})
238 }
239}
240
241#[cfg(test)]
242mod tests {
243 use super::*;
244
245 #[test]
246 fn stats_record_success() {
247 let stats = ToolStats::default();
248 stats.record_success(Duration::from_millis(5), Duration::from_millis(3));
249 assert_eq!(stats.call_count.load(Ordering::Relaxed), 1);
250 assert_eq!(stats.success_count.load(Ordering::Relaxed), 1);
251 assert_eq!(stats.success_rate(), 1.0);
252 }
253
254 #[test]
255 fn stats_record_failure() {
256 let stats = ToolStats::default();
257 stats.record_success(Duration::from_millis(5), Duration::from_millis(3));
258 stats.record_failure(Duration::from_millis(2));
259 assert_eq!(stats.call_count.load(Ordering::Relaxed), 2);
260 assert_eq!(stats.success_count.load(Ordering::Relaxed), 1);
261 assert_eq!(stats.success_rate(), 0.5);
262 }
263
264 #[test]
265 fn stats_should_retire() {
266 let stats = ToolStats::default();
267 for _ in 0..2 {
269 stats.record_success(Duration::from_millis(1), Duration::from_millis(1));
270 }
271 for _ in 0..13 {
272 stats.record_failure(Duration::from_millis(1));
273 }
274 assert!(stats.should_retire(10, 0.2));
275 }
276
277 #[test]
278 fn stats_is_hot() {
279 let stats = ToolStats::default();
280 for _ in 0..1001 {
281 stats.record_success(Duration::from_millis(1), Duration::from_millis(1));
282 }
283 assert!(stats.is_hot(1000));
284 }
285
286 #[test]
287 fn stats_snapshot_restore_roundtrip() {
288 let stats = ToolStats::default();
289 stats.record_success(Duration::from_millis(5), Duration::from_millis(3));
290 stats.record_failure(Duration::from_millis(2));
291 let snap = stats.snapshot();
292
293 let restored = ToolStats::default();
294 restored.restore(&snap);
295 assert_eq!(restored.call_count.load(Ordering::Relaxed), 2);
296 assert_eq!(restored.success_count.load(Ordering::Relaxed), 1);
297 assert_eq!(
298 restored.peak_latency_ns.load(Ordering::Relaxed),
299 snap.peak_latency_ns
300 );
301 assert!((restored.effectiveness_f32() - 0.5).abs() < f32::EPSILON);
302 }
303
304 #[test]
305 fn stats_track_peak_latency() {
306 let stats = ToolStats::default();
309 stats.record_success(Duration::from_millis(10), Duration::from_millis(1));
310 assert_eq!(stats.peak_latency_ns.load(Ordering::Relaxed), 10_000_000);
311 stats.record_failure(Duration::from_millis(25));
312 assert_eq!(stats.peak_latency_ns.load(Ordering::Relaxed), 25_000_000);
313 stats.record_success(Duration::from_millis(5), Duration::from_millis(1));
314 assert_eq!(
315 stats.peak_latency_ns.load(Ordering::Relaxed),
316 25_000_000,
317 "peak latency must never decrease"
318 );
319 }
320}