use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
pub type Args = serde_json::Value;
pub type Output = serde_json::Value;
#[derive(Debug, Default)]
pub struct ToolStats {
pub call_count: AtomicU64,
pub success_count: AtomicU64,
pub p50_latency_ns: AtomicU64,
pub peak_latency_ns: AtomicU64,
pub cpu_time_ns: AtomicU64,
pub lmdb_pages_touched: AtomicU64,
pub last_used_unix: AtomicU64,
pub effectiveness: std::sync::atomic::AtomicU32,
}
impl ToolStats {
pub fn record_success(&self, latency: Duration, cpu_time: Duration) {
self.call_count.fetch_add(1, Ordering::Relaxed);
self.success_count.fetch_add(1, Ordering::Relaxed);
let latency_ns = latency.as_nanos() as u64;
self.p50_latency_ns
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |old| {
Some(old / 2 + latency_ns / 2)
})
.ok();
self.peak_latency_ns
.fetch_max(latency_ns, Ordering::Relaxed);
self.cpu_time_ns
.fetch_add(cpu_time.as_nanos() as u64, Ordering::Relaxed);
self.last_used_unix.store(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
Ordering::Relaxed,
);
self.update_effectiveness();
}
pub fn record_failure(&self, latency: Duration) {
self.call_count.fetch_add(1, Ordering::Relaxed);
let latency_ns = latency.as_nanos() as u64;
self.p50_latency_ns
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |old| {
Some(old / 2 + latency_ns / 2)
})
.ok();
self.peak_latency_ns
.fetch_max(latency_ns, Ordering::Relaxed);
self.last_used_unix.store(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
Ordering::Relaxed,
);
self.update_effectiveness();
}
pub fn success_rate(&self) -> f64 {
let calls = self.call_count.load(Ordering::Relaxed);
if calls == 0 {
return 1.0;
}
let successes = self.success_count.load(Ordering::Relaxed);
successes as f64 / calls as f64
}
pub fn effectiveness_f32(&self) -> f32 {
f32::from_bits(self.effectiveness.load(Ordering::Relaxed))
}
pub fn set_effectiveness(&self, score: f32) {
self.effectiveness.store(score.to_bits(), Ordering::Relaxed);
}
fn update_effectiveness(&self) {
let calls = self.call_count.load(Ordering::Relaxed);
if calls == 0 {
return;
}
let successes = self.success_count.load(Ordering::Relaxed);
let rate = (successes as f32) / (calls as f32);
self.effectiveness.store(rate.to_bits(), Ordering::Relaxed);
}
pub fn should_retire(&self, min_calls: u64, threshold: f32) -> bool {
let calls = self.call_count.load(Ordering::Relaxed);
if calls < min_calls {
return false;
}
self.effectiveness_f32() < threshold
}
pub fn is_hot(&self, threshold: u64) -> bool {
self.call_count.load(Ordering::Relaxed) > threshold
}
pub fn restore(&self, snap: &ToolStatsSnapshot) {
self.call_count.store(snap.call_count, Ordering::Relaxed);
self.success_count
.store(snap.success_count, Ordering::Relaxed);
self.p50_latency_ns
.store(snap.p50_latency_ns, Ordering::Relaxed);
self.peak_latency_ns
.store(snap.peak_latency_ns, Ordering::Relaxed);
self.cpu_time_ns.store(snap.cpu_time_ns, Ordering::Relaxed);
self.lmdb_pages_touched
.store(snap.lmdb_pages_touched, Ordering::Relaxed);
self.last_used_unix
.store(snap.last_used_unix, Ordering::Relaxed);
self.effectiveness
.store(snap.effectiveness.to_bits(), Ordering::Relaxed);
}
pub fn snapshot(&self) -> ToolStatsSnapshot {
ToolStatsSnapshot {
call_count: self.call_count.load(Ordering::Relaxed),
success_count: self.success_count.load(Ordering::Relaxed),
p50_latency_ns: self.p50_latency_ns.load(Ordering::Relaxed),
peak_latency_ns: self.peak_latency_ns.load(Ordering::Relaxed),
cpu_time_ns: self.cpu_time_ns.load(Ordering::Relaxed),
lmdb_pages_touched: self.lmdb_pages_touched.load(Ordering::Relaxed),
last_used_unix: self.last_used_unix.load(Ordering::Relaxed),
effectiveness: self.effectiveness_f32(),
}
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ToolStatsSnapshot {
pub call_count: u64,
pub success_count: u64,
pub p50_latency_ns: u64,
pub peak_latency_ns: u64,
pub cpu_time_ns: u64,
pub lmdb_pages_touched: u64,
pub last_used_unix: u64,
pub effectiveness: f32,
}
#[async_trait]
pub trait Tool: Send + Sync {
fn name(&self) -> &str;
fn gana(&self) -> crate::Gana;
fn effects(&self) -> &crate::EffectRow;
async fn call(&self, ctx: &mut crate::Context, args: Args) -> crate::Result<Output>;
fn stats(&self) -> &ToolStats;
fn description(&self) -> &str {
self.gana().description()
}
fn input_schema(&self) -> serde_json::Value {
serde_json::json!({})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn stats_record_success() {
let stats = ToolStats::default();
stats.record_success(Duration::from_millis(5), Duration::from_millis(3));
assert_eq!(stats.call_count.load(Ordering::Relaxed), 1);
assert_eq!(stats.success_count.load(Ordering::Relaxed), 1);
assert_eq!(stats.success_rate(), 1.0);
}
#[test]
fn stats_record_failure() {
let stats = ToolStats::default();
stats.record_success(Duration::from_millis(5), Duration::from_millis(3));
stats.record_failure(Duration::from_millis(2));
assert_eq!(stats.call_count.load(Ordering::Relaxed), 2);
assert_eq!(stats.success_count.load(Ordering::Relaxed), 1);
assert_eq!(stats.success_rate(), 0.5);
}
#[test]
fn stats_should_retire() {
let stats = ToolStats::default();
for _ in 0..2 {
stats.record_success(Duration::from_millis(1), Duration::from_millis(1));
}
for _ in 0..13 {
stats.record_failure(Duration::from_millis(1));
}
assert!(stats.should_retire(10, 0.2));
}
#[test]
fn stats_is_hot() {
let stats = ToolStats::default();
for _ in 0..1001 {
stats.record_success(Duration::from_millis(1), Duration::from_millis(1));
}
assert!(stats.is_hot(1000));
}
#[test]
fn stats_snapshot_restore_roundtrip() {
let stats = ToolStats::default();
stats.record_success(Duration::from_millis(5), Duration::from_millis(3));
stats.record_failure(Duration::from_millis(2));
let snap = stats.snapshot();
let restored = ToolStats::default();
restored.restore(&snap);
assert_eq!(restored.call_count.load(Ordering::Relaxed), 2);
assert_eq!(restored.success_count.load(Ordering::Relaxed), 1);
assert_eq!(
restored.peak_latency_ns.load(Ordering::Relaxed),
snap.peak_latency_ns
);
assert!((restored.effectiveness_f32() - 0.5).abs() < f32::EPSILON);
}
#[test]
fn stats_track_peak_latency() {
let stats = ToolStats::default();
stats.record_success(Duration::from_millis(10), Duration::from_millis(1));
assert_eq!(stats.peak_latency_ns.load(Ordering::Relaxed), 10_000_000);
stats.record_failure(Duration::from_millis(25));
assert_eq!(stats.peak_latency_ns.load(Ordering::Relaxed), 25_000_000);
stats.record_success(Duration::from_millis(5), Duration::from_millis(1));
assert_eq!(
stats.peak_latency_ns.load(Ordering::Relaxed),
25_000_000,
"peak latency must never decrease"
);
}
}