#![allow(dead_code)]
use std::collections::HashMap;
use std::time::{Duration, Instant};
use tracing::{debug, info, warn};
#[derive(Debug, Clone)]
pub struct HealingConfig {
pub max_retries: u32,
pub base_delay_ms: u64,
pub max_delay_ms: u64,
pub jitter: f64,
pub circuit_breaker_threshold: u32,
pub circuit_breaker_recovery_secs: u64,
pub worker_dead_after_secs: u64,
pub max_failed_tasks: u32,
}
impl Default for HealingConfig {
fn default() -> Self {
Self {
max_retries: 3,
base_delay_ms: 1000,
max_delay_ms: 30000,
jitter: 0.3,
circuit_breaker_threshold: 5,
circuit_breaker_recovery_secs: 30,
worker_dead_after_secs: 60,
max_failed_tasks: 10,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HealingCircuitState {
Closed,
Open,
HalfOpen,
}
#[derive(Debug, Clone)]
pub struct HealingCircuitBreaker {
pub state: HealingCircuitState,
pub failure_count: u32,
pub last_failure_time: Option<Instant>,
pub open_duration: Duration,
pub threshold: u32,
}
impl HealingCircuitBreaker {
pub fn new(open_duration_secs: u64, threshold: u32) -> Self {
Self {
state: HealingCircuitState::Closed,
failure_count: 0,
last_failure_time: None,
open_duration: Duration::from_secs(open_duration_secs),
threshold,
}
}
pub fn record_success(&mut self) {
self.failure_count = 0;
self.state = HealingCircuitState::Closed;
}
pub fn record_failure(&mut self) {
self.failure_count += 1;
self.last_failure_time = Some(Instant::now());
if self.failure_count >= self.threshold {
self.state = HealingCircuitState::Open;
warn!(
failure_count = self.failure_count,
threshold = self.threshold,
"Circuit breaker opened"
);
}
}
pub fn can_execute(&mut self) -> bool {
match self.state {
HealingCircuitState::Closed => true,
HealingCircuitState::Open => {
if let Some(last) = self.last_failure_time {
if last.elapsed() >= self.open_duration {
self.state = HealingCircuitState::HalfOpen;
info!("Circuit breaker half-open — allowing test request");
true
} else {
false
}
} else {
false
}
}
HealingCircuitState::HalfOpen => true,
}
}
}
#[derive(Debug, Clone)]
pub struct FailureRecord {
pub agent_id: String,
pub failure_count: u32,
pub first_failure: Instant,
pub last_failure: Instant,
pub last_error: String,
pub is_dead: bool,
}
impl FailureRecord {
pub fn new(agent_id: &str, error: &str) -> Self {
Self {
agent_id: agent_id.to_string(),
failure_count: 1,
first_failure: Instant::now(),
last_failure: Instant::now(),
last_error: error.to_string(),
is_dead: false,
}
}
pub fn add_failure(&mut self, error: &str) {
self.failure_count += 1;
self.last_failure = Instant::now();
self.last_error = error.to_string();
}
pub fn record_success(&mut self) {
self.failure_count = 0;
self.is_dead = false;
}
pub fn time_since_last_failure(&self) -> Duration {
self.last_failure.elapsed()
}
}
#[derive(Debug)]
pub struct SelfHealingEngine {
pub config: HealingConfig,
failure_records: HashMap<String, FailureRecord>,
circuit_breakers: HashMap<String, HealingCircuitBreaker>,
}
impl SelfHealingEngine {
pub fn new() -> Self {
Self {
config: HealingConfig::default(),
failure_records: HashMap::new(),
circuit_breakers: HashMap::new(),
}
}
pub fn with_config(config: HealingConfig) -> Self {
Self {
config,
failure_records: HashMap::new(),
circuit_breakers: HashMap::new(),
}
}
pub fn record_failure(&mut self, agent_id: &str, error: &str) -> u32 {
let record = self
.failure_records
.entry(agent_id.to_string())
.and_modify(|r| r.add_failure(error))
.or_insert_with(|| FailureRecord::new(agent_id, error));
if record.failure_count >= self.config.max_failed_tasks {
record.is_dead = true;
warn!(
agent_id = %agent_id,
failure_count = record.failure_count,
max_failed = self.config.max_failed_tasks,
"Agent marked as dead"
);
}
let cb = self
.circuit_breakers
.entry(agent_id.to_string())
.or_insert_with(|| {
HealingCircuitBreaker::new(
self.config.circuit_breaker_recovery_secs,
self.config.circuit_breaker_threshold,
)
});
cb.record_failure();
record.failure_count
}
pub fn record_success(&mut self, agent_id: &str) {
if let Some(record) = self.failure_records.get_mut(agent_id) {
record.record_success();
}
if let Some(cb) = self.circuit_breakers.get_mut(agent_id) {
cb.record_success();
}
}
pub fn is_healthy(&mut self, agent_id: &str) -> bool {
if let Some(cb) = self.circuit_breakers.get_mut(agent_id) {
if !cb.can_execute() {
debug!(
agent_id = %agent_id,
state = ?cb.state,
"Agent blocked by circuit breaker"
);
return false;
}
}
if let Some(record) = self.failure_records.get(agent_id) {
if record.is_dead {
if record.time_since_last_failure()
> Duration::from_secs(self.config.worker_dead_after_secs)
{
debug!(
agent_id = %agent_id,
"Agent revived after dead timeout"
);
return true;
}
return false;
}
}
true
}
pub fn failure_count(&self, agent_id: &str) -> u32 {
self.failure_records
.get(agent_id)
.map(|r| r.failure_count)
.unwrap_or(0)
}
pub fn circuit_state(&self, agent_id: &str) -> Option<HealingCircuitState> {
self.circuit_breakers.get(agent_id).map(|cb| cb.state)
}
pub fn dead_workers(&self) -> Vec<String> {
self.failure_records
.iter()
.filter(|(_, r)| {
r.is_dead
&& r.time_since_last_failure()
> Duration::from_secs(self.config.worker_dead_after_secs)
})
.map(|(id, _)| id.clone())
.collect()
}
pub fn all_dead_workers(&self) -> Vec<String> {
self.failure_records
.iter()
.filter(|(_, r)| r.is_dead)
.map(|(id, _)| id.clone())
.collect()
}
pub fn delay_for_attempt(&self, attempt: u32) -> Duration {
use rand::Rng;
let exp = 2u64.pow(attempt);
let base = self.config.base_delay_ms * exp;
let capped = base.min(self.config.max_delay_ms);
let jitter_range = (capped as f64) * self.config.jitter;
let jitter = rand::thread_rng().gen_range(-jitter_range..=jitter_range) as u64;
let delay = capped.saturating_add(jitter).max(self.config.base_delay_ms);
Duration::from_millis(delay)
}
pub async fn retry_with_backoff<F, T>(
&self,
agent_id: &str,
is_transient: impl Fn(&str) -> bool,
mut operation: F,
) -> Result<T, String>
where
F: FnMut(u32) -> Result<T, String>,
{
let max_attempts = self.config.max_retries + 1; let mut last_error = String::from("Unknown error");
for attempt in 0..max_attempts {
if attempt > 0 {
let delay = self.delay_for_attempt(attempt - 1);
debug!(
agent_id = %agent_id,
attempt = attempt,
delay_ms = delay.as_millis(),
"Retrying after backoff"
);
tokio::time::sleep(delay).await;
}
match operation(attempt) {
Ok(result) => {
if attempt > 0 {
info!(
agent_id = %agent_id,
attempt = attempt,
"Operation succeeded after retry"
);
}
return Ok(result);
}
Err(e) => {
last_error = e.clone();
if !is_transient(&e) {
debug!(
agent_id = %agent_id,
attempt = attempt,
error = %e,
"Non-transient error — not retrying"
);
return Err(e);
}
warn!(
agent_id = %agent_id,
attempt = attempt,
max_attempts = max_attempts,
error = %e,
"Transient error — will retry"
);
}
}
}
Err(format!(
"All {} attempts failed for agent '{}': {}",
max_attempts, agent_id, last_error
))
}
pub fn reset(&mut self) {
self.failure_records.clear();
self.circuit_breakers.clear();
info!("Self-healing engine reset — all failure records cleared");
}
pub fn reset_agent(&mut self, agent_id: &str) {
self.failure_records.remove(agent_id);
self.circuit_breakers.remove(agent_id);
debug!(agent_id = %agent_id, "Self-healing record reset for agent");
}
pub fn tracked_agents(&self) -> usize {
self.failure_records.len()
}
pub fn dead_agent_count(&self) -> usize {
self.all_dead_workers().len()
}
pub fn open_circuit_count(&self) -> usize {
self.circuit_breakers
.values()
.filter(|cb| cb.state == HealingCircuitState::Open)
.count()
}
}
impl Default for SelfHealingEngine {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_healing_config_default() {
let config = HealingConfig::default();
assert_eq!(config.max_retries, 3);
assert_eq!(config.base_delay_ms, 1000);
assert_eq!(config.max_delay_ms, 30000);
assert_eq!(config.circuit_breaker_threshold, 5);
assert_eq!(config.circuit_breaker_recovery_secs, 30);
}
#[test]
fn test_circuit_breaker_initial_state() {
let mut cb = HealingCircuitBreaker::new(30, 5);
assert_eq!(cb.state, HealingCircuitState::Closed);
assert_eq!(cb.failure_count, 0);
assert!(cb.can_execute());
}
#[test]
fn test_circuit_breaker_opens_after_threshold() {
let mut cb = HealingCircuitBreaker::new(30, 3);
assert!(cb.can_execute());
cb.record_failure();
assert!(cb.can_execute());
assert_eq!(cb.failure_count, 1);
cb.record_failure();
assert!(cb.can_execute());
assert_eq!(cb.failure_count, 2);
cb.record_failure();
assert_eq!(cb.state, HealingCircuitState::Open);
assert!(!cb.can_execute());
assert_eq!(cb.failure_count, 3);
}
#[test]
fn test_circuit_breaker_success_resets() {
let mut cb = HealingCircuitBreaker::new(30, 3);
cb.record_failure();
cb.record_failure();
assert_eq!(cb.failure_count, 2);
cb.record_success();
assert_eq!(cb.state, HealingCircuitState::Closed);
assert_eq!(cb.failure_count, 0);
}
#[test]
fn test_failure_record_new() {
let record = FailureRecord::new("agent-1", "timeout");
assert_eq!(record.agent_id, "agent-1");
assert_eq!(record.failure_count, 1);
assert_eq!(record.last_error, "timeout");
assert!(!record.is_dead);
}
#[test]
fn test_failure_record_add_failure() {
let mut record = FailureRecord::new("agent-1", "timeout");
record.add_failure("connection refused");
assert_eq!(record.failure_count, 2);
assert_eq!(record.last_error, "connection refused");
}
#[test]
fn test_failure_record_success_resets() {
let mut record = FailureRecord::new("agent-1", "timeout");
record.add_failure("connection refused");
assert_eq!(record.failure_count, 2);
record.record_success();
assert_eq!(record.failure_count, 0);
assert!(!record.is_dead);
}
#[test]
fn test_self_healing_engine_new() {
let engine = SelfHealingEngine::new();
assert_eq!(engine.tracked_agents(), 0);
assert_eq!(engine.dead_agent_count(), 0);
}
#[test]
fn test_self_healing_engine_record_failure() {
let mut engine = SelfHealingEngine::new();
let count = engine.record_failure("agent-1", "timeout");
assert_eq!(count, 1);
assert_eq!(engine.failure_count("agent-1"), 1);
assert!(engine.is_healthy("agent-1"));
}
#[test]
fn test_self_healing_engine_marks_dead() {
let mut engine = SelfHealingEngine::new();
engine.config.max_failed_tasks = 3;
engine.record_failure("agent-1", "error 1");
engine.record_failure("agent-1", "error 2");
engine.record_failure("agent-1", "error 3");
assert!(!engine.is_healthy("agent-1"));
assert!(engine.all_dead_workers().contains(&"agent-1".to_string()));
}
#[test]
fn test_self_healing_engine_circuit_breaker() {
let mut engine = SelfHealingEngine::new();
engine.config.circuit_breaker_threshold = 3;
engine.record_failure("agent-1", "error 1");
engine.record_failure("agent-1", "error 2");
engine.record_failure("agent-1", "error 3");
assert_eq!(
engine.circuit_state("agent-1"),
Some(HealingCircuitState::Open)
);
assert!(!engine.is_healthy("agent-1"));
}
#[test]
fn test_self_healing_engine_success_resets() {
let mut engine = SelfHealingEngine::new();
engine.config.circuit_breaker_threshold = 3;
engine.record_failure("agent-1", "error 1");
engine.record_failure("agent-1", "error 2");
assert_eq!(engine.failure_count("agent-1"), 2);
engine.record_success("agent-1");
assert_eq!(engine.failure_count("agent-1"), 0);
assert_eq!(
engine.circuit_state("agent-1"),
Some(HealingCircuitState::Closed)
);
}
#[test]
fn test_delay_for_attempt() {
let engine = SelfHealingEngine::new();
let d0 = engine.delay_for_attempt(0);
assert!(d0.as_millis() >= 700); assert!(d0.as_millis() <= 1300);
let d1 = engine.delay_for_attempt(1);
assert!(d1.as_millis() >= 1400);
assert!(d1.as_millis() <= 2600);
let d4 = engine.delay_for_attempt(4);
assert!(d4.as_millis() <= 30000 + 9000); }
#[test]
fn test_reset() {
let mut engine = SelfHealingEngine::new();
engine.record_failure("agent-1", "error");
engine.record_failure("agent-2", "error");
assert_eq!(engine.tracked_agents(), 2);
engine.reset();
assert_eq!(engine.tracked_agents(), 0);
}
#[test]
fn test_reset_agent() {
let mut engine = SelfHealingEngine::new();
engine.record_failure("agent-1", "error");
engine.record_failure("agent-2", "error");
assert_eq!(engine.tracked_agents(), 2);
engine.reset_agent("agent-1");
assert_eq!(engine.tracked_agents(), 1);
assert_eq!(engine.failure_count("agent-1"), 0);
}
#[test]
fn test_dead_workers_empty_initially() {
let engine = SelfHealingEngine::new();
assert!(engine.dead_workers().is_empty());
assert!(engine.all_dead_workers().is_empty());
}
#[test]
fn test_open_circuit_count() {
let mut engine = SelfHealingEngine::new();
engine.config.circuit_breaker_threshold = 2;
engine.record_failure("agent-1", "error");
engine.record_failure("agent-1", "error");
engine.record_failure("agent-2", "error");
engine.record_failure("agent-2", "error");
assert_eq!(engine.open_circuit_count(), 2);
}
#[tokio::test]
async fn test_retry_with_backoff_succeeds_on_first_attempt() {
let engine = SelfHealingEngine::new();
let result = engine
.retry_with_backoff("agent-1", |_| true, |_| Ok::<_, String>(42))
.await;
assert_eq!(result, Ok(42));
}
#[tokio::test]
async fn test_retry_with_backoff_succeeds_after_retries() {
let engine = SelfHealingEngine::new();
let attempt = std::sync::atomic::AtomicU32::new(0);
let result = engine
.retry_with_backoff(
"agent-1",
|_| true,
|_| {
let prev = attempt.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if prev < 2 {
Err("transient error".to_string())
} else {
Ok(42)
}
},
)
.await;
assert_eq!(result, Ok(42));
}
#[tokio::test]
async fn test_retry_with_backoff_fails_on_non_transient() {
let engine = SelfHealingEngine::new();
let result: Result<i32, String> = engine
.retry_with_backoff(
"agent-1",
|e| !e.contains("non-transient"),
|_| Err("non-transient error".to_string()),
)
.await;
assert!(result.is_err());
assert!(result.unwrap_err().contains("non-transient"));
}
#[tokio::test]
async fn test_retry_with_backoff_exhausts_retries() {
let mut engine = SelfHealingEngine::new();
engine.config.max_retries = 2;
let result: Result<i32, String> = engine
.retry_with_backoff(
"agent-1",
|_| true,
|_| Err("persistent transient error".to_string()),
)
.await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.contains("All 3 attempts failed"));
assert!(err.contains("persistent transient error"));
}
}