codex_client/
retry.rs

1use codex_http_client::Request;
2use codex_http_client::TransportError;
3use rand::Rng;
4use std::future::Future;
5use std::time::Duration;
6use tokio::time::sleep;
7
8#[derive(Debug, Clone)]
9pub struct RetryPolicy {
10    pub max_attempts: u64,
11    pub base_delay: Duration,
12    pub retry_on: RetryOn,
13}
14
15#[derive(Debug, Clone)]
16pub struct RetryOn {
17    pub retry_429: bool,
18    pub retry_5xx: bool,
19    pub retry_transport: bool,
20}
21
22impl RetryOn {
23    pub fn should_retry(&self, err: &TransportError, attempt: u64, max_attempts: u64) -> bool {
24        if attempt >= max_attempts {
25            return false;
26        }
27        match err {
28            TransportError::Http { status, .. } => {
29                (self.retry_429 && status.as_u16() == 429)
30                    || (self.retry_5xx && status.is_server_error())
31            }
32            TransportError::Timeout | TransportError::Network(_) => self.retry_transport,
33            _ => false,
34        }
35    }
36}
37
38pub fn backoff(base: Duration, attempt: u64) -> Duration {
39    if attempt == 0 {
40        return base;
41    }
42    let exp = 2u64.saturating_pow(attempt as u32 - 1);
43    let millis = base.as_millis() as u64;
44    let raw = millis.saturating_mul(exp);
45    let jitter: f64 = rand::rng().random_range(0.9..1.1);
46    Duration::from_millis((raw as f64 * jitter) as u64)
47}
48
49pub async fn run_with_retry<T, F, Fut>(
50    policy: RetryPolicy,
51    mut make_req: impl FnMut() -> Request,
52    op: F,
53) -> Result<T, TransportError>
54where
55    F: Fn(Request, u64) -> Fut,
56    Fut: Future<Output = Result<T, TransportError>>,
57{
58    for attempt in 0..=policy.max_attempts {
59        let req = make_req();
60        match op(req, attempt).await {
61            Ok(resp) => return Ok(resp),
62            Err(err)
63                if policy
64                    .retry_on
65                    .should_retry(&err, attempt, policy.max_attempts) =>
66            {
67                sleep(backoff(policy.base_delay, attempt + 1)).await;
68            }
69            Err(err) => return Err(err),
70        }
71    }
72    Err(TransportError::RetryLimit)
73}
codex_client/retry.rs

codex_client/
retry.rs
