Skip to main content

ri_agent_graph/
retry.rs

1use crate::error::AgentGraphError;
2use std::sync::Arc;
3use std::time::Duration;
4
5/// Type alias for retry predicate function.
6pub type RetryPredicate = Arc<dyn Fn(&AgentGraphError) -> bool + Send + Sync>;
7
8/// Retry policy for node execution.
9#[derive(Clone)]
10pub struct RetryPolicy {
11    /// Maximum number of attempts (including first try)
12    pub max_attempts: usize,
13    /// Initial delay between retries
14    pub initial_interval: Duration,
15    /// Multiplicative backoff factor
16    pub backoff_factor: f64,
17    /// Maximum delay between retries
18    pub max_interval: Duration,
19    /// Whether to add random jitter to delays
20    pub jitter: bool,
21    /// Optional predicate to determine if an error is retryable
22    pub retry_on: Option<RetryPredicate>,
23}
24
25impl Default for RetryPolicy {
26    fn default() -> Self {
27        Self {
28            max_attempts: 3,
29            initial_interval: Duration::from_secs(1),
30            backoff_factor: 2.0,
31            max_interval: Duration::from_secs(60),
32            jitter: true,
33            retry_on: None,
34        }
35    }
36}
37
38impl RetryPolicy {
39    pub fn new() -> Self {
40        Self::default()
41    }
42
43    pub fn with_max_attempts(mut self, n: usize) -> Self {
44        self.max_attempts = n;
45        self
46    }
47
48    pub fn with_initial_interval(mut self, d: Duration) -> Self {
49        self.initial_interval = d;
50        self
51    }
52
53    pub fn with_backoff_factor(mut self, f: f64) -> Self {
54        self.backoff_factor = f;
55        self
56    }
57
58    pub fn with_max_interval(mut self, d: Duration) -> Self {
59        self.max_interval = d;
60        self
61    }
62
63    pub fn with_jitter(mut self, jitter: bool) -> Self {
64        self.jitter = jitter;
65        self
66    }
67
68    pub fn with_retry_on(
69        mut self,
70        predicate: impl Fn(&AgentGraphError) -> bool + Send + Sync + 'static,
71    ) -> Self {
72        self.retry_on = Some(Arc::new(predicate));
73        self
74    }
75
76    /// Check if a given error should be retried
77    pub fn should_retry(&self, error: &AgentGraphError) -> bool {
78        match &self.retry_on {
79            Some(predicate) => predicate(error),
80            None => true,
81        }
82    }
83
84    /// Calculate delay for a given attempt number (0-indexed)
85    pub fn delay_for_attempt(&self, attempt: usize) -> Duration {
86        let base = self.initial_interval.as_secs_f64() * self.backoff_factor.powi(attempt as i32);
87        let capped = base.min(self.max_interval.as_secs_f64());
88        if self.jitter {
89            let jitter_factor = jitter_factor();
90            Duration::from_secs_f64(capped * jitter_factor)
91        } else {
92            Duration::from_secs_f64(capped)
93        }
94    }
95}
96
97/// Simple deterministic jitter factor using hash-based pseudo-randomness.
98fn jitter_factor() -> f64 {
99    use std::collections::hash_map::DefaultHasher;
100    use std::hash::{Hash, Hasher};
101    use std::time::SystemTime;
102
103    let mut hasher = DefaultHasher::new();
104    SystemTime::now()
105        .duration_since(SystemTime::UNIX_EPOCH)
106        .unwrap_or_default()
107        .as_nanos()
108        .hash(&mut hasher);
109    std::thread::current().id().hash(&mut hasher);
110    let hash = hasher.finish();
111    // Normalize to [0.5, 1.0] range for reasonable jitter
112    0.5 + (hash as f64 / u64::MAX as f64) * 0.5
113}
114
115impl std::fmt::Debug for RetryPolicy {
116    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117        f.debug_struct("RetryPolicy")
118            .field("max_attempts", &self.max_attempts)
119            .field("initial_interval", &self.initial_interval)
120            .field("backoff_factor", &self.backoff_factor)
121            .field("max_interval", &self.max_interval)
122            .field("jitter", &self.jitter)
123            .field("retry_on", &self.retry_on.is_some())
124            .finish()
125    }
126}