Skip to main content

molo_core/provider/
retry.rs

1//! Retry wrapper: `RetryProvider` implements [`Provider`] by wrapping an inner
2//! provider and retrying recoverable failures per policy — rate limits /
3//! network / timeouts / 5xx.
4//!
5//! Retrying is a wrapper layer outside the Provider: the interface itself has
6//! no retry logic, and any Provider implementation (OpenAi / Fake / a
7//! user-written one) gains retry by being wrapped, with zero changes to
8//! callers or the loop layer.
9//!
10//! # Examples
11//!
12//! ```rust
13//! # extern crate molo_core as molo;
14//! # #[tokio::main]
15//! # async fn main() -> Result<(), molo::provider::ProviderError> {
16//! use molo::message::Message;
17//! use molo::provider::{FakeProvider, FakeReply, Provider, RetryProvider, ProviderError};
18//!
19//! // Script: the first call is rate limited, the second succeeds — RetryProvider
20//! // retries automatically and the caller sees nothing.
21//! let inner = FakeProvider::new([
22//!     FakeReply::Error(ProviderError::RateLimited { retry_after: None }),
23//!     FakeReply::Text("hi".into()),
24//! ]);
25//! let provider = RetryProvider::new(inner);
26//! let resp = provider.chat(molo::provider::ChatRequest::default()).await?;
27//! assert_eq!(resp.message, Message::assistant("hi"));
28//! # Ok(())
29//! # }
30//! ```
31
32use super::{
33    ChatRequest, ChatResponse, Provider, ProviderCapabilities, ProviderError,
34    ProviderRequestContext, StreamEvent, TimeoutStage,
35};
36use async_trait::async_trait;
37use futures::stream::BoxStream;
38use std::sync::Arc;
39use std::sync::atomic::{AtomicU64, Ordering};
40use std::time::{Duration, SystemTime, UNIX_EPOCH};
41
42/// Retry policy (with defaults; `Default` is "exponential backoff + jitter,
43/// 3 attempts").
44///
45/// Default backoff: initial 0.5s / factor 2.0 / cap 10s / jitter on (full
46/// jitter, to prevent thundering herds).
47#[derive(Debug, Clone, PartialEq)]
48#[non_exhaustive]
49pub struct RetryPolicy {
50    /// Total attempts (including the first); default 3 (i.e. at most 2
51    /// retries after a failure).
52    pub(crate) max_attempts: usize,
53    /// Backoff strategy (how long to wait before retrying after each
54    /// failure).
55    pub(crate) backoff: Backoff,
56    /// Which errors are retryable; the default is [`Retryable::Default`].
57    pub(crate) retryable: Retryable,
58    /// When rate limited, prefer waiting the vendor's `Retry-After` duration
59    /// (overrides backoff); default true.
60    pub(crate) respect_retry_after: bool,
61}
62
63impl Default for RetryPolicy {
64    fn default() -> Self {
65        Self {
66            max_attempts: 3,
67            backoff: Backoff::Exponential {
68                initial: Duration::from_millis(500),
69                factor: 2.0,
70                max: Duration::from_secs(10),
71                jitter: true,
72            },
73            retryable: Retryable::Default,
74            respect_retry_after: true,
75        }
76    }
77}
78
79impl RetryPolicy {
80    /// Constructs the default retry policy.
81    pub fn new() -> Self {
82        Self::default()
83    }
84
85    /// Total attempts, including the first call.
86    pub fn max_attempts(&self) -> usize {
87        self.max_attempts
88    }
89
90    /// Returns a policy with an updated attempt limit.
91    pub fn with_max_attempts(mut self, max_attempts: usize) -> Self {
92        self.max_attempts = max_attempts;
93        self
94    }
95
96    /// Backoff strategy used between retry attempts.
97    pub fn backoff(&self) -> &Backoff {
98        &self.backoff
99    }
100
101    /// Returns a policy with an updated backoff strategy.
102    pub fn with_backoff(mut self, backoff: Backoff) -> Self {
103        self.backoff = backoff;
104        self
105    }
106
107    /// Retryability predicate.
108    pub fn retryable(&self) -> &Retryable {
109        &self.retryable
110    }
111
112    /// Returns a policy with an updated retryability predicate.
113    pub fn with_retryable(mut self, retryable: Retryable) -> Self {
114        self.retryable = retryable;
115        self
116    }
117
118    /// Whether vendor-provided `Retry-After` durations override backoff.
119    pub fn respect_retry_after(&self) -> bool {
120        self.respect_retry_after
121    }
122
123    /// Returns a policy with updated `Retry-After` handling.
124    pub fn with_respect_retry_after(mut self, respect_retry_after: bool) -> Self {
125        self.respect_retry_after = respect_retry_after;
126        self
127    }
128}
129
130/// Backoff strategy: how long to wait after each failure before the next
131/// attempt.
132#[derive(Debug, Clone, PartialEq)]
133pub enum Backoff {
134    /// Fixed interval.
135    Fixed(Duration),
136    /// Exponential backoff: `initial * factor^attempt`, capped at `max`;
137    /// when `jitter` is on, draw full jitter in `[0, computed value)`
138    /// (thundering-herd protection — requests failing at the same time do
139    /// not all retry at the same moment).
140    Exponential {
141        /// Initial wait duration.
142        initial: Duration,
143        /// Exponential factor.
144        factor: f64,
145        /// Cap duration.
146        max: Duration,
147        /// Whether to use full jitter (`[0, computed value)`,
148        /// thundering-herd protection).
149        jitter: bool,
150    },
151}
152
153impl Backoff {
154    /// The duration to wait after the `attempt`-th failure (0 = the first
155    /// failure).
156    fn delay(&self, attempt: usize) -> Duration {
157        match self {
158            Self::Fixed(d) => *d,
159            Self::Exponential {
160                initial,
161                factor,
162                max,
163                jitter,
164            } => {
165                let base =
166                    (initial.as_secs_f64() * factor.powf(attempt as f64)).min(max.as_secs_f64());
167                if *jitter {
168                    // Full jitter: LCG pseudo-random (zero new dependencies;
169                    // the state is atomically shared so concurrent draws
170                    // advance sequentially — concurrently failing requests
171                    // get dispersed delays, making the herd protection work).
172                    Duration::from_secs_f64(base * random01())
173                } else {
174                    Duration::from_secs_f64(base)
175                }
176            }
177        }
178    }
179}
180
181/// `Retry-After` wait cap (5 minutes): if the vendor's instruction exceeds
182/// this, wait this long instead.
183/// Backoff itself is capped by `max`, but the Retry-After path does not go
184/// through backoff — an absurd value from a broken endpoint (e.g. 136 years)
185/// would make the caller wait forever, hence the separate cap.
186const RETRY_AFTER_CAP: Duration = Duration::from_secs(300);
187
188/// Lightweight pseudo-random source (LCG; zero new dependencies, serves only
189/// the thundering-herd jitter, not for cryptography).
190///
191/// Seed = clock nanoseconds at the first draw (full 64 bits); each subsequent
192/// draw advances the LCG state (atomically shared, so concurrent draws
193/// advance sequentially). The clock appears only in the initial seed and does
194/// not carry the randomness itself — when a rate-limit burst hits,
195/// simultaneously failing requests draw from sequentially advanced
196/// independent states, so delays are highly dispersed.
197fn random01() -> f64 {
198    const A: u64 = 6364136223846793005; // LCG constants (Knuth / MMIX)
199    const C: u64 = 1442695040888963407;
200    static STATE: AtomicU64 = AtomicU64::new(0);
201    loop {
202        let current = STATE.load(Ordering::Relaxed);
203        let next = if current == 0 {
204            // First bootstrap: seed with full epoch nanoseconds (the clock
205            // only contributes the initial difference and does not take part
206            // in later draws), and **advance the LCG immediately** — otherwise
207            // the high 24 bits of the seed are often 0 and the first draw
208            // always returns 0 (the first retry would have no backoff,
209            // breaking herd protection at the very first burst).
210            let seed = SystemTime::now()
211                .duration_since(UNIX_EPOCH)
212                .unwrap_or_default()
213                .as_nanos() as u64;
214            seed.wrapping_mul(A).wrapping_add(C)
215        } else {
216            current.wrapping_mul(A).wrapping_add(C)
217        };
218        if STATE
219            .compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed)
220            .is_ok()
221        {
222            // Take the high 24 bits mapped to [0, 1).
223            return (next >> 40) as f64 / (1u64 << 24) as f64;
224        }
225    }
226}
227
228/// Retry judgment: which errors are worth retrying.
229///
230/// `Custom` closures are not comparable, so `PartialEq` is implemented by
231/// hand (only `Default` values equal each other); `Debug` is also
232/// hand-written (a closure only prints its variant name).
233#[derive(Clone)]
234pub enum Retryable {
235    /// Default: Network / Timeout / RateLimited / Api (status >= 500);
236    /// 4xx (auth / invalid arguments / quota exhausted) are not retried —
237    /// retrying would not change the outcome.
238    Default,
239    /// Custom judgment (e.g. only retry rate limits, or exclude certain 4xx
240    /// by message content).
241    Custom(Arc<dyn Fn(&ProviderError) -> bool + Send + Sync>),
242}
243
244impl std::fmt::Debug for Retryable {
245    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
246        match self {
247            Self::Default => f.write_str("Default"),
248            Self::Custom(_) => f.write_str("Custom(_)"),
249        }
250    }
251}
252
253impl PartialEq for Retryable {
254    fn eq(&self, other: &Self) -> bool {
255        matches!((self, other), (Self::Default, Self::Default))
256    }
257}
258
259impl Retryable {
260    fn is_retryable(&self, error: &ProviderError) -> bool {
261        match self {
262            Self::Default => {
263                matches!(
264                    error,
265                    ProviderError::Network(_)
266                        | ProviderError::Timeout(_)
267                        | ProviderError::RateLimited { .. }
268                ) || matches!(error, ProviderError::Api { status, .. } if *status >= 500)
269            }
270            Self::Custom(f) => f(error),
271        }
272    }
273}
274
275/// Retry wrapper: implements [`Provider`] and retries inner failures per
276/// [`RetryPolicy`].
277///
278/// **Streaming semantics**: only when `stream_chat` returns `Err` (connection
279/// setup failure) is the whole call retried; **errors inside the stream are
280/// passed through verbatim and never retried** — once the stream is
281/// established this implementation cannot tell an "interruption before the
282/// first event" from an "interruption after part of the output was
283/// delivered", and retrying would duplicate or corrupt output.
284/// Hence "retry before the first event" only covers method-level failures;
285/// interruptions within the stream (including after setup but before the
286/// first event) are left to the caller (the Agent) to handle.
287///
288/// Timeouts are the inner provider's (for example `OpenAiProvider`)
289/// responsibility: each attempt may hit the inner `Timeout` error, and
290/// `Retryable::Default` includes Timeout, so "timeouts are retried too"
291/// falls out naturally.
292///
293/// # Errors
294///
295/// After retries are exhausted, the **last** error is returned (neither the
296/// first nor an aggregate); non-retryable errors are returned immediately on
297/// the first failure.
298///
299/// # Cancellation semantics
300///
301/// Dropping the future while waiting on backoff cancels the whole call; no
302/// further attempts are made. An already-issued inner request is cancelled
303/// together with the future (whether the network request continues at the
304/// transport layer depends on the underlying client).
305pub struct RetryProvider<I> {
306    inner: I,
307    policy: RetryPolicy,
308}
309
310impl<I> std::fmt::Debug for RetryProvider<I> {
311    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
312        // inner is a generic Provider without Debug; print the type name and
313        // policy (in the same style as ReActAgent).
314        f.debug_struct("RetryProvider")
315            .field("inner", &std::any::type_name::<I>())
316            .field("policy", &self.policy)
317            .finish()
318    }
319}
320
321impl<I> RetryProvider<I> {
322    /// Wraps with the default policy (3 attempts / exponential backoff +
323    /// jitter / default retry judgment).
324    pub fn new(inner: I) -> Self {
325        Self {
326            inner,
327            policy: RetryPolicy::default(),
328        }
329    }
330
331    /// Replaces the retry policy.
332    ///
333    /// # Examples
334    ///
335    /// For tests / local simulation: fixed short waits, at most 2 attempts,
336    /// ignoring the vendor's `Retry-After`:
337    ///
338    /// ```
339    /// # extern crate molo_core as molo;
340    /// use std::time::Duration;
341    /// use molo::{Backoff, FakeProvider, RetryPolicy, RetryProvider};
342    ///
343    /// let provider = RetryProvider::new(FakeProvider::new([])).with_policy(
344    ///     RetryPolicy::default()
345    ///         .with_max_attempts(2)
346    ///         .with_backoff(Backoff::Fixed(Duration::from_millis(10)))
347    ///         .with_respect_retry_after(false),
348    /// );
349    /// ```
350    pub fn with_policy(mut self, policy: RetryPolicy) -> Self {
351        self.policy = policy;
352        self
353    }
354}
355
356impl<I> RetryProvider<I> {
357    /// Retry decision and wait for one failure: whether another attempt is
358    /// possible, and if so how long to wait.
359    ///
360    /// `Some(delay)` = wait then retry; `None` = give up (not retryable or
361    /// attempts exhausted).
362    fn retry_decision(&self, error: &ProviderError, attempts: usize) -> Option<Duration> {
363        if !self.policy.retryable.is_retryable(error) || attempts + 1 >= self.policy.max_attempts {
364            return None;
365        }
366        let delay = match (error, self.policy.respect_retry_after) {
367            // Retry-After cap: a broken endpoint may return an absurd value
368            // (e.g. 136 years), and backoff's own `max` cap does not apply on
369            // this path — cap at [`RETRY_AFTER_CAP`].
370            (
371                ProviderError::RateLimited {
372                    retry_after: Some(d),
373                },
374                true,
375            ) => (*d).min(RETRY_AFTER_CAP),
376            _ => self.policy.backoff.delay(attempts),
377        };
378        #[cfg(feature = "tracing")]
379        tracing::warn!(
380            attempt = attempts + 1,
381            max_attempts = self.policy.max_attempts,
382            delay = ?delay,
383            error = %error,
384            "provider call failed, retrying",
385        );
386        Some(delay)
387    }
388}
389
390#[async_trait]
391impl<I: Provider + Send + Sync> Provider for RetryProvider<I> {
392    fn model(&self) -> Option<&str> {
393        self.inner.model()
394    }
395
396    fn capabilities(&self) -> ProviderCapabilities {
397        self.inner.capabilities()
398    }
399
400    async fn chat(&self, request: ChatRequest) -> Result<ChatResponse, ProviderError> {
401        let mut attempts = 0usize;
402        loop {
403            match self.inner.chat(request.clone()).await {
404                Ok(response) => return Ok(response),
405                Err(error) => match self.retry_decision(&error, attempts) {
406                    Some(delay) => {
407                        tokio::time::sleep(delay).await;
408                        attempts += 1;
409                    }
410                    None => return Err(error),
411                },
412            }
413        }
414    }
415
416    async fn chat_with_context(
417        &self,
418        request: ChatRequest,
419        context: &ProviderRequestContext,
420    ) -> Result<ChatResponse, ProviderError> {
421        let mut attempts = 0usize;
422        loop {
423            check_context(context)?;
424            match self.inner.chat_with_context(request.clone(), context).await {
425                Ok(response) => return Ok(response),
426                Err(error) => match self.retry_decision(&error, attempts) {
427                    Some(delay) => {
428                        sleep_with_context(delay, context).await?;
429                        attempts += 1;
430                    }
431                    None => return Err(error),
432                },
433            }
434        }
435    }
436
437    async fn stream_chat(
438        &self,
439        request: ChatRequest,
440    ) -> Result<BoxStream<'static, Result<StreamEvent, ProviderError>>, ProviderError> {
441        // Retry only when the method returns Err; errors inside the stream
442        // are always passed through (see the type-level docs).
443        let mut attempts = 0usize;
444        loop {
445            match self.inner.stream_chat(request.clone()).await {
446                Ok(stream) => return Ok(stream),
447                Err(error) => match self.retry_decision(&error, attempts) {
448                    Some(delay) => {
449                        tokio::time::sleep(delay).await;
450                        attempts += 1;
451                    }
452                    None => return Err(error),
453                },
454            }
455        }
456    }
457
458    async fn stream_chat_with_context(
459        &self,
460        request: ChatRequest,
461        context: &ProviderRequestContext,
462    ) -> Result<BoxStream<'static, Result<StreamEvent, ProviderError>>, ProviderError> {
463        let mut attempts = 0usize;
464        loop {
465            check_context(context)?;
466            match self
467                .inner
468                .stream_chat_with_context(request.clone(), context)
469                .await
470            {
471                Ok(stream) => return Ok(stream),
472                Err(error) => match self.retry_decision(&error, attempts) {
473                    Some(delay) => {
474                        sleep_with_context(delay, context).await?;
475                        attempts += 1;
476                    }
477                    None => return Err(error),
478                },
479            }
480        }
481    }
482}
483
484fn check_context(context: &ProviderRequestContext) -> Result<(), ProviderError> {
485    if context.is_cancelled() {
486        Err(ProviderError::Cancelled)
487    } else if context.is_expired() {
488        Err(ProviderError::Timeout(TimeoutStage::Request))
489    } else {
490        Ok(())
491    }
492}
493
494async fn sleep_with_context(
495    delay: Duration,
496    context: &ProviderRequestContext,
497) -> Result<(), ProviderError> {
498    check_context(context)?;
499    match context.remaining() {
500        Some(remaining) if remaining.is_zero() => {
501            Err(ProviderError::Timeout(TimeoutStage::Request))
502        }
503        Some(remaining) => {
504            tokio::select! {
505                _ = context.cancellation.cancelled() => Err(ProviderError::Cancelled),
506                _ = tokio::time::sleep(remaining) => Err(ProviderError::Timeout(TimeoutStage::Request)),
507                _ = tokio::time::sleep(delay) => Ok(()),
508            }
509        }
510        None => {
511            tokio::select! {
512                _ = context.cancellation.cancelled() => Err(ProviderError::Cancelled),
513                _ = tokio::time::sleep(delay) => Ok(()),
514            }
515        }
516    }
517}
518
519#[cfg(test)]
520mod tests {
521    use super::*;
522    use crate::provider::{FakeProvider, FakeReply};
523    use std::sync::Arc;
524
525    fn rate_limited() -> ProviderError {
526        ProviderError::RateLimited { retry_after: None }
527    }
528
529    fn api(status: u16) -> ProviderError {
530        ProviderError::Api {
531            status,
532            code: None,
533            message: "boom".into(),
534        }
535    }
536
537    fn network() -> ProviderError {
538        ProviderError::Network("connection refused".into())
539    }
540
541    #[tokio::test]
542    async fn retries_then_succeeds() {
543        // Rate-limit failure → retry succeeds; the caller only sees the
544        // final result.
545        let inner = Arc::new(FakeProvider::new([
546            FakeReply::Error(rate_limited()),
547            FakeReply::Text("ok".into()),
548        ]));
549        let provider = RetryProvider::new(inner.clone());
550        let resp = provider.chat(ChatRequest::default()).await.unwrap();
551        assert_eq!(resp.message, crate::message::Message::assistant("ok"));
552        // Assert attempt count: first failure + one retry = 2 calls.
553        assert_eq!(inner.requests().len(), 2);
554    }
555
556    #[tokio::test]
557    async fn gives_up_after_max_attempts() {
558        // Always fails: return the last error after max_attempts (default 3)
559        // attempts.
560        let inner = Arc::new(FakeProvider::new([
561            FakeReply::Error(network()),
562            FakeReply::Error(network()),
563            FakeReply::Error(network()),
564        ]));
565        let provider = RetryProvider::new(inner.clone());
566        let err = provider.chat(ChatRequest::default()).await.unwrap_err();
567        assert!(matches!(err, ProviderError::Network(_)));
568        assert_eq!(inner.requests().len(), 3);
569    }
570
571    #[tokio::test]
572    async fn api_4xx_not_retried() {
573        // 4xx business error: the default judgment does not retry, giving up
574        // after a single call.
575        let inner = Arc::new(FakeProvider::new([FakeReply::Error(api(400))]));
576        let provider = RetryProvider::new(inner.clone());
577        let err = provider.chat(ChatRequest::default()).await.unwrap_err();
578        assert!(matches!(err, ProviderError::Api { status: 400, .. }));
579        assert_eq!(inner.requests().len(), 1);
580    }
581
582    #[tokio::test]
583    async fn api_5xx_retried() {
584        let inner = Arc::new(FakeProvider::new([
585            FakeReply::Error(api(503)),
586            FakeReply::Text("ok".into()),
587        ]));
588        let provider = RetryProvider::new(inner.clone());
589        provider.chat(ChatRequest::default()).await.unwrap();
590        assert_eq!(inner.requests().len(), 2);
591    }
592
593    #[tokio::test]
594    async fn max_attempts_one_disables_retry() {
595        let inner = Arc::new(FakeProvider::new([FakeReply::Error(network())]));
596        let provider = RetryProvider::new(inner.clone()).with_policy(RetryPolicy {
597            max_attempts: 1,
598            ..Default::default()
599        });
600        provider.chat(ChatRequest::default()).await.unwrap_err();
601        assert_eq!(inner.requests().len(), 1);
602    }
603
604    /// Retry-After cap: a vendor instruction of 3600s is capped at
605    /// `RETRY_AFTER_CAP` (300s) rather than waiting indefinitely.
606    #[test]
607    fn retry_after_capped_at_cap() {
608        let provider = RetryProvider::new(FakeProvider::new([]));
609        let delay = provider.retry_decision(
610            &ProviderError::RateLimited {
611                retry_after: Some(Duration::from_secs(3600)),
612            },
613            0,
614        );
615        assert_eq!(delay, Some(RETRY_AFTER_CAP));
616    }
617
618    #[tokio::test]
619    async fn retry_after_respected_when_present() {
620        // The vendor instructs a 1s wait: the policy prefers it over backoff;
621        // the retry succeeds after the wait.
622        let inner = Arc::new(FakeProvider::new([
623            FakeReply::Error(ProviderError::RateLimited {
624                retry_after: Some(Duration::from_millis(1)),
625            }),
626            FakeReply::Text("ok".into()),
627        ]));
628        let provider = RetryProvider::new(inner.clone());
629        provider.chat(ChatRequest::default()).await.unwrap();
630        assert_eq!(inner.requests().len(), 2);
631    }
632
633    /// Verifies the Retry-After **priority** mechanism: backoff is 50ms
634    /// (factor 1.0) while Retry-After is 200ms — the actual wait must be
635    /// ≈200ms rather than 50ms (if respect were broken, the retry would finish
636    /// after 50ms). 200ms is the shortest reliably measurable window for the
637    /// suite (tokio timers only fire late, never early; the 150ms lower bound
638    /// always holds; the 600ms upper bound leaves 400ms of slack).
639    #[tokio::test]
640    async fn retry_after_delay_overrides_backoff() {
641        use std::time::Instant;
642
643        let inner = Arc::new(FakeProvider::new([
644            FakeReply::Error(ProviderError::RateLimited {
645                retry_after: Some(Duration::from_millis(200)),
646            }),
647            FakeReply::Text("ok".into()),
648        ]));
649        let provider = RetryProvider::new(inner.clone()).with_policy(RetryPolicy {
650            backoff: Backoff::Exponential {
651                initial: Duration::from_millis(50),
652                factor: 1.0,
653                max: Duration::from_secs(1),
654                jitter: false,
655            },
656            ..Default::default()
657        });
658        let start = Instant::now();
659        provider.chat(ChatRequest::default()).await.unwrap();
660        let elapsed = start.elapsed();
661        assert!(
662            elapsed >= Duration::from_millis(150) && elapsed < Duration::from_millis(600),
663            "should wait Retry-After (200ms) rather than 50ms backoff, elapsed {elapsed:?}"
664        );
665    }
666
667    #[tokio::test]
668    async fn custom_retryable_predicate() {
669        // Custom judgment: retry only 4xx other than 429 (reversing the
670        // default).
671        let inner = Arc::new(FakeProvider::new([
672            FakeReply::Error(api(400)),
673            FakeReply::Text("ok".into()),
674        ]));
675        let provider = RetryProvider::new(inner.clone()).with_policy(RetryPolicy {
676            retryable: Retryable::Custom(Arc::new(
677                |e| matches!(e, ProviderError::Api { status, .. } if *status == 400),
678            )),
679            ..Default::default()
680        });
681        provider.chat(ChatRequest::default()).await.unwrap();
682        assert_eq!(inner.requests().len(), 2);
683    }
684
685    #[tokio::test]
686    async fn stream_retries_only_before_first_event() {
687        // Streaming: first setup failure → retry the whole call; after
688        // success, events inside the stream pass through.
689        let inner = Arc::new(FakeProvider::new([
690            FakeReply::Error(network()),
691            FakeReply::Text("ok".into()),
692        ]));
693        let provider = RetryProvider::new(inner.clone());
694        let mut stream = provider.stream_chat(ChatRequest::default()).await.unwrap();
695        use futures::StreamExt;
696        let first = stream.next().await.unwrap().unwrap();
697        assert_eq!(first, StreamEvent::Delta("ok".into()));
698        assert_eq!(inner.requests().len(), 2);
699    }
700
701    #[tokio::test]
702    async fn stream_failure_after_max_attempts() {
703        let inner = Arc::new(FakeProvider::new([FakeReply::Error(network())]));
704        let provider = RetryProvider::new(inner.clone()).with_policy(RetryPolicy {
705            max_attempts: 1,
706            ..Default::default()
707        });
708        // BoxStream has no Debug, so unwrap_err is unavailable; match to get
709        // the error.
710        let err = match provider.stream_chat(ChatRequest::default()).await {
711            Err(e) => e,
712            Ok(_) => panic!("expected error"),
713        };
714        assert!(matches!(err, ProviderError::Network(_)));
715        assert_eq!(inner.requests().len(), 1);
716    }
717
718    #[test]
719    fn exponential_backoff_grows_and_caps() {
720        let backoff = Backoff::Exponential {
721            initial: Duration::from_millis(500),
722            factor: 2.0,
723            max: Duration::from_secs(10),
724            jitter: false,
725        };
726        assert_eq!(backoff.delay(0), Duration::from_millis(500));
727        assert_eq!(backoff.delay(1), Duration::from_secs(1));
728        assert_eq!(backoff.delay(2), Duration::from_secs(2));
729        // Cap: even after 5 failures the delay never exceeds max.
730        assert_eq!(backoff.delay(10), Duration::from_secs(10));
731    }
732
733    #[test]
734    fn jitter_stays_within_bounds() {
735        let backoff = Backoff::Exponential {
736            initial: Duration::from_secs(1),
737            factor: 1.0,
738            max: Duration::from_secs(10),
739            jitter: true,
740        };
741        // Full jitter in [0, base): base is always 1s (factor 1.0), so the
742        // upper bound must be < 1s — a limit of 10s would let a degenerate
743        // implementation (e.g. always returning 0.9) pass.
744        for attempt in 0..50 {
745            let d = backoff.delay(attempt);
746            assert!(d < Duration::from_secs(1), "jitter out of bounds: {d:?}");
747        }
748    }
749
750    /// Verifies jitter's randomness mechanics: the first draw (LCG bootstrap)
751    /// is not always 0 — if it were, the first retry in the process would
752    /// have no backoff; multiple draws must yield multiple distinct values
753    /// (state advancing).
754    #[test]
755    fn jitter_first_draw_nonzero_and_dispersed() {
756        let first = random01();
757        assert!(
758            first > 0.0,
759            "first draw must not be 0 (bootstrap must advance LCG first)"
760        );
761        assert!(first < 1.0);
762
763        let mut seen = std::collections::HashSet::new();
764        for _ in 0..64 {
765            let d = random01();
766            assert!(d > 0.0 && d < 1.0);
767            seen.insert(d.to_bits());
768        }
769        assert!(
770            seen.len() > 1,
771            "LCG draws must be dispersed (state advancing)"
772        );
773    }
774}