use serde::{Deserialize, Serialize};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct LoopModeConfig {
#[serde(default = "default_interval_secs")]
pub interval_secs: u64,
#[serde(default)]
pub max_iterations: Option<u64>,
#[serde(default = "default_halt_threshold")]
pub halt_on_consecutive_failures: u32,
#[serde(default)]
pub fallback_prompt: Option<String>,
#[serde(default = "default_count_failures")]
pub count_failed_iterations: bool,
}
fn default_interval_secs() -> u64 {
180
}
fn default_halt_threshold() -> u32 {
3
}
fn default_count_failures() -> bool {
true
}
impl LoopModeConfig {
pub fn interval(&self) -> Duration {
Duration::from_secs(self.interval_secs).max(Duration::from_millis(1))
}
}
impl Default for LoopModeConfig {
fn default() -> Self {
Self {
interval_secs: default_interval_secs(),
max_iterations: None,
halt_on_consecutive_failures: default_halt_threshold(),
fallback_prompt: None,
count_failed_iterations: default_count_failures(),
}
}
}
#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct LoopModeState {
pub iterations_run: u64,
pub iterations_succeeded: u64,
pub iterations_failed: u64,
pub consecutive_failures: u32,
pub started_at_epoch_secs: u64,
pub last_tick_epoch_secs: u64,
}
impl LoopModeState {
pub fn record_success(&mut self, now_epoch_secs: u64) {
self.iterations_run += 1;
self.iterations_succeeded += 1;
self.consecutive_failures = 0;
self.last_tick_epoch_secs = now_epoch_secs;
}
pub fn record_failure(&mut self, now_epoch_secs: u64) {
self.iterations_run += 1;
self.iterations_failed += 1;
self.consecutive_failures += 1;
self.last_tick_epoch_secs = now_epoch_secs;
}
pub fn should_halt(&self, config: &LoopModeConfig) -> Option<LoopFinishReason> {
if self.consecutive_failures >= config.halt_on_consecutive_failures {
return Some(LoopFinishReason::ConsecutiveFailures);
}
if let Some(max) = config.max_iterations {
let counted = if config.count_failed_iterations {
self.iterations_run
} else {
self.iterations_succeeded
};
if counted >= max {
return Some(LoopFinishReason::MaxIterationsReached);
}
}
None
}
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum LoopFinishReason {
MaxIterationsReached,
ConsecutiveFailures,
ExternalStop,
FatalError,
}
pub type TickFn =
Box<dyn Fn() -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync>;
pub struct LoopRunner {
pub config: LoopModeConfig,
pub state: LoopModeState,
stop: Arc<AtomicBool>,
}
impl LoopRunner {
pub fn new(config: LoopModeConfig) -> Self {
Self {
config,
state: LoopModeState::default(),
stop: Arc::new(AtomicBool::new(false)),
}
}
pub fn stop_handle(&self) -> Arc<AtomicBool> {
self.stop.clone()
}
pub async fn run(&mut self, tick: &TickFn) -> LoopFinishReason {
self.run_with_state_observer(tick, |_| async {}).await
}
pub async fn run_with_state_observer<F, Fut>(
&mut self,
tick: &TickFn,
mut observe: F,
) -> LoopFinishReason
where
F: FnMut(LoopModeState) -> Fut,
Fut: Future<Output = ()>,
{
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
self.state.started_at_epoch_secs = now;
observe(self.state.clone()).await;
let mut interval = tokio::time::interval(self.config.interval());
interval.tick().await;
loop {
if self.stop.load(Ordering::Relaxed) {
return LoopFinishReason::ExternalStop;
}
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
match tick().await {
Ok(()) => self.state.record_success(now),
Err(_) => self.state.record_failure(now),
}
observe(self.state.clone()).await;
if let Some(reason) = self.state.should_halt(&self.config) {
return reason;
}
interval.tick().await;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn config_defaults_are_sensible() {
let cfg = LoopModeConfig::default();
assert_eq!(cfg.interval_secs, 180);
assert_eq!(cfg.interval(), Duration::from_secs(180));
assert_eq!(cfg.max_iterations, None);
assert_eq!(cfg.halt_on_consecutive_failures, 3);
assert!(cfg.fallback_prompt.is_none());
assert!(cfg.count_failed_iterations);
}
#[test]
fn config_round_trips_through_toml() {
let cfg = LoopModeConfig {
interval_secs: 60,
max_iterations: Some(100),
halt_on_consecutive_failures: 5,
fallback_prompt: Some("expand test coverage".into()),
count_failed_iterations: false,
};
let serialized = toml::to_string(&cfg).expect("serialize");
let parsed: LoopModeConfig = toml::from_str(&serialized).expect("parse");
assert_eq!(cfg, parsed);
}
#[test]
fn state_record_success_resets_consecutive_failures() {
let mut state = LoopModeState::default();
state.record_failure(10);
state.record_failure(20);
assert_eq!(state.consecutive_failures, 2);
state.record_success(30);
assert_eq!(state.consecutive_failures, 0);
assert_eq!(state.iterations_run, 3);
assert_eq!(state.iterations_succeeded, 1);
assert_eq!(state.iterations_failed, 2);
assert_eq!(state.last_tick_epoch_secs, 30);
}
#[test]
fn should_halt_on_consecutive_failures() {
let cfg = LoopModeConfig {
halt_on_consecutive_failures: 3,
..LoopModeConfig::default()
};
let mut state = LoopModeState::default();
state.record_failure(10);
state.record_failure(20);
assert_eq!(state.should_halt(&cfg), None);
state.record_failure(30);
assert_eq!(
state.should_halt(&cfg),
Some(LoopFinishReason::ConsecutiveFailures)
);
}
#[test]
fn should_halt_on_max_iterations_counting_failures() {
let cfg = LoopModeConfig {
max_iterations: Some(2),
halt_on_consecutive_failures: 999,
count_failed_iterations: true,
..LoopModeConfig::default()
};
let mut state = LoopModeState::default();
state.record_success(10);
assert_eq!(state.should_halt(&cfg), None);
state.record_failure(20);
assert_eq!(
state.should_halt(&cfg),
Some(LoopFinishReason::MaxIterationsReached)
);
}
#[test]
fn should_halt_on_max_iterations_ignoring_failures() {
let cfg = LoopModeConfig {
max_iterations: Some(2),
halt_on_consecutive_failures: 999,
count_failed_iterations: false,
..LoopModeConfig::default()
};
let mut state = LoopModeState::default();
state.record_failure(10);
state.record_failure(20);
state.record_failure(30);
assert_eq!(state.should_halt(&cfg), None);
state.record_success(40);
state.record_success(50);
assert_eq!(
state.should_halt(&cfg),
Some(LoopFinishReason::MaxIterationsReached)
);
}
#[test]
fn unbounded_loop_never_halts_without_failures() {
let cfg = LoopModeConfig::default();
let mut state = LoopModeState::default();
for i in 0..1000 {
state.record_success(i);
}
assert_eq!(state.should_halt(&cfg), None);
}
#[test]
fn finish_reason_serializes_snake_case() {
assert_eq!(
serde_json::to_string(&LoopFinishReason::MaxIterationsReached).unwrap(),
"\"max_iterations_reached\""
);
assert_eq!(
serde_json::to_string(&LoopFinishReason::ConsecutiveFailures).unwrap(),
"\"consecutive_failures\""
);
}
#[tokio::test]
async fn loop_runner_halts_on_max_iterations() {
let config = LoopModeConfig {
interval_secs: 0,
max_iterations: Some(3),
halt_on_consecutive_failures: 999,
..LoopModeConfig::default()
};
let mut runner = LoopRunner::new(config);
let tick: TickFn = Box::new(|| Box::pin(async { Ok(()) }));
let reason = runner.run(&tick).await;
assert_eq!(reason, LoopFinishReason::MaxIterationsReached);
assert_eq!(runner.state.iterations_run, 3);
}
#[tokio::test]
async fn loop_runner_observes_started_and_tick_state() {
let config = LoopModeConfig {
interval_secs: 0,
max_iterations: Some(2),
halt_on_consecutive_failures: 999,
..LoopModeConfig::default()
};
let mut runner = LoopRunner::new(config);
let observed = std::sync::Arc::new(tokio::sync::Mutex::new(Vec::new()));
let observed_for_callback = observed.clone();
let tick: TickFn = Box::new(|| Box::pin(async { Ok(()) }));
let reason = runner
.run_with_state_observer(&tick, move |state| {
let observed = observed_for_callback.clone();
async move {
observed.lock().await.push(state);
}
})
.await;
assert_eq!(reason, LoopFinishReason::MaxIterationsReached);
let snapshots = observed.lock().await;
assert!(snapshots[0].started_at_epoch_secs > 0);
assert_eq!(snapshots[0].iterations_run, 0);
assert_eq!(snapshots[1].iterations_run, 1);
assert_eq!(snapshots[2].iterations_run, 2);
}
#[tokio::test]
async fn loop_runner_halts_on_consecutive_failures() {
let config = LoopModeConfig {
interval_secs: 0,
max_iterations: None,
halt_on_consecutive_failures: 2,
..LoopModeConfig::default()
};
let mut runner = LoopRunner::new(config);
let tick: TickFn = Box::new(|| Box::pin(async { Err("fail".into()) }));
let reason = runner.run(&tick).await;
assert_eq!(reason, LoopFinishReason::ConsecutiveFailures);
assert_eq!(runner.state.consecutive_failures, 2);
}
#[tokio::test]
async fn loop_runner_external_stop() {
let config = LoopModeConfig {
interval_secs: 0,
max_iterations: None,
halt_on_consecutive_failures: 999,
..LoopModeConfig::default()
};
let mut runner = LoopRunner::new(config);
let stop = runner.stop_handle();
stop.store(true, std::sync::atomic::Ordering::Relaxed);
let tick: TickFn = Box::new(|| Box::pin(async { Ok(()) }));
let reason = runner.run(&tick).await;
assert_eq!(reason, LoopFinishReason::ExternalStop);
}
}