agent_heartbeat/
lib.rs

1/*!
2agent-heartbeat: periodic heartbeat recorder and stall detector for AI agents.
3
4Agents call `ping()` at regular intervals. If no ping arrives within
5the configured window, `is_stale(now_ms)` returns `true`.
6
7```rust
8use agent_heartbeat::Heartbeat;
9
10let mut hb = Heartbeat::new("worker-1").interval_ms(5_000);
11hb.ping(1000);
12assert!(!hb.is_stale(3000));
13assert!(hb.is_stale(9000));
14```
15*/
16
17/// Tracks periodic pings from an agent and detects stalls.
18#[derive(Debug, Clone)]
19pub struct Heartbeat {
20    agent_id: String,
21    interval_ms: u64,
22    last_ping: Option<u64>,
23    ping_count: u64,
24}
25
26impl Heartbeat {
27    /// Create a new heartbeat tracker for `agent_id`.
28    pub fn new(agent_id: impl Into<String>) -> Self {
29        Self {
30            agent_id: agent_id.into(),
31            interval_ms: 10_000,
32            last_ping: None,
33            ping_count: 0,
34        }
35    }
36
37    /// Set the maximum allowed gap between pings (milliseconds).
38    pub fn interval_ms(mut self, ms: u64) -> Self {
39        self.interval_ms = ms;
40        self
41    }
42
43    /// Record a heartbeat at timestamp `now_ms`.
44    pub fn ping(&mut self, now_ms: u64) {
45        self.last_ping = Some(now_ms);
46        self.ping_count += 1;
47    }
48
49    /// Returns `true` if no ping has been received or the last ping
50    /// was more than `interval_ms` ago.
51    pub fn is_stale(&self, now_ms: u64) -> bool {
52        match self.last_ping {
53            None => true,
54            Some(last) => now_ms.saturating_sub(last) > self.interval_ms,
55        }
56    }
57
58    /// Milliseconds since the last ping, or `None` if never pinged.
59    pub fn elapsed_ms(&self, now_ms: u64) -> Option<u64> {
60        self.last_ping.map(|last| now_ms.saturating_sub(last))
61    }
62
63    /// Total number of pings recorded.
64    pub fn ping_count(&self) -> u64 {
65        self.ping_count
66    }
67
68    /// Timestamp of the last ping, or `None`.
69    pub fn last_ping_ms(&self) -> Option<u64> {
70        self.last_ping
71    }
72
73    /// Agent ID this heartbeat tracks.
74    pub fn agent_id(&self) -> &str {
75        &self.agent_id
76    }
77
78    /// Configured stale interval in milliseconds.
79    pub fn get_interval_ms(&self) -> u64 {
80        self.interval_ms
81    }
82
83    /// Reset ping history (count and last timestamp).
84    pub fn reset(&mut self) {
85        self.last_ping = None;
86        self.ping_count = 0;
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn new_is_stale() {
96        let hb = Heartbeat::new("a");
97        assert!(hb.is_stale(0));
98    }
99
100    #[test]
101    fn ping_clears_stale() {
102        let mut hb = Heartbeat::new("a").interval_ms(1000);
103        hb.ping(500);
104        assert!(!hb.is_stale(1000));
105    }
106
107    #[test]
108    fn stale_after_interval() {
109        let mut hb = Heartbeat::new("a").interval_ms(1000);
110        hb.ping(0);
111        assert!(hb.is_stale(1001));
112    }
113
114    #[test]
115    fn exactly_at_boundary_not_stale() {
116        let mut hb = Heartbeat::new("a").interval_ms(1000);
117        hb.ping(0);
118        assert!(!hb.is_stale(1000));
119    }
120
121    #[test]
122    fn elapsed_ms_none_before_ping() {
123        let hb = Heartbeat::new("a");
124        assert!(hb.elapsed_ms(500).is_none());
125    }
126
127    #[test]
128    fn elapsed_ms_after_ping() {
129        let mut hb = Heartbeat::new("a");
130        hb.ping(100);
131        assert_eq!(hb.elapsed_ms(350), Some(250));
132    }
133
134    #[test]
135    fn ping_count_increments() {
136        let mut hb = Heartbeat::new("a");
137        assert_eq!(hb.ping_count(), 0);
138        hb.ping(1);
139        hb.ping(2);
140        hb.ping(3);
141        assert_eq!(hb.ping_count(), 3);
142    }
143
144    #[test]
145    fn last_ping_ms() {
146        let mut hb = Heartbeat::new("a");
147        assert!(hb.last_ping_ms().is_none());
148        hb.ping(42);
149        assert_eq!(hb.last_ping_ms(), Some(42));
150        hb.ping(99);
151        assert_eq!(hb.last_ping_ms(), Some(99));
152    }
153
154    #[test]
155    fn agent_id() {
156        let hb = Heartbeat::new("worker-42");
157        assert_eq!(hb.agent_id(), "worker-42");
158    }
159
160    #[test]
161    fn get_interval_ms() {
162        let hb = Heartbeat::new("a").interval_ms(3000);
163        assert_eq!(hb.get_interval_ms(), 3000);
164    }
165
166    #[test]
167    fn reset_clears_state() {
168        let mut hb = Heartbeat::new("a").interval_ms(1000);
169        hb.ping(500);
170        hb.reset();
171        assert_eq!(hb.ping_count(), 0);
172        assert!(hb.last_ping_ms().is_none());
173        assert!(hb.is_stale(600));
174    }
175
176    #[test]
177    fn multiple_pings_only_last_matters() {
178        let mut hb = Heartbeat::new("a").interval_ms(500);
179        hb.ping(100);
180        hb.ping(200);
181        hb.ping(800);
182        assert!(!hb.is_stale(1200));
183        assert!(hb.is_stale(1400));
184    }
185
186    #[test]
187    fn default_interval_is_10s() {
188        let hb = Heartbeat::new("a");
189        assert_eq!(hb.get_interval_ms(), 10_000);
190    }
191}
agent_heartbeat/lib.rs

agent_heartbeat/
lib.rs
