typesafe_ai_rs/
retry.rs

1//! Retry policy, exponential backoff, and server retry hints.
2
3use std::{
4    collections::BTreeSet,
5    fmt,
6    sync::Arc,
7    time::{Duration, Instant, SystemTime},
8};
9
10use reqwest::header::HeaderMap;
11
12use crate::Error;
13
14/// An additional rule that can opt an error into retries.
15pub type RetryPredicate = Arc<dyn Fn(&Error) -> bool + Send + Sync>;
16
17/// Configuration shared by synchronous and asynchronous request retries.
18///
19/// `max_retries` counts retries after the initial attempt. The optional `timeout`
20/// budget includes requests and waits; a retry whose delay reaches the budget is
21/// skipped. Request timeouts are configured separately on the client.
22#[derive(Clone)]
23pub struct RetryPolicy {
24    /// Maximum retries after the initial attempt. Zero disables retries.
25    pub max_retries: u32,
26    /// First exponential backoff delay. Zero disables backoff.
27    pub backoff_initial: Duration,
28    /// Maximum exponential backoff delay. Zero disables backoff.
29    pub backoff_max: Duration,
30    /// Fraction of each backoff randomly subtracted, between zero and one.
31    pub backoff_jitter: f64,
32    /// HTTP status codes that trigger a retry.
33    pub http_statuses: BTreeSet<u16>,
34    /// Honor `Retry-After` and `retry-after-ms` response headers.
35    pub respect_retry_after: bool,
36    /// Maximum accepted server delay. Larger hints fall back to backoff.
37    /// `None` accepts any representable server delay.
38    pub max_retry_after: Option<Duration>,
39    /// Retry request connection and response-body delivery failures.
40    pub api_connection_error: bool,
41    /// Retry request timeouts, independently from connection failures.
42    pub api_timeout_error: bool,
43    /// Optional additional retry rule; built-in rules still apply.
44    /// Cancellation is never retried, including by this predicate.
45    pub predicate: Option<RetryPredicate>,
46    /// Total retry budget including the first request, or `None` for no limit.
47    pub timeout: Option<Duration>,
48}
49
50impl Default for RetryPolicy {
51    fn default() -> Self {
52        Self {
53            max_retries: 2,
54            backoff_initial: Duration::from_millis(500),
55            backoff_max: Duration::from_secs(5),
56            backoff_jitter: 0.25,
57            http_statuses: [408, 429].into_iter().chain(500..600).collect(),
58            respect_retry_after: true,
59            max_retry_after: Some(Duration::from_secs(60)),
60            api_connection_error: true,
61            api_timeout_error: true,
62            predicate: None,
63            timeout: Some(Duration::from_secs(30)),
64        }
65    }
66}
67
68impl fmt::Debug for RetryPolicy {
69    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
70        f.debug_struct("RetryPolicy")
71            .field("max_retries", &self.max_retries)
72            .field("backoff_initial", &self.backoff_initial)
73            .field("backoff_max", &self.backoff_max)
74            .field("backoff_jitter", &self.backoff_jitter)
75            .field("http_statuses", &self.http_statuses)
76            .field("respect_retry_after", &self.respect_retry_after)
77            .field("max_retry_after", &self.max_retry_after)
78            .field("api_connection_error", &self.api_connection_error)
79            .field("api_timeout_error", &self.api_timeout_error)
80            .field("predicate", &self.predicate.as_ref().map(|_| "<predicate>"))
81            .field("timeout", &self.timeout)
82            .finish()
83    }
84}
85
86impl RetryPolicy {
87    /// A policy with no retries, retaining the other defaults.
88    pub fn no_retries() -> Self {
89        Self {
90            max_retries: 0,
91            ..Self::default()
92        }
93    }
94
95    /// Validate jitter, status codes, and timer-compatible delays and budget.
96    pub fn validate(&self) -> Result<(), Error> {
97        if !(0.0..=1.0).contains(&self.backoff_jitter) {
98            return Err(Error::Configuration(
99                "backoff_jitter must be between zero and one".into(),
100            ));
101        }
102        if self.timeout.is_some_and(|timeout| timeout.is_zero()) {
103            return Err(Error::Configuration(
104                "retry timeout must be greater than zero".into(),
105            ));
106        }
107        if self
108            .http_statuses
109            .iter()
110            .any(|status| !(100..=999).contains(status))
111        {
112            return Err(Error::Configuration(
113                "retry http_statuses must contain HTTP codes between 100 and 999".into(),
114            ));
115        }
116        let now = Instant::now();
117        for (name, duration) in [
118            ("backoff_initial", Some(self.backoff_initial)),
119            ("backoff_max", Some(self.backoff_max)),
120            ("max_retry_after", self.max_retry_after),
121            ("retry timeout", self.timeout),
122        ] {
123            if duration.is_some_and(|duration| now.checked_add(duration).is_none()) {
124                return Err(Error::Configuration(format!(
125                    "{name} is too large for the platform's timer"
126                )));
127            }
128        }
129        Ok(())
130    }
131
132    /// Whether the error matches a built-in or additional retry rule.
133    ///
134    /// This checks error classification only. The request loop separately applies
135    /// the attempt limit and retry budget.
136    pub fn should_retry(&self, error: &Error) -> bool {
137        let builtin = match error {
138            Error::Cancelled => return false,
139            Error::Timeout { .. } => self.api_timeout_error,
140            Error::Connection(_) => self.api_connection_error,
141            Error::Api(error) => self.http_statuses.contains(&error.status.as_u16()),
142            Error::ResponseValidation { response, .. } => {
143                self.http_statuses.contains(&response.status.as_u16())
144            }
145            _ => false,
146        };
147        builtin
148            || self
149                .predicate
150                .as_ref()
151                .is_some_and(|predicate| predicate(error))
152    }
153
154    /// Delay before a zero-based retry attempt, honoring permitted server hints.
155    ///
156    /// Attempt zero uses `backoff_initial`, attempt one twice that value, and so
157    /// on up to `backoff_max`, with random downward jitter. Very large attempt
158    /// numbers and header values are handled without arithmetic overflow.
159    pub fn delay(&self, attempt: u32, headers: Option<&HeaderMap>) -> Duration {
160        if self.respect_retry_after {
161            if let Some(delay) = headers.and_then(parse_retry_after) {
162                if self.max_retry_after.is_none_or(|maximum| delay <= maximum) {
163                    return delay;
164                }
165            }
166        }
167        self.backoff(attempt, fastrand::f64())
168    }
169
170    fn backoff(&self, attempt: u32, random: f64) -> Duration {
171        let nanoseconds = self
172            .backoff_initial
173            .as_nanos()
174            .saturating_mul(2_u128.saturating_pow(attempt))
175            .min(self.backoff_max.as_nanos());
176        let exponential = Duration::new(
177            (nanoseconds / 1_000_000_000) as u64,
178            (nanoseconds % 1_000_000_000) as u32,
179        );
180        let seconds = exponential.as_secs_f64() * (1.0 - random * self.backoff_jitter);
181        Duration::try_from_secs_f64(seconds)
182            .unwrap_or(exponential)
183            .min(exponential)
184    }
185}
186
187/// Parse a server retry delay, preferring `retry-after-ms` over `Retry-After`.
188///
189/// Supports nonnegative fractional milliseconds or seconds and HTTP dates. An
190/// expired date means zero delay. Invalid, non-finite, and unrepresentable values
191/// are ignored; an invalid millisecond hint falls back to the seconds/date hint.
192pub fn parse_retry_after(headers: &HeaderMap) -> Option<Duration> {
193    parse_retry_after_at(headers, SystemTime::now())
194}
195
196fn parse_retry_after_at(headers: &HeaderMap, now: SystemTime) -> Option<Duration> {
197    if let Some(delay) = headers
198        .get("retry-after-ms")
199        .and_then(|value| value.to_str().ok())
200        .and_then(|raw| parse_seconds(raw, 0.001))
201    {
202        return Some(delay);
203    }
204    let raw = headers.get("retry-after")?.to_str().ok()?.trim();
205    if let Some(delay) = parse_seconds(raw, 1.0) {
206        return Some(delay);
207    }
208    let date = httpdate::parse_http_date(raw).ok()?;
209    timer_duration(date.duration_since(now).unwrap_or(Duration::ZERO))
210}
211
212fn parse_seconds(raw: &str, multiplier: f64) -> Option<Duration> {
213    let raw = raw.trim();
214    let value = if raw.is_empty() {
215        0.0
216    } else {
217        raw.parse::<f64>().ok()?
218    };
219    if !value.is_finite() || value < 0.0 {
220        return None;
221    }
222    timer_duration(Duration::try_from_secs_f64(value * multiplier).ok()?)
223}
224
225fn timer_duration(duration: Duration) -> Option<Duration> {
226    Instant::now().checked_add(duration).map(|_| duration)
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use crate::ApiError;
233    use reqwest::StatusCode;
234    use serde_json::Value;
235
236    fn headers(ms: Option<&str>, seconds: Option<&str>) -> HeaderMap {
237        let mut headers = HeaderMap::new();
238        if let Some(ms) = ms {
239            headers.insert("retry-after-ms", ms.parse().unwrap());
240        }
241        if let Some(seconds) = seconds {
242            headers.insert("retry-after", seconds.parse().unwrap());
243        }
244        headers
245    }
246
247    #[test]
248    fn parses_fractional_delays_and_prefers_milliseconds() {
249        assert_eq!(
250            parse_retry_after(&headers(Some("250.5"), Some("3"))),
251            Some(Duration::from_micros(250_500))
252        );
253        assert_eq!(
254            parse_retry_after(&headers(None, Some("1.25"))),
255            Some(Duration::from_millis(1250))
256        );
257        assert_eq!(
258            parse_retry_after(&headers(Some(" "), Some("3"))),
259            Some(Duration::ZERO)
260        );
261        assert_eq!(
262            parse_retry_after(&headers(Some("invalid"), Some("3"))),
263            Some(Duration::from_secs(3))
264        );
265    }
266
267    #[test]
268    fn handles_hostile_header_numbers_without_panicking() {
269        for raw in ["NaN", "inf", "-inf", "1e300", "-1", "1e99999999999"] {
270            assert_eq!(parse_retry_after(&headers(None, Some(raw))), None, "{raw}");
271            assert_eq!(
272                parse_retry_after(&headers(Some(raw), Some("2"))),
273                Some(Duration::from_secs(2)),
274                "{raw}"
275            );
276        }
277        assert_eq!(
278            parse_retry_after(&headers(None, Some("18446744073709549568"))),
279            None
280        );
281    }
282
283    #[test]
284    fn parses_http_dates_and_clamps_past_dates() {
285        let date = "Wed, 21 Oct 2015 07:28:00 GMT";
286        let timestamp = httpdate::parse_http_date(date).unwrap();
287        let headers = headers(None, Some(date));
288        assert_eq!(
289            parse_retry_after_at(&headers, timestamp - Duration::from_secs(5)),
290            Some(Duration::from_secs(5))
291        );
292        assert_eq!(
293            parse_retry_after_at(&headers, timestamp + Duration::from_secs(5)),
294            Some(Duration::ZERO)
295        );
296    }
297
298    #[test]
299    fn exponential_backoff_caps_and_supports_disabled_waits() {
300        let mut policy = RetryPolicy {
301            backoff_jitter: 0.0,
302            ..RetryPolicy::default()
303        };
304        for (attempt, milliseconds) in [(0, 500), (1, 1000), (2, 2000), (4, 5000), (u32::MAX, 5000)]
305        {
306            assert_eq!(
307                policy.delay(attempt, None),
308                Duration::from_millis(milliseconds)
309            );
310        }
311        policy.backoff_initial = Duration::ZERO;
312        assert_eq!(policy.delay(u32::MAX, None), Duration::ZERO);
313        policy.backoff_initial = Duration::MAX;
314        policy.backoff_max = Duration::MAX;
315        assert_eq!(policy.delay(u32::MAX, None), Duration::MAX);
316        policy.backoff_max = Duration::ZERO;
317        assert_eq!(policy.delay(u32::MAX, None), Duration::ZERO);
318    }
319
320    #[test]
321    fn jitter_only_reduces_capped_backoff() {
322        let policy = RetryPolicy::default();
323        assert_eq!(policy.backoff(0, 0.0), Duration::from_millis(500));
324        assert_eq!(policy.backoff(0, 1.0), Duration::from_millis(375));
325        assert_eq!(policy.backoff(10, 1.0), Duration::from_millis(3750));
326    }
327
328    #[test]
329    fn server_hints_respect_the_configured_limit() {
330        let mut policy = RetryPolicy {
331            backoff_jitter: 0.0,
332            ..RetryPolicy::default()
333        };
334        assert_eq!(
335            policy.delay(0, Some(&headers(None, Some("60")))),
336            Duration::from_secs(60)
337        );
338        assert_eq!(
339            policy.delay(0, Some(&headers(None, Some("61")))),
340            Duration::from_millis(500)
341        );
342        policy.max_retry_after = None;
343        assert_eq!(
344            policy.delay(0, Some(&headers(None, Some("61")))),
345            Duration::from_secs(61)
346        );
347        policy.respect_retry_after = false;
348        assert_eq!(
349            policy.delay(0, Some(&headers(None, Some("10")))),
350            Duration::from_millis(500)
351        );
352    }
353
354    #[test]
355    fn validates_jitter_and_budget() {
356        for jitter in [f64::NAN, f64::INFINITY, -0.01, 1.01] {
357            assert!(RetryPolicy {
358                backoff_jitter: jitter,
359                ..RetryPolicy::default()
360            }
361            .validate()
362            .is_err());
363        }
364        assert!(RetryPolicy {
365            timeout: Some(Duration::ZERO),
366            ..RetryPolicy::default()
367        }
368        .validate()
369        .is_err());
370        assert!(RetryPolicy {
371            timeout: None,
372            backoff_jitter: 1.0,
373            ..RetryPolicy::default()
374        }
375        .validate()
376        .is_ok());
377        assert!(RetryPolicy {
378            backoff_initial: Duration::MAX,
379            ..RetryPolicy::default()
380        }
381        .validate()
382        .is_err());
383        assert!(RetryPolicy {
384            backoff_max: Duration::MAX,
385            ..RetryPolicy::default()
386        }
387        .validate()
388        .is_err());
389        assert!(RetryPolicy {
390            max_retry_after: Some(Duration::MAX),
391            ..RetryPolicy::default()
392        }
393        .validate()
394        .is_err());
395        assert!(RetryPolicy {
396            timeout: Some(Duration::MAX),
397            ..RetryPolicy::default()
398        }
399        .validate()
400        .is_err());
401        assert!(RetryPolicy {
402            http_statuses: [99, 1000].into_iter().collect(),
403            ..RetryPolicy::default()
404        }
405        .validate()
406        .is_err());
407    }
408
409    #[test]
410    fn status_rules_and_custom_predicates_are_additive_but_never_cancelled() {
411        let mut policy = RetryPolicy::default();
412        for (status, retry) in [
413            (400, false),
414            (401, false),
415            (408, true),
416            (429, true),
417            (500, true),
418            (599, true),
419        ] {
420            let error = ApiError::new(
421                StatusCode::from_u16(status).unwrap(),
422                Value::Null,
423                HeaderMap::new(),
424                None,
425            )
426            .into();
427            assert_eq!(policy.should_retry(&error), retry, "{status}");
428        }
429        let error = Error::InvalidRequest("custom".into());
430        assert!(!policy.should_retry(&error));
431        policy.predicate = Some(Arc::new(|_| true));
432        assert!(policy.should_retry(&error));
433        assert!(!policy.should_retry(&Error::Cancelled));
434    }
435
436    #[test]
437    fn successful_validation_errors_require_explicit_retry_rules() {
438        let error = Error::ResponseValidation {
439            field_path: "answers".into(),
440            response: Box::new(crate::RawResponse {
441                status: StatusCode::OK,
442                headers: HeaderMap::new(),
443                body: Default::default(),
444            }),
445        };
446        let mut policy = RetryPolicy::default();
447        assert!(!policy.should_retry(&error));
448        policy.http_statuses.insert(200);
449        assert!(policy.should_retry(&error));
450    }
451}
typesafe_ai_rs/retry.rs

typesafe_ai_rs/
retry.rs
