1use 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#[derive(Debug, Clone, PartialEq)]
48#[non_exhaustive]
49pub struct RetryPolicy {
50 pub(crate) max_attempts: usize,
53 pub(crate) backoff: Backoff,
56 pub(crate) retryable: Retryable,
58 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 pub fn new() -> Self {
82 Self::default()
83 }
84
85 pub fn max_attempts(&self) -> usize {
87 self.max_attempts
88 }
89
90 pub fn with_max_attempts(mut self, max_attempts: usize) -> Self {
92 self.max_attempts = max_attempts;
93 self
94 }
95
96 pub fn backoff(&self) -> &Backoff {
98 &self.backoff
99 }
100
101 pub fn with_backoff(mut self, backoff: Backoff) -> Self {
103 self.backoff = backoff;
104 self
105 }
106
107 pub fn retryable(&self) -> &Retryable {
109 &self.retryable
110 }
111
112 pub fn with_retryable(mut self, retryable: Retryable) -> Self {
114 self.retryable = retryable;
115 self
116 }
117
118 pub fn respect_retry_after(&self) -> bool {
120 self.respect_retry_after
121 }
122
123 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#[derive(Debug, Clone, PartialEq)]
133pub enum Backoff {
134 Fixed(Duration),
136 Exponential {
141 initial: Duration,
143 factor: f64,
145 max: Duration,
147 jitter: bool,
150 },
151}
152
153impl Backoff {
154 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 Duration::from_secs_f64(base * random01())
173 } else {
174 Duration::from_secs_f64(base)
175 }
176 }
177 }
178 }
179}
180
181const RETRY_AFTER_CAP: Duration = Duration::from_secs(300);
187
188fn random01() -> f64 {
198 const A: u64 = 6364136223846793005; 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 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 return (next >> 40) as f64 / (1u64 << 24) as f64;
224 }
225 }
226}
227
228#[derive(Clone)]
234pub enum Retryable {
235 Default,
239 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
275pub 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 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 pub fn new(inner: I) -> Self {
325 Self {
326 inner,
327 policy: RetryPolicy::default(),
328 }
329 }
330
331 pub fn with_policy(mut self, policy: RetryPolicy) -> Self {
351 self.policy = policy;
352 self
353 }
354}
355
356impl<I> RetryProvider<I> {
357 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 (
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 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 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_eq!(inner.requests().len(), 2);
554 }
555
556 #[tokio::test]
557 async fn gives_up_after_max_attempts() {
558 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 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 #[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 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 #[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 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 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 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 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 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 #[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}