use std::time::Duration;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HeartbeatConfig {
pub assessment_timeout: Duration,
pub action_timeout: Duration,
pub context_mode: HeartbeatContextMode,
pub consecutive_clear_backoff_threshold: u32,
pub max_backoff_multiplier: u32,
}
impl Default for HeartbeatConfig {
fn default() -> Self {
Self {
assessment_timeout: Duration::from_secs(30),
action_timeout: Duration::from_secs(120),
context_mode: HeartbeatContextMode::EphemeralWithSummary,
consecutive_clear_backoff_threshold: 5,
max_backoff_multiplier: 4,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum HeartbeatContextMode {
SharedPersistent,
#[default]
EphemeralWithSummary,
FullyEphemeral,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum HeartbeatSeverity {
Info,
Warning,
Critical,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum HeartbeatAssessment {
NeedsAction {
reason: String,
severity: HeartbeatSeverity,
data: Option<serde_json::Value>,
},
AllClear { summary: String },
Error { message: String },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HeartbeatState {
pub consecutive_clear_count: u32,
pub current_backoff: u32,
pub last_assessment_summary: Option<String>,
pub last_beat_at: Option<DateTime<Utc>>,
pub total_beats: u64,
pub total_actions: u64,
}
impl Default for HeartbeatState {
fn default() -> Self {
Self {
consecutive_clear_count: 0,
current_backoff: 1,
last_assessment_summary: None,
last_beat_at: None,
total_beats: 0,
total_actions: 0,
}
}
}
impl HeartbeatState {
pub fn record_assessment(
&mut self,
assessment: &HeartbeatAssessment,
config: &HeartbeatConfig,
) -> u32 {
self.total_beats += 1;
self.last_beat_at = Some(Utc::now());
match assessment {
HeartbeatAssessment::AllClear { summary } => {
self.consecutive_clear_count += 1;
self.last_assessment_summary = Some(summary.clone());
if self.consecutive_clear_count >= config.consecutive_clear_backoff_threshold {
self.current_backoff =
(self.current_backoff * 2).min(config.max_backoff_multiplier);
}
}
HeartbeatAssessment::NeedsAction { reason, .. } => {
self.consecutive_clear_count = 0;
self.current_backoff = 1; self.total_actions += 1;
self.last_assessment_summary = Some(reason.clone());
}
HeartbeatAssessment::Error { message } => {
self.consecutive_clear_count = 0;
self.last_assessment_summary = Some(format!("ERROR: {}", message));
}
}
self.current_backoff
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_heartbeat_config() {
let config = HeartbeatConfig::default();
assert_eq!(config.assessment_timeout, Duration::from_secs(30));
assert_eq!(config.action_timeout, Duration::from_secs(120));
assert_eq!(
config.context_mode,
HeartbeatContextMode::EphemeralWithSummary
);
assert_eq!(config.consecutive_clear_backoff_threshold, 5);
assert_eq!(config.max_backoff_multiplier, 4);
}
#[test]
fn heartbeat_state_tracks_clears() {
let config = HeartbeatConfig {
consecutive_clear_backoff_threshold: 3,
max_backoff_multiplier: 4,
..Default::default()
};
let mut state = HeartbeatState::default();
for _ in 0..3 {
state.record_assessment(
&HeartbeatAssessment::AllClear {
summary: "ok".to_string(),
},
&config,
);
}
assert_eq!(state.consecutive_clear_count, 3);
assert_eq!(state.current_backoff, 2);
assert_eq!(state.total_beats, 3);
assert_eq!(state.total_actions, 0);
}
#[test]
fn heartbeat_state_resets_on_action() {
let config = HeartbeatConfig::default();
let mut state = HeartbeatState::default();
state.consecutive_clear_count = 10;
state.current_backoff = 4;
state.record_assessment(
&HeartbeatAssessment::NeedsAction {
reason: "drift detected".to_string(),
severity: HeartbeatSeverity::Warning,
data: None,
},
&config,
);
assert_eq!(state.consecutive_clear_count, 0);
assert_eq!(state.current_backoff, 1);
assert_eq!(state.total_actions, 1);
}
#[test]
fn backoff_caps_at_max() {
let config = HeartbeatConfig {
consecutive_clear_backoff_threshold: 1,
max_backoff_multiplier: 4,
..Default::default()
};
let mut state = HeartbeatState::default();
for _ in 0..20 {
state.record_assessment(
&HeartbeatAssessment::AllClear {
summary: "ok".to_string(),
},
&config,
);
}
assert_eq!(state.current_backoff, 4);
}
#[test]
fn error_does_not_reset_backoff() {
let config = HeartbeatConfig::default();
let mut state = HeartbeatState::default();
state.current_backoff = 3;
state.record_assessment(
&HeartbeatAssessment::Error {
message: "timeout".to_string(),
},
&config,
);
assert_eq!(state.current_backoff, 3); assert_eq!(state.consecutive_clear_count, 0);
}
#[test]
fn heartbeat_assessment_serialization() {
let assessment = HeartbeatAssessment::NeedsAction {
reason: "policy drift".to_string(),
severity: HeartbeatSeverity::Critical,
data: Some(serde_json::json!({"drift_count": 3})),
};
let json = serde_json::to_string(&assessment).unwrap();
let deserialized: HeartbeatAssessment = serde_json::from_str(&json).unwrap();
match deserialized {
HeartbeatAssessment::NeedsAction {
reason, severity, ..
} => {
assert_eq!(reason, "policy drift");
assert_eq!(severity, HeartbeatSeverity::Critical);
}
_ => panic!("expected NeedsAction"),
}
}
}