poimen-memory/crates/mem-cli/src/agent/observability.rs

//! Observability and Metrics

use serde::{Deserialize, Serialize};
use std::sync::{Arc, RwLock};
use std::collections::HashMap;

/// Agent metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentMetrics {
    pub agent_id: String,
    pub requests_total: u64,
    pub requests_success: u64,
    pub requests_failed: u64,
    pub average_latency_ms: f32,
    pub p95_latency_ms: f32,
    pub p99_latency_ms: f32,
    pub capabilities_used: HashMap<String, u64>,
    pub last_updated: String,
}

impl Default for AgentMetrics {
    fn default() -> Self {
        AgentMetrics {
            agent_id: "unknown".to_string(),
            requests_total: 0,
            requests_success: 0,
            requests_failed: 0,
            average_latency_ms: 0.0,
            p95_latency_ms: 0.0,
            p99_latency_ms: 0.0,
            capabilities_used: HashMap::new(),
            last_updated: chrono::Utc::now().to_rfc3339(),
        }
    }
}

/// Metrics collector (thread-safe with RwLock for better read concurrency)
pub struct MetricsCollector {
    metrics: Arc<std::sync::RwLock<HashMap<String, AgentMetrics>>>,
    latencies: Arc<std::sync::RwLock<HashMap<String, Vec<f32>>>>,
}

impl MetricsCollector {
    pub fn new() -> Self {
        MetricsCollector {
            metrics: Arc::new(RwLock::new(HashMap::new())),
            latencies: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Record request
    pub fn record_request(
        &self,
        agent_id: &str,
        success: bool,
        latency_ms: f32,
        capability: Option<&str>,
    ) {
        let mut metrics = self.metrics.write().unwrap();
        let mut lats = self.latencies.write().unwrap();

        let metric = metrics
            .entry(agent_id.to_string())
            .or_insert_with(|| AgentMetrics {
                agent_id: agent_id.to_string(),
                ..Default::default()
            });

        metric.requests_total += 1;
        if success {
            metric.requests_success += 1;
        } else {
            metric.requests_failed += 1;
        }

        if let Some(cap) = capability {
            *metric
                .capabilities_used
                .entry(cap.to_string())
                .or_insert(0) += 1;
        }

        metric.last_updated = chrono::Utc::now().to_rfc3339();

        // Track latency
        let lat_vec = lats
            .entry(agent_id.to_string())
            .or_insert_with(Vec::new);
        lat_vec.push(latency_ms);

        // Update percentiles
        if lat_vec.len() >= 20 {
            lat_vec.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
            metric.average_latency_ms = lat_vec.iter().sum::<f32>() / lat_vec.len() as f32;
            metric.p95_latency_ms = lat_vec[(lat_vec.len() * 95) / 100];
            metric.p99_latency_ms = lat_vec[(lat_vec.len() * 99) / 100];
        }
    }

    /// Get metrics for agent (read-only lock, better concurrency)
    pub fn get_metrics(&self, agent_id: &str) -> Option<AgentMetrics> {
        self.metrics.read().unwrap().get(agent_id).cloned()
    }

    /// Get all metrics (read-only lock)
    pub fn get_all_metrics(&self) -> Vec<AgentMetrics> {
        self.metrics.read().unwrap().values().cloned().collect()
    }

    /// Reset metrics for agent (write lock)
    pub fn reset(&self, agent_id: &str) {
        self.metrics.write().unwrap().remove(agent_id);
        self.latencies.write().unwrap().remove(agent_id);
    }
}

impl Default for MetricsCollector {
    fn default() -> Self {
        Self::new()
    }
}

// QUALITY IMPROVEMENTS:
// - Changed from Mutex to RwLock: readers don't block each other
// - Multiple get_metrics() calls concurrent (common pattern)
// - Only record_request() needs exclusive write lock
// - Performance improvement for high-read scenarios