ares_agent/
loop_mode.rs

1//! Long-running iteration mode for ares agents.
2//!
3//! Provides a data model for agents that execute on a fixed interval
4//! (cron-like) rather than responding to a single request. Each "beat"
5//! runs one unit of work; the agent continues until it hits a halt
6//! condition (max iterations, consecutive failures, external stop).
7//!
8//! Phase 1 (this module): pure data types and configuration. The
9//! runtime scheduler, tick dispatcher, and integration with
10//! `orchestrator` land in follow-up phases.
11
12use serde::{Deserialize, Serialize};
13use std::future::Future;
14use std::pin::Pin;
15use std::sync::atomic::{AtomicBool, Ordering};
16use std::sync::Arc;
17use std::time::Duration;
18
19/// Configuration for a long-running iteration-mode agent.
20#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
21pub struct LoopModeConfig {
22    /// Wall-clock interval between iterations, in seconds.
23    #[serde(default = "default_interval_secs")]
24    pub interval_secs: u64,
25    /// Maximum total iterations before the agent halts. `None` = unbounded.
26    #[serde(default)]
27    pub max_iterations: Option<u64>,
28    /// Halt if this many consecutive iterations fail.
29    #[serde(default = "default_halt_threshold")]
30    pub halt_on_consecutive_failures: u32,
31    /// Prompt used when the agent has nothing picked for the current iteration
32    /// (anti-idle fallback).
33    #[serde(default)]
34    pub fallback_prompt: Option<String>,
35    /// Whether failed iterations count against `max_iterations`.
36    #[serde(default = "default_count_failures")]
37    pub count_failed_iterations: bool,
38}
39
40fn default_interval_secs() -> u64 {
41    180
42}
43fn default_halt_threshold() -> u32 {
44    3
45}
46fn default_count_failures() -> bool {
47    true
48}
49
50impl LoopModeConfig {
51    /// The interval as a `std::time::Duration` (minimum 1ms for scheduler safety).
52    pub fn interval(&self) -> Duration {
53        Duration::from_secs(self.interval_secs).max(Duration::from_millis(1))
54    }
55}
56
57impl Default for LoopModeConfig {
58    fn default() -> Self {
59        Self {
60            interval_secs: default_interval_secs(),
61            max_iterations: None,
62            halt_on_consecutive_failures: default_halt_threshold(),
63            fallback_prompt: None,
64            count_failed_iterations: default_count_failures(),
65        }
66    }
67}
68
69/// Runtime state of a long-running iteration-mode agent.
70#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
71pub struct LoopModeState {
72    pub iterations_run: u64,
73    pub iterations_succeeded: u64,
74    pub iterations_failed: u64,
75    pub consecutive_failures: u32,
76    pub started_at_epoch_secs: u64,
77    pub last_tick_epoch_secs: u64,
78}
79
80impl LoopModeState {
81    /// Record a successful iteration. Resets the consecutive-failure counter.
82    pub fn record_success(&mut self, now_epoch_secs: u64) {
83        self.iterations_run += 1;
84        self.iterations_succeeded += 1;
85        self.consecutive_failures = 0;
86        self.last_tick_epoch_secs = now_epoch_secs;
87    }
88
89    /// Record a failed iteration. Increments the consecutive-failure counter.
90    pub fn record_failure(&mut self, now_epoch_secs: u64) {
91        self.iterations_run += 1;
92        self.iterations_failed += 1;
93        self.consecutive_failures += 1;
94        self.last_tick_epoch_secs = now_epoch_secs;
95    }
96
97    /// Check whether the state should halt given the provided config.
98    pub fn should_halt(&self, config: &LoopModeConfig) -> Option<LoopFinishReason> {
99        if self.consecutive_failures >= config.halt_on_consecutive_failures {
100            return Some(LoopFinishReason::ConsecutiveFailures);
101        }
102        if let Some(max) = config.max_iterations {
103            let counted = if config.count_failed_iterations {
104                self.iterations_run
105            } else {
106                self.iterations_succeeded
107            };
108            if counted >= max {
109                return Some(LoopFinishReason::MaxIterationsReached);
110            }
111        }
112        None
113    }
114}
115
116/// Reason a long-running iteration-mode agent stopped.
117#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
118#[serde(rename_all = "snake_case")]
119pub enum LoopFinishReason {
120    /// Hit `max_iterations` from config.
121    MaxIterationsReached,
122    /// Hit `halt_on_consecutive_failures`.
123    ConsecutiveFailures,
124    /// External stop signal (user, supervisor, SIGTERM).
125    ExternalStop,
126    /// Runtime error that the agent could not recover from.
127    FatalError,
128}
129
130/// Boxed async tick function: returns Ok(()) on success, Err on failure.
131pub type TickFn =
132    Box<dyn Fn() -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync>;
133
134/// Runtime scheduler that drives an iteration-mode agent.
135pub struct LoopRunner {
136    pub config: LoopModeConfig,
137    pub state: LoopModeState,
138    stop: Arc<AtomicBool>,
139}
140
141impl LoopRunner {
142    pub fn new(config: LoopModeConfig) -> Self {
143        Self {
144            config,
145            state: LoopModeState::default(),
146            stop: Arc::new(AtomicBool::new(false)),
147        }
148    }
149
150    pub fn stop_handle(&self) -> Arc<AtomicBool> {
151        self.stop.clone()
152    }
153
154    pub async fn run(&mut self, tick: &TickFn) -> LoopFinishReason {
155        self.run_with_state_observer(tick, |_| async {}).await
156    }
157
158    pub async fn run_with_state_observer<F, Fut>(
159        &mut self,
160        tick: &TickFn,
161        mut observe: F,
162    ) -> LoopFinishReason
163    where
164        F: FnMut(LoopModeState) -> Fut,
165        Fut: Future<Output = ()>,
166    {
167        let now = std::time::SystemTime::now()
168            .duration_since(std::time::UNIX_EPOCH)
169            .unwrap_or_default()
170            .as_secs();
171        self.state.started_at_epoch_secs = now;
172        observe(self.state.clone()).await;
173
174        let mut interval = tokio::time::interval(self.config.interval());
175        interval.tick().await; // first tick fires immediately
176
177        loop {
178            if self.stop.load(Ordering::Relaxed) {
179                return LoopFinishReason::ExternalStop;
180            }
181
182            let now = std::time::SystemTime::now()
183                .duration_since(std::time::UNIX_EPOCH)
184                .unwrap_or_default()
185                .as_secs();
186
187            match tick().await {
188                Ok(()) => self.state.record_success(now),
189                Err(_) => self.state.record_failure(now),
190            }
191            observe(self.state.clone()).await;
192
193            if let Some(reason) = self.state.should_halt(&self.config) {
194                return reason;
195            }
196
197            interval.tick().await;
198        }
199    }
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[test]
207    fn config_defaults_are_sensible() {
208        let cfg = LoopModeConfig::default();
209        assert_eq!(cfg.interval_secs, 180);
210        assert_eq!(cfg.interval(), Duration::from_secs(180));
211        assert_eq!(cfg.max_iterations, None);
212        assert_eq!(cfg.halt_on_consecutive_failures, 3);
213        assert!(cfg.fallback_prompt.is_none());
214        assert!(cfg.count_failed_iterations);
215    }
216
217    #[test]
218    fn config_round_trips_through_toml() {
219        let cfg = LoopModeConfig {
220            interval_secs: 60,
221            max_iterations: Some(100),
222            halt_on_consecutive_failures: 5,
223            fallback_prompt: Some("expand test coverage".into()),
224            count_failed_iterations: false,
225        };
226        let serialized = toml::to_string(&cfg).expect("serialize");
227        let parsed: LoopModeConfig = toml::from_str(&serialized).expect("parse");
228        assert_eq!(cfg, parsed);
229    }
230
231    #[test]
232    fn state_record_success_resets_consecutive_failures() {
233        let mut state = LoopModeState::default();
234        state.record_failure(10);
235        state.record_failure(20);
236        assert_eq!(state.consecutive_failures, 2);
237        state.record_success(30);
238        assert_eq!(state.consecutive_failures, 0);
239        assert_eq!(state.iterations_run, 3);
240        assert_eq!(state.iterations_succeeded, 1);
241        assert_eq!(state.iterations_failed, 2);
242        assert_eq!(state.last_tick_epoch_secs, 30);
243    }
244
245    #[test]
246    fn should_halt_on_consecutive_failures() {
247        let cfg = LoopModeConfig {
248            halt_on_consecutive_failures: 3,
249            ..LoopModeConfig::default()
250        };
251        let mut state = LoopModeState::default();
252        state.record_failure(10);
253        state.record_failure(20);
254        assert_eq!(state.should_halt(&cfg), None);
255        state.record_failure(30);
256        assert_eq!(
257            state.should_halt(&cfg),
258            Some(LoopFinishReason::ConsecutiveFailures)
259        );
260    }
261
262    #[test]
263    fn should_halt_on_max_iterations_counting_failures() {
264        let cfg = LoopModeConfig {
265            max_iterations: Some(2),
266            halt_on_consecutive_failures: 999,
267            count_failed_iterations: true,
268            ..LoopModeConfig::default()
269        };
270        let mut state = LoopModeState::default();
271        state.record_success(10);
272        assert_eq!(state.should_halt(&cfg), None);
273        state.record_failure(20);
274        assert_eq!(
275            state.should_halt(&cfg),
276            Some(LoopFinishReason::MaxIterationsReached)
277        );
278    }
279
280    #[test]
281    fn should_halt_on_max_iterations_ignoring_failures() {
282        let cfg = LoopModeConfig {
283            max_iterations: Some(2),
284            halt_on_consecutive_failures: 999,
285            count_failed_iterations: false,
286            ..LoopModeConfig::default()
287        };
288        let mut state = LoopModeState::default();
289        state.record_failure(10);
290        state.record_failure(20);
291        state.record_failure(30);
292        assert_eq!(state.should_halt(&cfg), None);
293        state.record_success(40);
294        state.record_success(50);
295        assert_eq!(
296            state.should_halt(&cfg),
297            Some(LoopFinishReason::MaxIterationsReached)
298        );
299    }
300
301    #[test]
302    fn unbounded_loop_never_halts_without_failures() {
303        let cfg = LoopModeConfig::default();
304        let mut state = LoopModeState::default();
305        for i in 0..1000 {
306            state.record_success(i);
307        }
308        assert_eq!(state.should_halt(&cfg), None);
309    }
310
311    #[test]
312    fn finish_reason_serializes_snake_case() {
313        assert_eq!(
314            serde_json::to_string(&LoopFinishReason::MaxIterationsReached).unwrap(),
315            "\"max_iterations_reached\""
316        );
317        assert_eq!(
318            serde_json::to_string(&LoopFinishReason::ConsecutiveFailures).unwrap(),
319            "\"consecutive_failures\""
320        );
321    }
322
323    #[tokio::test]
324    async fn loop_runner_halts_on_max_iterations() {
325        let config = LoopModeConfig {
326            interval_secs: 0,
327            max_iterations: Some(3),
328            halt_on_consecutive_failures: 999,
329            ..LoopModeConfig::default()
330        };
331        let mut runner = LoopRunner::new(config);
332        let tick: TickFn = Box::new(|| Box::pin(async { Ok(()) }));
333        let reason = runner.run(&tick).await;
334        assert_eq!(reason, LoopFinishReason::MaxIterationsReached);
335        assert_eq!(runner.state.iterations_run, 3);
336    }
337
338    #[tokio::test]
339    async fn loop_runner_observes_started_and_tick_state() {
340        let config = LoopModeConfig {
341            interval_secs: 0,
342            max_iterations: Some(2),
343            halt_on_consecutive_failures: 999,
344            ..LoopModeConfig::default()
345        };
346        let mut runner = LoopRunner::new(config);
347        let observed = std::sync::Arc::new(tokio::sync::Mutex::new(Vec::new()));
348        let observed_for_callback = observed.clone();
349        let tick: TickFn = Box::new(|| Box::pin(async { Ok(()) }));
350
351        let reason = runner
352            .run_with_state_observer(&tick, move |state| {
353                let observed = observed_for_callback.clone();
354                async move {
355                    observed.lock().await.push(state);
356                }
357            })
358            .await;
359
360        assert_eq!(reason, LoopFinishReason::MaxIterationsReached);
361        let snapshots = observed.lock().await;
362        assert!(snapshots[0].started_at_epoch_secs > 0);
363        assert_eq!(snapshots[0].iterations_run, 0);
364        assert_eq!(snapshots[1].iterations_run, 1);
365        assert_eq!(snapshots[2].iterations_run, 2);
366    }
367
368    #[tokio::test]
369    async fn loop_runner_halts_on_consecutive_failures() {
370        let config = LoopModeConfig {
371            interval_secs: 0,
372            max_iterations: None,
373            halt_on_consecutive_failures: 2,
374            ..LoopModeConfig::default()
375        };
376        let mut runner = LoopRunner::new(config);
377        let tick: TickFn = Box::new(|| Box::pin(async { Err("fail".into()) }));
378        let reason = runner.run(&tick).await;
379        assert_eq!(reason, LoopFinishReason::ConsecutiveFailures);
380        assert_eq!(runner.state.consecutive_failures, 2);
381    }
382
383    #[tokio::test]
384    async fn loop_runner_external_stop() {
385        let config = LoopModeConfig {
386            interval_secs: 0,
387            max_iterations: None,
388            halt_on_consecutive_failures: 999,
389            ..LoopModeConfig::default()
390        };
391        let mut runner = LoopRunner::new(config);
392        let stop = runner.stop_handle();
393        stop.store(true, std::sync::atomic::Ordering::Relaxed);
394        let tick: TickFn = Box::new(|| Box::pin(async { Ok(()) }));
395        let reason = runner.run(&tick).await;
396        assert_eq!(reason, LoopFinishReason::ExternalStop);
397    }
398}
ares_agent/loop_mode.rs

ares_agent/
loop_mode.rs
