Skip to main content

claude_codex/
retry.rs

1use std::time::Duration;
2
3pub const RETRY_INITIAL_DELAY_MS: u64 = 2000;
4pub const RETRY_MAX_DELAY_MS: u64 = 30_000;
5pub const RETRY_BACKOFF_FACTOR: u64 = 2;
6pub const MAX_RATE_LIMIT_RETRIES: u32 = 3;
7
8#[derive(Debug, Clone, Copy)]
9pub struct BackoffOutcome {
10    pub wait_ms: u64,
11    pub exceeds_budget: bool,
12}
13
14pub fn should_retry_status(status: u16) -> bool {
15    matches!(status, 429 | 500 | 502 | 503 | 504)
16}
17
18pub fn compute_backoff_delay(attempt: u32, retry_after: Option<&str>) -> BackoffOutcome {
19    if let Some(raw) = retry_after
20        && let Ok(raw_secs) = raw.parse::<f64>()
21    {
22        let target_ms = (raw_secs * 1000.0).ceil() as u64;
23        return BackoffOutcome {
24            wait_ms: target_ms.min(RETRY_MAX_DELAY_MS),
25            exceeds_budget: target_ms > RETRY_MAX_DELAY_MS,
26        };
27    }
28
29    let mut exp =
30        RETRY_INITIAL_DELAY_MS.saturating_mul(RETRY_BACKOFF_FACTOR.saturating_pow(attempt));
31    if exp > RETRY_MAX_DELAY_MS {
32        exp = RETRY_MAX_DELAY_MS;
33    }
34    let jitter = exp / 2;
35    let wait_ms = (exp / 2) + (jitter / 2);
36    BackoffOutcome {
37        wait_ms,
38        exceeds_budget: false,
39    }
40}
41
42pub async fn sleep(ms: u64) {
43    tokio::time::sleep(Duration::from_millis(ms)).await;
44}
45
46#[cfg(test)]
47pub async fn retry_on_statuses<T, E, F>(mut next: F) -> Result<T, E>
48where
49    E: std::fmt::Debug,
50    F: FnMut(u32) -> Result<T, E>,
51{
52    let mut attempt = 0;
53    loop {
54        attempt += 1;
55        if attempt > MAX_RATE_LIMIT_RETRIES + 1 {
56            break;
57        }
58        match next(attempt) {
59            Ok(value) => return Ok(value),
60            Err(err) if attempt <= MAX_RATE_LIMIT_RETRIES + 1 => {
61                if attempt > MAX_RATE_LIMIT_RETRIES {
62                    return Err(err);
63                }
64                sleep(compute_backoff_delay(attempt, None).wait_ms).await;
65            }
66            Err(err) => return Err(err),
67        }
68    }
69    unreachable!()
70}