agent-heartbeat 0.1.0

/*!
agent-heartbeat: periodic heartbeat recorder and stall detector for AI agents.

Agents call `ping()` at regular intervals. If no ping arrives within
the configured window, `is_stale(now_ms)` returns `true`.

```rust
use agent_heartbeat::Heartbeat;

let mut hb = Heartbeat::new("worker-1").interval_ms(5_000);
hb.ping(1000);
assert!(!hb.is_stale(3000));
assert!(hb.is_stale(9000));
```
*/

/// Tracks periodic pings from an agent and detects stalls.
#[derive(Debug, Clone)]
pub struct Heartbeat {
    agent_id: String,
    interval_ms: u64,
    last_ping: Option<u64>,
    ping_count: u64,
}

impl Heartbeat {
    /// Create a new heartbeat tracker for `agent_id`.
    pub fn new(agent_id: impl Into<String>) -> Self {
        Self {
            agent_id: agent_id.into(),
            interval_ms: 10_000,
            last_ping: None,
            ping_count: 0,
        }
    }

    /// Set the maximum allowed gap between pings (milliseconds).
    pub fn interval_ms(mut self, ms: u64) -> Self {
        self.interval_ms = ms;
        self
    }

    /// Record a heartbeat at timestamp `now_ms`.
    pub fn ping(&mut self, now_ms: u64) {
        self.last_ping = Some(now_ms);
        self.ping_count += 1;
    }

    /// Returns `true` if no ping has been received or the last ping
    /// was more than `interval_ms` ago.
    pub fn is_stale(&self, now_ms: u64) -> bool {
        match self.last_ping {
            None => true,
            Some(last) => now_ms.saturating_sub(last) > self.interval_ms,
        }
    }

    /// Milliseconds since the last ping, or `None` if never pinged.
    pub fn elapsed_ms(&self, now_ms: u64) -> Option<u64> {
        self.last_ping.map(|last| now_ms.saturating_sub(last))
    }

    /// Total number of pings recorded.
    pub fn ping_count(&self) -> u64 {
        self.ping_count
    }

    /// Timestamp of the last ping, or `None`.
    pub fn last_ping_ms(&self) -> Option<u64> {
        self.last_ping
    }

    /// Agent ID this heartbeat tracks.
    pub fn agent_id(&self) -> &str {
        &self.agent_id
    }

    /// Configured stale interval in milliseconds.
    pub fn get_interval_ms(&self) -> u64 {
        self.interval_ms
    }

    /// Reset ping history (count and last timestamp).
    pub fn reset(&mut self) {
        self.last_ping = None;
        self.ping_count = 0;
    }
}

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

    #[test]
    fn new_is_stale() {
        let hb = Heartbeat::new("a");
        assert!(hb.is_stale(0));
    }

    #[test]
    fn ping_clears_stale() {
        let mut hb = Heartbeat::new("a").interval_ms(1000);
        hb.ping(500);
        assert!(!hb.is_stale(1000));
    }

    #[test]
    fn stale_after_interval() {
        let mut hb = Heartbeat::new("a").interval_ms(1000);
        hb.ping(0);
        assert!(hb.is_stale(1001));
    }

    #[test]
    fn exactly_at_boundary_not_stale() {
        let mut hb = Heartbeat::new("a").interval_ms(1000);
        hb.ping(0);
        assert!(!hb.is_stale(1000));
    }

    #[test]
    fn elapsed_ms_none_before_ping() {
        let hb = Heartbeat::new("a");
        assert!(hb.elapsed_ms(500).is_none());
    }

    #[test]
    fn elapsed_ms_after_ping() {
        let mut hb = Heartbeat::new("a");
        hb.ping(100);
        assert_eq!(hb.elapsed_ms(350), Some(250));
    }

    #[test]
    fn ping_count_increments() {
        let mut hb = Heartbeat::new("a");
        assert_eq!(hb.ping_count(), 0);
        hb.ping(1);
        hb.ping(2);
        hb.ping(3);
        assert_eq!(hb.ping_count(), 3);
    }

    #[test]
    fn last_ping_ms() {
        let mut hb = Heartbeat::new("a");
        assert!(hb.last_ping_ms().is_none());
        hb.ping(42);
        assert_eq!(hb.last_ping_ms(), Some(42));
        hb.ping(99);
        assert_eq!(hb.last_ping_ms(), Some(99));
    }

    #[test]
    fn agent_id() {
        let hb = Heartbeat::new("worker-42");
        assert_eq!(hb.agent_id(), "worker-42");
    }

    #[test]
    fn get_interval_ms() {
        let hb = Heartbeat::new("a").interval_ms(3000);
        assert_eq!(hb.get_interval_ms(), 3000);
    }

    #[test]
    fn reset_clears_state() {
        let mut hb = Heartbeat::new("a").interval_ms(1000);
        hb.ping(500);
        hb.reset();
        assert_eq!(hb.ping_count(), 0);
        assert!(hb.last_ping_ms().is_none());
        assert!(hb.is_stale(600));
    }

    #[test]
    fn multiple_pings_only_last_matters() {
        let mut hb = Heartbeat::new("a").interval_ms(500);
        hb.ping(100);
        hb.ping(200);
        hb.ping(800);
        assert!(!hb.is_stale(1200));
        assert!(hb.is_stale(1400));
    }

    #[test]
    fn default_interval_is_10s() {
        let hb = Heartbeat::new("a");
        assert_eq!(hb.get_interval_ms(), 10_000);
    }
}