use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
pub const AGENT_HEALTH_COLLECTION: &str = "_ai_agent_health";
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(rename_all = "lowercase")]
pub enum CircuitState {
#[default]
Closed,
Open,
HalfOpen,
}
impl std::fmt::Display for CircuitState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
CircuitState::Closed => write!(f, "closed"),
CircuitState::Open => write!(f, "open"),
CircuitState::HalfOpen => write!(f, "half_open"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentHealthMetrics {
#[serde(rename = "_key")]
pub agent_id: String,
#[serde(default)]
pub circuit_state: CircuitState,
#[serde(default)]
pub consecutive_failures: u32,
#[serde(default)]
pub window_failures: u32,
#[serde(default)]
pub window_requests: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub last_heartbeat: Option<DateTime<Utc>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub circuit_state_changed_at: Option<DateTime<Utc>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub circuit_retry_at: Option<DateTime<Utc>>,
#[serde(default)]
pub active_tasks: u32,
#[serde(default)]
pub tasks_recovered: u32,
pub updated_at: DateTime<Utc>,
}
impl AgentHealthMetrics {
pub fn new(agent_id: String) -> Self {
Self {
agent_id,
circuit_state: CircuitState::Closed,
consecutive_failures: 0,
window_failures: 0,
window_requests: 0,
last_heartbeat: Some(Utc::now()),
circuit_state_changed_at: Some(Utc::now()),
circuit_retry_at: None,
active_tasks: 0,
tasks_recovered: 0,
updated_at: Utc::now(),
}
}
pub fn record_success(&mut self) {
self.consecutive_failures = 0;
self.window_requests += 1;
self.updated_at = Utc::now();
if self.circuit_state == CircuitState::HalfOpen {
self.transition_to_closed();
}
}
pub fn record_failure(&mut self, failure_threshold: u32, failure_rate_threshold: f64) {
self.consecutive_failures += 1;
self.window_failures += 1;
self.window_requests += 1;
self.updated_at = Utc::now();
let should_open = match self.circuit_state {
CircuitState::Closed => {
self.consecutive_failures >= failure_threshold
|| (self.window_requests >= 10 && self.failure_rate() >= failure_rate_threshold)
}
CircuitState::HalfOpen => {
true
}
CircuitState::Open => false,
};
if should_open {
self.transition_to_open();
}
}
pub fn failure_rate(&self) -> f64 {
if self.window_requests == 0 {
0.0
} else {
self.window_failures as f64 / self.window_requests as f64
}
}
pub fn transition_to_closed(&mut self) {
self.circuit_state = CircuitState::Closed;
self.circuit_state_changed_at = Some(Utc::now());
self.circuit_retry_at = None;
self.consecutive_failures = 0;
self.window_failures = 0;
self.window_requests = 0;
self.updated_at = Utc::now();
}
pub fn transition_to_open(&mut self) {
self.circuit_state = CircuitState::Open;
self.circuit_state_changed_at = Some(Utc::now());
self.updated_at = Utc::now();
}
pub fn set_retry_at(&mut self, retry_at: DateTime<Utc>) {
self.circuit_retry_at = Some(retry_at);
self.updated_at = Utc::now();
}
pub fn transition_to_half_open(&mut self) {
self.circuit_state = CircuitState::HalfOpen;
self.circuit_state_changed_at = Some(Utc::now());
self.circuit_retry_at = None;
self.updated_at = Utc::now();
}
pub fn should_try_half_open(&self) -> bool {
if self.circuit_state != CircuitState::Open {
return false;
}
match self.circuit_retry_at {
Some(retry_at) => Utc::now() >= retry_at,
None => false,
}
}
pub fn is_healthy(&self, heartbeat_timeout_secs: u64) -> bool {
match self.last_heartbeat {
Some(last) => {
let elapsed = (Utc::now() - last).num_seconds() as u64;
elapsed < heartbeat_timeout_secs
}
None => false,
}
}
pub fn update_heartbeat(&mut self) {
self.last_heartbeat = Some(Utc::now());
self.updated_at = Utc::now();
}
pub fn allows_requests(&self) -> bool {
matches!(
self.circuit_state,
CircuitState::Closed | CircuitState::HalfOpen
)
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct RecoverySystemStatus {
pub total_agents: usize,
pub agents_circuit_open: usize,
pub agents_unhealthy: usize,
pub stalled_tasks: usize,
pub stuck_contributions: usize,
pub last_scan: Option<DateTime<Utc>>,
pub recent_events: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_agent_health_new() {
let health = AgentHealthMetrics::new("agent-001".to_string());
assert_eq!(health.agent_id, "agent-001");
assert_eq!(health.circuit_state, CircuitState::Closed);
assert_eq!(health.consecutive_failures, 0);
assert!(health.last_heartbeat.is_some());
}
#[test]
fn test_record_success() {
let mut health = AgentHealthMetrics::new("agent-001".to_string());
health.consecutive_failures = 3;
health.record_success();
assert_eq!(health.consecutive_failures, 0);
assert_eq!(health.window_requests, 1);
}
#[test]
fn test_record_failure_opens_circuit() {
let mut health = AgentHealthMetrics::new("agent-001".to_string());
for _ in 0..5 {
health.record_failure(5, 0.5);
}
assert_eq!(health.circuit_state, CircuitState::Open);
assert_eq!(health.consecutive_failures, 5);
}
#[test]
fn test_failure_rate() {
let mut health = AgentHealthMetrics::new("agent-001".to_string());
health.window_requests = 10;
health.window_failures = 3;
assert!((health.failure_rate() - 0.3).abs() < 0.001);
}
#[test]
fn test_half_open_success_closes() {
let mut health = AgentHealthMetrics::new("agent-001".to_string());
health.transition_to_half_open();
assert_eq!(health.circuit_state, CircuitState::HalfOpen);
health.record_success();
assert_eq!(health.circuit_state, CircuitState::Closed);
}
#[test]
fn test_half_open_failure_reopens() {
let mut health = AgentHealthMetrics::new("agent-001".to_string());
health.transition_to_half_open();
health.record_failure(5, 0.5);
assert_eq!(health.circuit_state, CircuitState::Open);
}
#[test]
fn test_allows_requests() {
let mut health = AgentHealthMetrics::new("agent-001".to_string());
assert!(health.allows_requests());
health.transition_to_half_open();
assert!(health.allows_requests());
health.transition_to_open();
assert!(!health.allows_requests()); }
#[test]
fn test_circuit_state_display() {
assert_eq!(CircuitState::Closed.to_string(), "closed");
assert_eq!(CircuitState::Open.to_string(), "open");
assert_eq!(CircuitState::HalfOpen.to_string(), "half_open");
}
}