af-llm 0.2.0

//! In-process circuit breaker for LLM calls.
//!
//! Direct port of `agent_core/llm/circuit_breaker.py`, made `Send + Sync` so it
//! can be shared across Tokio tasks behind an `Arc`. The state machine and the
//! exact thresholds (including the 120s cooldown lesson from the DeepSeek
//! overload incident) are preserved.
//!
//! ```text
//! CLOSED    → normal operation, requests pass through
//! OPEN      → provider is down, requests fail fast (no LLM call)
//! HALF_OPEN → testing if provider recovered (allow 1 request through)
//!
//! CLOSED    → OPEN:      `failure_threshold` consecutive failures within window
//! OPEN      → HALF_OPEN: after `cooldown` elapses
//! HALF_OPEN → CLOSED:    test request succeeds
//! HALF_OPEN → OPEN:      test request fails
//! ```
//!
//! State is in-process. For multi-worker fan-out, back it with Redis (same as
//! the note in the Python original).

use std::sync::Mutex;
use std::time::{Duration, Instant};

pub const FAILURE_THRESHOLD: u32 = 5;
pub const FAILURE_WINDOW: Duration = Duration::from_secs(60);
// 120s, not 30s: a shorter cooldown fired the half-open probe while DeepSeek was
// still overloaded, which re-opened the circuit and looped.
pub const COOLDOWN: Duration = Duration::from_secs(120);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum State {
    Closed,
    Open,
    HalfOpen,
}

impl State {
    pub fn as_str(self) -> &'static str {
        match self {
            State::Closed => "closed",
            State::Open => "open",
            State::HalfOpen => "half_open",
        }
    }
}

#[derive(Debug)]
struct Inner {
    state: State,
    failure_count: u32,
    last_failure: Option<Instant>,
    opened_at: Option<Instant>,
}

/// Thread-safe circuit breaker. Cheap to clone the `Arc`; the lock is held only
/// for trivial state transitions.
#[derive(Debug)]
pub struct CircuitBreaker {
    failure_threshold: u32,
    failure_window: Duration,
    cooldown: Duration,
    inner: Mutex<Inner>,
}

impl Default for CircuitBreaker {
    fn default() -> Self {
        Self::new(FAILURE_THRESHOLD, FAILURE_WINDOW, COOLDOWN)
    }
}

impl CircuitBreaker {
    pub fn new(failure_threshold: u32, failure_window: Duration, cooldown: Duration) -> Self {
        Self {
            failure_threshold,
            failure_window,
            cooldown,
            inner: Mutex::new(Inner {
                state: State::Closed,
                failure_count: 0,
                last_failure: None,
                opened_at: None,
            }),
        }
    }

    /// Current state, applying the OPEN → HALF_OPEN cooldown transition lazily
    /// (matches the Python `state` property side effect).
    pub fn state(&self) -> State {
        let mut g = self.inner.lock().unwrap();
        if g.state == State::Open {
            if let Some(opened_at) = g.opened_at {
                if opened_at.elapsed() >= self.cooldown {
                    g.state = State::HalfOpen;
                    tracing::info!(target: "llm.circuit", "circuit_breaker_half_open");
                }
            }
        }
        g.state
    }

    /// Whether the circuit is open (caller should fail fast).
    pub fn is_open(&self) -> bool {
        self.state() == State::Open
    }

    /// Record a successful call — closes the circuit and resets the counter.
    pub fn record_success(&self) {
        let mut g = self.inner.lock().unwrap();
        if g.state == State::HalfOpen {
            tracing::info!(target: "llm.circuit", reason = "half_open_success", "circuit_breaker_closed");
        }
        g.state = State::Closed;
        g.failure_count = 0;
    }

    /// Record a failed call — may open or re-open the circuit.
    pub fn record_failure(&self) {
        let now = Instant::now();
        let mut g = self.inner.lock().unwrap();

        // Reset the counter if the previous failure is outside the window.
        if let Some(last) = g.last_failure {
            if now.duration_since(last) > self.failure_window {
                g.failure_count = 0;
            }
        }

        g.failure_count += 1;
        g.last_failure = Some(now);

        if g.state == State::HalfOpen {
            g.state = State::Open;
            g.opened_at = Some(now);
            tracing::warn!(target: "llm.circuit", failures = g.failure_count, "circuit_breaker_reopened");
        } else if g.failure_count >= self.failure_threshold {
            g.state = State::Open;
            g.opened_at = Some(now);
            tracing::warn!(target: "llm.circuit", failures = g.failure_count, "circuit_breaker_opened");
        }
    }

    /// Snapshot for health checks.
    pub fn status(&self) -> Status {
        let state = self.state();
        let g = self.inner.lock().unwrap();
        Status {
            state,
            failure_count: g.failure_count,
            threshold: self.failure_threshold,
            cooldown: self.cooldown,
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct Status {
    pub state: State,
    pub failure_count: u32,
    pub threshold: u32,
    pub cooldown: Duration,
}

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

    #[test]
    fn opens_after_threshold_failures() {
        let cb = CircuitBreaker::default();
        assert!(!cb.is_open());
        for _ in 0..FAILURE_THRESHOLD {
            cb.record_failure();
        }
        assert!(cb.is_open());
        assert_eq!(cb.state(), State::Open);
    }

    #[test]
    fn success_resets() {
        let cb = CircuitBreaker::default();
        cb.record_failure();
        cb.record_failure();
        cb.record_success();
        assert_eq!(cb.status().failure_count, 0);
        assert!(!cb.is_open());
    }

    #[test]
    fn cooldown_moves_to_half_open_then_closes_on_success() {
        // Tiny cooldown so the test does not sleep for two minutes.
        let cb = CircuitBreaker::new(2, FAILURE_WINDOW, Duration::from_millis(10));
        cb.record_failure();
        cb.record_failure();
        assert_eq!(cb.state(), State::Open);
        std::thread::sleep(Duration::from_millis(20));
        assert_eq!(cb.state(), State::HalfOpen);
        cb.record_success();
        assert_eq!(cb.state(), State::Closed);
    }

    #[test]
    fn half_open_failure_reopens() {
        let cb = CircuitBreaker::new(2, FAILURE_WINDOW, Duration::from_millis(10));
        cb.record_failure();
        cb.record_failure();
        std::thread::sleep(Duration::from_millis(20));
        assert_eq!(cb.state(), State::HalfOpen);
        cb.record_failure();
        assert_eq!(cb.state(), State::Open);
    }
}