claude_codex/
retry.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2use std::time::Duration;
3
4pub const RETRY_INITIAL_DELAY_MS: u64 = 2000;
5pub const RETRY_MAX_DELAY_MS: u64 = 30_000;
6pub const RETRY_BACKOFF_FACTOR: u64 = 2;
7pub const MAX_RATE_LIMIT_RETRIES: u32 = 3;
8
9#[derive(Debug, Clone, Copy)]
10pub struct BackoffOutcome {
11    pub wait_ms: u64,
12    pub exceeds_budget: bool,
13}
14
15pub fn should_retry_status(status: u16) -> bool {
16    matches!(status, 429 | 500 | 502 | 503 | 504)
17}
18
19pub fn compute_backoff_delay(attempt: u32, retry_after: Option<&str>) -> BackoffOutcome {
20    if let Some(raw) = retry_after
21        && let Ok(raw_secs) = raw.parse::<f64>()
22    {
23        let target_ms = (raw_secs * 1000.0).ceil() as u64;
24        return BackoffOutcome {
25            wait_ms: target_ms.min(RETRY_MAX_DELAY_MS),
26            exceeds_budget: target_ms > RETRY_MAX_DELAY_MS,
27        };
28    }
29
30    let mut exp =
31        RETRY_INITIAL_DELAY_MS.saturating_mul(RETRY_BACKOFF_FACTOR.saturating_pow(attempt));
32    if exp > RETRY_MAX_DELAY_MS {
33        exp = RETRY_MAX_DELAY_MS;
34    }
35    let jitter = exp / 2;
36    let wait_ms = (exp / 2) + (jitter / 2);
37    BackoffOutcome {
38        wait_ms,
39        exceeds_budget: false,
40    }
41}
42
43static ZERO_RETRY_DELAY_FOR_TESTS: AtomicBool = AtomicBool::new(false);
44
45/// Make retry sleeps return immediately so exhaustion paths can be exercised
46/// in tests without waiting out the real backoff schedule.
47pub fn set_zero_retry_delay_for_tests(enabled: bool) {
48    ZERO_RETRY_DELAY_FOR_TESTS.store(enabled, Ordering::SeqCst);
49}
50
51pub async fn sleep(ms: u64) {
52    if ZERO_RETRY_DELAY_FOR_TESTS.load(Ordering::SeqCst) {
53        return;
54    }
55    tokio::time::sleep(Duration::from_millis(ms)).await;
56}
57
58#[cfg(test)]
59pub async fn retry_on_statuses<T, E, F>(mut next: F) -> Result<T, E>
60where
61    E: std::fmt::Debug,
62    F: FnMut(u32) -> Result<T, E>,
63{
64    let mut attempt = 0;
65    loop {
66        attempt += 1;
67        if attempt > MAX_RATE_LIMIT_RETRIES + 1 {
68            break;
69        }
70        match next(attempt) {
71            Ok(value) => return Ok(value),
72            Err(err) if attempt <= MAX_RATE_LIMIT_RETRIES + 1 => {
73                if attempt > MAX_RATE_LIMIT_RETRIES {
74                    return Err(err);
75                }
76                sleep(compute_backoff_delay(attempt, None).wait_ms).await;
77            }
78            Err(err) => return Err(err),
79        }
80    }
81    unreachable!()
82}
claude_codex/retry.rs

claude_codex/
retry.rs
