Skip to main content

vtcode_commons/
retry.rs

1//! Canonical retry policy shared across the workspace.
2//!
3//! This module owns the retry *policy math*: attempt budgets, exponential
4//! backoff with an optional deterministic jitter, and category-based retry
5//! decisions built on [`ErrorCategory::is_retryable`]. Domain-specific
6//! adapters (typed error downcasts, tool-aware timeout rules, LLM
7//! `Retry-After` extraction) live in `vtcode-core::retry` as an extension
8//! trait over this policy.
9//!
10//! Wire-level HTTP clients that only need "should I retry this call?" use
11//! [`RetryPolicy::classify_anyhow`] / [`RetryPolicy::classify_status`];
12//! richer loops use [`RetryPolicy::decision_for_category`].
13
14use std::time::Duration;
15
16use crate::error_category::{ErrorCategory, classify_anyhow_error};
17
18/// Typed retry policy shared across runtime layers.
19#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize)]
20pub struct RetryPolicy {
21    /// Maximum number of total attempts, including the initial call.
22    pub max_attempts: u32,
23    pub initial_delay: Duration,
24    pub max_delay: Duration,
25    pub multiplier: f64,
26    pub jitter: f64,
27}
28
29impl RetryPolicy {
30    pub const DEFAULT: Self = Self::from_retries(2, Duration::from_secs(1), Duration::from_secs(60), 2.0);
31
32    pub const fn new(max_attempts: u32, initial_delay: Duration, max_delay: Duration, multiplier: f64) -> Self {
33        Self {
34            max_attempts: if max_attempts < 1 { 1 } else { max_attempts },
35            initial_delay,
36            max_delay,
37            multiplier: if multiplier < 1.0 { 1.0 } else { multiplier },
38            jitter: 0.0,
39        }
40    }
41
42    pub const fn from_retries(max_retries: u32, initial_delay: Duration, max_delay: Duration, multiplier: f64) -> Self {
43        Self::new(max_retries.saturating_add(1), initial_delay, max_delay, multiplier)
44    }
45
46    /// Millisecond-based constructor for wire clients.
47    ///
48    /// Uses a 2.0 multiplier and no jitter, so
49    /// [`Self::delay_for_attempt`] reproduces the classic
50    /// `base_ms << attempt` doubling curve capped at `max_delay_ms`.
51    pub fn simple(max_retries: u32, base_delay_ms: u64, max_delay_ms: u64) -> Self {
52        Self::from_retries(max_retries, Duration::from_millis(base_delay_ms), Duration::from_millis(max_delay_ms), 2.0)
53    }
54
55    pub fn delay_for_attempt(&self, attempt_index: u32) -> Duration {
56        let multiplier = self.multiplier.powi(attempt_index as i32);
57        let base_delay = Duration::try_from_secs_f64(self.initial_delay.as_secs_f64() * multiplier)
58            .unwrap_or(self.max_delay)
59            .min(self.max_delay);
60
61        if !self.jitter.is_finite() || self.jitter <= 0.0 {
62            return base_delay;
63        }
64
65        #[allow(clippy::cast_sign_loss)]
66        let max_jitter_ms = (base_delay.as_millis() as f64 * self.jitter)
67            .round()
68            .clamp(0.0, u64::MAX as f64) as u64;
69        if max_jitter_ms == 0 {
70            return base_delay;
71        }
72
73        let offset = (u64::from(attempt_index) * 31) % max_jitter_ms.saturating_add(1);
74        base_delay.saturating_add(Duration::from_millis(offset))
75    }
76
77    pub fn decision_for_category(
78        &self,
79        category: ErrorCategory,
80        attempt_index: u32,
81        retry_after: Option<Duration>,
82    ) -> RetryDecision {
83        let has_remaining_attempts = attempt_index.saturating_add(1) < self.max_attempts;
84        if !category.is_retryable() || !has_remaining_attempts {
85            return RetryDecision {
86                category,
87                retryable: false,
88                delay: None,
89                retry_after,
90            };
91        }
92
93        let delay = retry_after.unwrap_or_else(|| self.delay_for_attempt(attempt_index));
94        RetryDecision {
95            category,
96            retryable: true,
97            delay: Some(delay),
98            retry_after,
99        }
100    }
101
102    /// Classify an `anyhow::Error` for retry eligibility.
103    ///
104    /// Attempt-agnostic: `retryable` reflects only the error category, not
105    /// the remaining attempt budget. Wire clients that manage their own
106    /// attempt counting use this; loops that want budget-aware decisions
107    /// use [`Self::decision_for_category`].
108    pub fn classify_anyhow(&self, error: &anyhow::Error) -> RetryDecision {
109        let category = classify_anyhow_error(error);
110        RetryDecision {
111            category,
112            retryable: category.is_retryable(),
113            delay: None,
114            retry_after: None,
115        }
116    }
117
118    /// Classify an HTTP status code for retry eligibility.
119    ///
120    /// Attempt-agnostic, like [`Self::classify_anyhow`].
121    pub fn classify_status(&self, status: u16) -> RetryDecision {
122        let category = match status {
123            429 => ErrorCategory::RateLimit,
124            500 | 502 | 504 => ErrorCategory::Network,
125            503 => ErrorCategory::ServiceUnavailable,
126            401 | 403 => ErrorCategory::Authentication,
127            _ => ErrorCategory::ExecutionError,
128        };
129        RetryDecision {
130            category,
131            retryable: category.is_retryable(),
132            delay: None,
133            retry_after: None,
134        }
135    }
136}
137
138impl Default for RetryPolicy {
139    fn default() -> Self {
140        Self::from_retries(2, Duration::from_secs(1), Duration::from_secs(60), 2.0)
141    }
142}
143
144/// Result of classifying a failure for retry handling.
145#[derive(Debug, Clone, PartialEq, Eq)]
146pub struct RetryDecision {
147    pub category: ErrorCategory,
148    pub retryable: bool,
149    pub delay: Option<Duration>,
150    pub retry_after: Option<Duration>,
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[test]
158    fn default_policy_allows_two_retries() {
159        let policy = RetryPolicy::default();
160        assert_eq!(policy.max_attempts, 3);
161        assert_eq!(policy.initial_delay, Duration::from_secs(1));
162        assert_eq!(policy.max_delay, Duration::from_secs(60));
163    }
164
165    #[test]
166    fn classify_status_rate_limit() {
167        let policy = RetryPolicy::default();
168        let decision = policy.classify_status(429);
169        assert!(decision.retryable);
170        assert_eq!(decision.category, ErrorCategory::RateLimit);
171    }
172
173    #[test]
174    fn classify_status_server_error() {
175        let policy = RetryPolicy::default();
176        let decision = policy.classify_status(503);
177        assert!(decision.retryable);
178        assert_eq!(decision.category, ErrorCategory::ServiceUnavailable);
179    }
180
181    #[test]
182    fn classify_status_auth_not_retryable() {
183        let policy = RetryPolicy::default();
184        let decision = policy.classify_status(401);
185        assert!(!decision.retryable);
186        assert_eq!(decision.category, ErrorCategory::Authentication);
187    }
188
189    #[test]
190    fn classify_anyhow_network_error() {
191        let policy = RetryPolicy::default();
192        let err = anyhow::anyhow!("connection refused");
193        let decision = policy.classify_anyhow(&err);
194        assert!(decision.retryable);
195    }
196
197    #[test]
198    fn simple_policy_matches_bit_shift_doubling() {
199        // Parity with the historical `base_ms << attempt` curve used by
200        // wire clients before consolidation.
201        let policy = RetryPolicy::simple(10, 1000, 5000);
202        let legacy = |attempt: u32| -> u64 { 1000u64.saturating_mul(1u64 << attempt.min(16)).min(5000) };
203        for attempt in 0..6 {
204            assert_eq!(
205                policy.delay_for_attempt(attempt),
206                Duration::from_millis(legacy(attempt)),
207                "delay mismatch at attempt {attempt}"
208            );
209        }
210    }
211
212    #[test]
213    fn delay_for_attempt_clamps_overflowing_backoff_to_max_delay() {
214        let policy = RetryPolicy::from_retries(3, Duration::from_secs(1), Duration::from_secs(8), f64::MAX);
215
216        assert_eq!(policy.delay_for_attempt(2), Duration::from_secs(8));
217    }
218
219    #[test]
220    fn delay_for_attempt_ignores_non_finite_jitter() {
221        let mut policy = RetryPolicy::from_retries(3, Duration::from_secs(1), Duration::from_secs(8), 2.0);
222        policy.jitter = f64::INFINITY;
223
224        assert_eq!(policy.delay_for_attempt(1), Duration::from_secs(2));
225    }
226
227    #[test]
228    fn delay_for_attempt_handles_huge_finite_jitter() {
229        let mut policy = RetryPolicy::from_retries(3, Duration::from_secs(1), Duration::from_secs(8), 2.0);
230        policy.jitter = f64::MAX;
231
232        assert!(policy.delay_for_attempt(1) >= Duration::from_secs(2));
233    }
234
235    #[test]
236    fn decision_for_category_respects_attempt_budget() {
237        let policy = RetryPolicy::from_retries(1, Duration::from_secs(1), Duration::from_secs(8), 2.0);
238
239        let first = policy.decision_for_category(ErrorCategory::Network, 0, None);
240        assert!(first.retryable);
241        assert_eq!(first.delay, Some(Duration::from_secs(1)));
242
243        let exhausted = policy.decision_for_category(ErrorCategory::Network, 1, None);
244        assert!(!exhausted.retryable);
245        assert!(exhausted.delay.is_none());
246    }
247
248    #[test]
249    fn decision_for_category_prefers_retry_after() {
250        let policy = RetryPolicy::from_retries(3, Duration::from_secs(1), Duration::from_secs(8), 2.0);
251
252        let decision = policy.decision_for_category(ErrorCategory::RateLimit, 0, Some(Duration::from_secs(7)));
253        assert!(decision.retryable);
254        assert_eq!(decision.delay, Some(Duration::from_secs(7)));
255        assert_eq!(decision.retry_after, Some(Duration::from_secs(7)));
256    }
257}