openlatch-client 0.1.14

The open-source security layer for AI agents — client forwarder
use std::collections::HashMap;
use std::time::{Duration, Instant};

const BACKOFF_BASE: [u64; 5] = [1, 4, 16, 64, 256];
const CIRCUIT_OPEN_THRESHOLD: u32 = 5;
const CIRCUIT_OPEN_WINDOW: Duration = Duration::from_secs(600);
const CIRCUIT_AUTO_CLOSE: Duration = Duration::from_secs(3600);
const HEALTHY_RESET: Duration = Duration::from_secs(600);
const JITTER_PERCENT: f64 = 0.20;

#[derive(Debug, Clone, PartialEq)]
pub enum HealState {
    Healthy,
    Drifted,
    Healing,
    Backoff { until: Instant },
    CircuitOpen { opened_at: Instant },
}

#[derive(Debug)]
pub struct EntryHealTracker {
    pub state: HealState,
    pub attempt: u32,
    pub last_healthy: Option<Instant>,
    recent_heals: Vec<Instant>,
}

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

impl EntryHealTracker {
    pub fn new() -> Self {
        Self {
            state: HealState::Healthy,
            attempt: 0,
            last_healthy: Some(Instant::now()),
            recent_heals: Vec::new(),
        }
    }

    pub fn mark_healthy(&mut self) {
        let now = Instant::now();
        if let Some(last) = self.last_healthy {
            if now.duration_since(last) >= HEALTHY_RESET {
                self.attempt = 0;
                self.recent_heals.clear();
            }
        }
        self.state = HealState::Healthy;
        self.last_healthy = Some(now);
    }

    pub fn mark_drifted(&mut self) {
        if matches!(self.state, HealState::CircuitOpen { .. }) {
            return;
        }
        self.state = HealState::Drifted;
        self.last_healthy = None;
    }

    pub fn should_heal(&self) -> bool {
        matches!(self.state, HealState::Drifted)
    }

    pub fn is_circuit_open(&self) -> bool {
        if let HealState::CircuitOpen { opened_at } = self.state {
            if Instant::now().duration_since(opened_at) >= CIRCUIT_AUTO_CLOSE {
                return false;
            }
            return true;
        }
        false
    }

    pub fn record_heal_attempt(&mut self) {
        let now = Instant::now();
        self.attempt += 1;
        self.recent_heals
            .retain(|t| now.duration_since(*t) < CIRCUIT_OPEN_WINDOW);
        self.recent_heals.push(now);

        if self.recent_heals.len() as u32 >= CIRCUIT_OPEN_THRESHOLD {
            self.state = HealState::CircuitOpen { opened_at: now };
            tracing::warn!(
                attempts = self.recent_heals.len(),
                "circuit breaker opened — auto-heal paused for this entry"
            );
            return;
        }

        self.state = HealState::Healing;
    }

    pub fn record_heal_success(&mut self) {
        if matches!(self.state, HealState::CircuitOpen { .. }) {
            return;
        }
        self.state = HealState::Healthy;
        self.last_healthy = Some(Instant::now());
    }

    pub fn record_heal_failure(&mut self) {
        let now = Instant::now();
        let idx = (self.attempt as usize).min(BACKOFF_BASE.len()) - 1;
        let base_secs = BACKOFF_BASE[idx];
        let jitter = (base_secs as f64 * JITTER_PERCENT * (rand_fraction() * 2.0 - 1.0)) as i64;
        let delay_secs = (base_secs as i64 + jitter).max(1) as u64;

        self.state = HealState::Backoff {
            until: now + Duration::from_secs(delay_secs),
        };
    }

    pub fn force_close_circuit(&mut self) {
        self.state = HealState::Healthy;
        self.attempt = 0;
        self.recent_heals.clear();
        self.last_healthy = Some(Instant::now());
    }

    pub fn backoff_remaining_ms(&self) -> Option<u64> {
        if let HealState::Backoff { until } = self.state {
            let now = Instant::now();
            if until > now {
                return Some(until.duration_since(now).as_millis() as u64);
            }
        }
        None
    }
}

fn rand_fraction() -> f64 {
    let bytes = uuid::Uuid::new_v4().as_bytes()[0];
    bytes as f64 / 255.0
}

#[derive(Debug, Default)]
pub struct HealManager {
    trackers: HashMap<String, EntryHealTracker>,
}

impl HealManager {
    pub fn new() -> Self {
        Self {
            trackers: HashMap::new(),
        }
    }

    pub fn tracker(&mut self, entry_id: &str) -> &mut EntryHealTracker {
        self.trackers.entry(entry_id.to_string()).or_default()
    }

    pub fn open_circuits(&self) -> Vec<&str> {
        self.trackers
            .iter()
            .filter(|(_, t)| t.is_circuit_open())
            .map(|(id, _)| id.as_str())
            .collect()
    }

    pub fn force_close(&mut self, entry_id: &str) -> bool {
        if let Some(tracker) = self.trackers.get_mut(entry_id) {
            tracker.force_close_circuit();
            true
        } else {
            false
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new_tracker_is_healthy() {
        let t = EntryHealTracker::new();
        assert_eq!(t.state, HealState::Healthy);
        assert_eq!(t.attempt, 0);
    }

    #[test]
    fn drifted_then_should_heal() {
        let mut t = EntryHealTracker::new();
        t.mark_drifted();
        assert!(t.should_heal());
    }

    #[test]
    fn circuit_opens_after_rapid_heals() {
        let mut t = EntryHealTracker::new();
        for _ in 0..5 {
            t.mark_drifted();
            t.record_heal_attempt();
            t.record_heal_success();
            t.mark_drifted();
        }
        assert!(t.is_circuit_open());
    }

    #[test]
    fn circuit_blocks_drift() {
        let mut t = EntryHealTracker::new();
        for _ in 0..5 {
            t.mark_drifted();
            t.record_heal_attempt();
            t.record_heal_success();
            t.mark_drifted();
        }
        assert!(t.is_circuit_open());
        t.mark_drifted();
        assert!(!t.should_heal());
    }

    #[test]
    fn force_close_resets_circuit() {
        let mut t = EntryHealTracker::new();
        for _ in 0..5 {
            t.mark_drifted();
            t.record_heal_attempt();
            t.record_heal_success();
            t.mark_drifted();
        }
        assert!(t.is_circuit_open());
        t.force_close_circuit();
        assert!(!t.is_circuit_open());
        assert_eq!(t.attempt, 0);
    }

    #[test]
    fn heal_manager_tracks_multiple_entries() {
        let mut m = HealManager::new();
        m.tracker("entry-1").mark_drifted();
        m.tracker("entry-2").mark_healthy();
        assert!(m.tracker("entry-1").should_heal());
        assert!(!m.tracker("entry-2").should_heal());
    }

    #[test]
    fn heal_manager_open_circuits() {
        let mut m = HealManager::new();
        for _ in 0..5 {
            m.tracker("entry-1").mark_drifted();
            m.tracker("entry-1").record_heal_attempt();
            m.tracker("entry-1").record_heal_success();
            m.tracker("entry-1").mark_drifted();
        }
        let open = m.open_circuits();
        assert!(open.contains(&"entry-1"));
    }

    #[test]
    fn force_close_via_manager() {
        let mut m = HealManager::new();
        for _ in 0..5 {
            m.tracker("entry-1").mark_drifted();
            m.tracker("entry-1").record_heal_attempt();
            m.tracker("entry-1").record_heal_success();
            m.tracker("entry-1").mark_drifted();
        }
        assert!(m.force_close("entry-1"));
        assert!(m.open_circuits().is_empty());
    }
}