1use std::collections::HashMap;
36use std::future::Future;
37use std::pin::Pin;
38use std::time::Duration;
39
40use crate::error::{ClientError, ClientResult};
41use crate::streaming::EventStream;
42use crate::transport::Transport;
43
44#[derive(Debug, Clone)]
89#[non_exhaustive]
90pub struct RetryPolicy {
91 pub max_retries: u32,
93 pub initial_backoff: Duration,
95 pub max_backoff: Duration,
97 pub backoff_multiplier: f64,
99}
100
101impl Default for RetryPolicy {
102 fn default() -> Self {
103 Self {
104 max_retries: 3,
105 initial_backoff: Duration::from_millis(500),
106 max_backoff: Duration::from_secs(30),
107 backoff_multiplier: 2.0,
108 }
109 }
110}
111
112impl RetryPolicy {
113 #[must_use]
115 pub const fn with_max_retries(mut self, max_retries: u32) -> Self {
116 self.max_retries = max_retries;
117 self
118 }
119
120 #[must_use]
122 pub const fn with_initial_backoff(mut self, backoff: Duration) -> Self {
123 self.initial_backoff = backoff;
124 self
125 }
126
127 #[must_use]
129 pub const fn with_max_backoff(mut self, max: Duration) -> Self {
130 self.max_backoff = max;
131 self
132 }
133
134 #[must_use]
136 pub const fn with_backoff_multiplier(mut self, multiplier: f64) -> Self {
137 self.backoff_multiplier = multiplier;
138 self
139 }
140}
141
142impl ClientError {
145 #[must_use]
154 pub const fn is_retryable(&self) -> bool {
155 match self {
156 Self::Http(_)
157 | Self::HttpClient(_)
158 | Self::Timeout(_)
159 | Self::TooManyPendingRequests { .. } => true,
160 Self::UnexpectedStatus { status, .. } => {
161 matches!(status, 429 | 502 | 503 | 504)
162 }
163 Self::Serialization(_)
165 | Self::Protocol(_)
166 | Self::Transport(_)
167 | Self::InvalidEndpoint(_)
168 | Self::AuthRequired { .. }
169 | Self::ProtocolBindingMismatch(_) => false,
170 }
171 }
172}
173
174fn is_idempotent_method(method: &str) -> bool {
186 matches!(
187 method,
188 "GetTask"
189 | "ListTasks"
190 | "CancelTask"
191 | "SubscribeToTask"
192 | "GetExtendedAgentCard"
193 | "GetTaskPushNotificationConfig"
194 | "ListTaskPushNotificationConfigs"
195 | "DeleteTaskPushNotificationConfig"
196 )
197}
198
199const fn safe_to_retry_non_idempotent(error: &ClientError) -> bool {
209 matches!(
210 error,
211 ClientError::UnexpectedStatus {
212 status: 429 | 503,
213 ..
214 }
215 )
216}
217
218fn retry_delay(
221 last_err: Option<&ClientError>,
222 backoff: Duration,
223 max_backoff: Duration,
224) -> Duration {
225 last_err
226 .and_then(ClientError::retry_after)
227 .map_or_else(|| jittered(backoff), |after| after.min(max_backoff))
228}
229
230pub(crate) struct RetryTransport {
235 inner: Box<dyn Transport>,
236 policy: RetryPolicy,
237}
238
239impl RetryTransport {
240 pub(crate) fn new(inner: Box<dyn Transport>, policy: RetryPolicy) -> Self {
242 Self { inner, policy }
243 }
244}
245
246impl Transport for RetryTransport {
247 fn send_request<'a>(
248 &'a self,
249 method: &'a str,
250 params: serde_json::Value,
251 extra_headers: &'a HashMap<String, String>,
252 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
253 Box::pin(async move {
254 let mut last_err: Option<ClientError> = None;
255 let mut backoff = self.policy.initial_backoff;
256 let idempotent = is_idempotent_method(method);
257
258 let serialized = serde_json::to_vec(¶ms).map_err(ClientError::Serialization)?;
261
262 for attempt in 0..=self.policy.max_retries {
263 if attempt > 0 {
264 let delay = retry_delay(last_err.as_ref(), backoff, self.policy.max_backoff);
265 trace_info!(method, attempt, ?delay, "retrying after backoff");
266 tokio::time::sleep(delay).await;
267 backoff = cap_backoff(
268 backoff,
269 self.policy.backoff_multiplier,
270 self.policy.max_backoff,
271 );
272 }
273
274 let attempt_params: serde_json::Value =
275 serde_json::from_slice(&serialized).map_err(ClientError::Serialization)?;
276
277 match self
278 .inner
279 .send_request(method, attempt_params, extra_headers)
280 .await
281 {
282 Ok(result) => return Ok(result),
283 Err(e)
290 if e.is_retryable() && (idempotent || safe_to_retry_non_idempotent(&e)) =>
291 {
292 trace_warn!(method, attempt, error = %e, "transient error, will retry");
293 last_err = Some(e);
294 }
295 Err(e) => return Err(e),
296 }
297 }
298
299 Err(last_err.expect("at least one attempt was made"))
300 })
301 }
302
303 fn send_streaming_request<'a>(
304 &'a self,
305 method: &'a str,
306 params: serde_json::Value,
307 extra_headers: &'a HashMap<String, String>,
308 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
309 Box::pin(async move {
310 let mut last_err: Option<ClientError> = None;
311 let mut backoff = self.policy.initial_backoff;
312 let idempotent = is_idempotent_method(method);
313
314 let serialized = serde_json::to_vec(¶ms).map_err(ClientError::Serialization)?;
317
318 for attempt in 0..=self.policy.max_retries {
319 if attempt > 0 {
320 let delay = retry_delay(last_err.as_ref(), backoff, self.policy.max_backoff);
321 trace_info!(
322 method,
323 attempt,
324 ?delay,
325 "retrying stream connect after backoff"
326 );
327 tokio::time::sleep(delay).await;
328 backoff = cap_backoff(
329 backoff,
330 self.policy.backoff_multiplier,
331 self.policy.max_backoff,
332 );
333 }
334
335 let attempt_params: serde_json::Value =
336 serde_json::from_slice(&serialized).map_err(ClientError::Serialization)?;
337
338 match self
339 .inner
340 .send_streaming_request(method, attempt_params, extra_headers)
341 .await
342 {
343 Ok(stream) => return Ok(stream),
344 Err(e)
348 if e.is_retryable() && (idempotent || safe_to_retry_non_idempotent(&e)) =>
349 {
350 trace_warn!(method, attempt, error = %e, "transient error, will retry");
351 last_err = Some(e);
352 }
353 Err(e) => return Err(e),
354 }
355 }
356
357 Err(last_err.expect("at least one attempt was made"))
358 })
359 }
360}
361
362fn cap_backoff(current: Duration, multiplier: f64, max: Duration) -> Duration {
369 let next_secs = current.as_secs_f64() * multiplier;
370 Duration::try_from_secs_f64(next_secs).map_or(max, |next| {
375 std::cmp::min(next, max)
380 })
381}
382
383#[allow(clippy::cast_precision_loss)] fn jitter_factor_from_bits(random_bits: u64) -> f64 {
390 (random_bits as f64 / u64::MAX as f64).mul_add(0.5, 0.5)
391}
392
393fn apply_jitter(backoff: Duration, factor: f64) -> Duration {
398 let jittered_secs = backoff.as_secs_f64() * factor;
399 Duration::try_from_secs_f64(jittered_secs).unwrap_or(backoff)
405}
406
407fn jittered(backoff: Duration) -> Duration {
414 use std::hash::{BuildHasher, Hasher};
415 let mut hasher = std::collections::hash_map::RandomState::new().build_hasher();
416 hasher.write_u128(backoff.as_nanos());
418 let factor = jitter_factor_from_bits(hasher.finish());
419 apply_jitter(backoff, factor)
420}
421
422#[cfg(test)]
425mod tests {
426 use super::*;
427
428 #[test]
429 fn http_errors_are_retryable() {
430 let e = ClientError::HttpClient("connection refused".into());
431 assert!(e.is_retryable());
432 }
433
434 #[test]
435 fn timeout_is_retryable() {
436 let e = ClientError::Timeout("request timed out".into());
437 assert!(e.is_retryable());
438 }
439
440 #[test]
441 fn status_503_is_retryable() {
442 let e = ClientError::UnexpectedStatus {
443 status: 503,
444 body: "Service Unavailable".into(),
445 retry_after: None,
446 };
447 assert!(e.is_retryable());
448 }
449
450 #[test]
451 fn status_429_is_retryable() {
452 let e = ClientError::UnexpectedStatus {
453 status: 429,
454 body: "Too Many Requests".into(),
455 retry_after: None,
456 };
457 assert!(e.is_retryable());
458 }
459
460 #[test]
461 fn status_404_is_not_retryable() {
462 let e = ClientError::UnexpectedStatus {
463 status: 404,
464 body: "Not Found".into(),
465 retry_after: None,
466 };
467 assert!(!e.is_retryable());
468 }
469
470 #[test]
471 fn serialization_error_is_not_retryable() {
472 let e = ClientError::Serialization(serde_json::from_str::<String>("not json").unwrap_err());
473 assert!(!e.is_retryable());
474 }
475
476 #[test]
477 fn protocol_error_is_not_retryable() {
478 let e = ClientError::Protocol(a2a_protocol_types::A2aError::task_not_found("t1"));
479 assert!(!e.is_retryable());
480 }
481
482 #[test]
483 fn default_retry_policy() {
484 let p = RetryPolicy::default();
485 assert_eq!(p.max_retries, 3);
486 assert_eq!(p.initial_backoff, Duration::from_millis(500));
487 assert_eq!(p.max_backoff, Duration::from_secs(30));
488 assert!((p.backoff_multiplier - 2.0).abs() < f64::EPSILON);
489 }
490
491 #[test]
492 fn cap_backoff_works() {
493 let result = cap_backoff(Duration::from_secs(1), 2.0, Duration::from_secs(5));
494 assert_eq!(result, Duration::from_secs(2));
495
496 let result = cap_backoff(Duration::from_secs(4), 2.0, Duration::from_secs(5));
497 assert_eq!(result, Duration::from_secs(5));
498 }
499
500 #[test]
501 fn status_502_is_retryable() {
502 let e = ClientError::UnexpectedStatus {
503 status: 502,
504 body: "Bad Gateway".into(),
505 retry_after: None,
506 };
507 assert!(e.is_retryable());
508 }
509
510 #[test]
511 fn status_504_is_retryable() {
512 let e = ClientError::UnexpectedStatus {
513 status: 504,
514 body: "Gateway Timeout".into(),
515 retry_after: None,
516 };
517 assert!(e.is_retryable());
518 }
519
520 #[test]
522 fn status_boundary_not_retryable() {
523 for status in [428, 430, 500, 501, 505] {
524 let e = ClientError::UnexpectedStatus {
525 status,
526 body: String::new(),
527 retry_after: None,
528 };
529 assert!(!e.is_retryable(), "status {status} should not be retryable");
530 }
531 }
532
533 #[test]
534 fn retry_policy_builder_methods() {
535 let p = RetryPolicy::default()
536 .with_max_retries(5)
537 .with_initial_backoff(Duration::from_secs(1))
538 .with_max_backoff(Duration::from_secs(60))
539 .with_backoff_multiplier(3.0);
540 assert_eq!(p.max_retries, 5);
541 assert_eq!(p.initial_backoff, Duration::from_secs(1));
542 assert_eq!(p.max_backoff, Duration::from_secs(60));
543 assert!((p.backoff_multiplier - 3.0).abs() < f64::EPSILON);
544 }
545
546 #[test]
547 fn cap_backoff_exact_boundary() {
548 let result = cap_backoff(Duration::from_secs(5), 1.0, Duration::from_secs(5));
550 assert_eq!(result, Duration::from_secs(5));
551
552 let result = cap_backoff(Duration::from_millis(1), 2.0, Duration::from_secs(5));
554 assert_eq!(result, Duration::from_millis(2));
555 }
556
557 #[test]
558 fn cap_backoff_infinity_returns_max() {
559 let max = Duration::from_secs(30);
561 let result = cap_backoff(Duration::from_secs(u64::MAX / 2), f64::MAX, max);
562 assert_eq!(result, max, "infinity should clamp to max");
563 }
564
565 #[test]
566 fn cap_backoff_finite_overflow_returns_max() {
567 let max = Duration::from_secs(30);
572 let result = cap_backoff(Duration::from_secs(10_000_000_000_000_000_000), 10.0, max);
573 assert_eq!(
574 result, max,
575 "finite-but-overflowing backoff should clamp to max"
576 );
577 }
578
579 #[test]
581 fn jittered_backoff_in_expected_range() {
582 let backoff = Duration::from_secs(2);
583 for _ in 0..100 {
585 let result = jittered(backoff);
586 assert!(
587 result >= Duration::from_secs(1),
588 "jittered backoff should be >= backoff/2, got {result:?}"
589 );
590 assert!(
591 result <= backoff,
592 "jittered backoff should be <= backoff, got {result:?}"
593 );
594 }
595 }
596
597 #[test]
599 fn jittered_zero_backoff() {
600 let result = jittered(Duration::ZERO);
601 assert_eq!(result, Duration::ZERO);
602 }
603
604 #[test]
605 fn cap_backoff_nan_returns_max() {
606 let max = Duration::from_secs(30);
607 let result = cap_backoff(Duration::from_secs(0), f64::NAN, max);
608 assert_eq!(result, max, "NaN should clamp to max");
609 }
610
611 #[test]
616 fn jitter_factor_from_bits_zero() {
617 let f = jitter_factor_from_bits(0);
618 assert!(
619 (f - 0.5).abs() < f64::EPSILON,
620 "factor(0) should be 0.5, got {f}"
621 );
622 }
623
624 #[test]
626 fn jitter_factor_from_bits_midpoint() {
627 let f = jitter_factor_from_bits(u64::MAX / 2);
628 assert!(
630 (0.74..=0.76).contains(&f),
631 "factor(u64::MAX/2) should be ~0.75, got {f}"
632 );
633 }
634
635 #[test]
637 fn jitter_factor_from_bits_max() {
638 let f = jitter_factor_from_bits(u64::MAX);
639 assert!(
642 (0.9..=1.0).contains(&f),
643 "factor(u64::MAX) should be ~1.0, got {f}"
644 );
645 }
646
647 #[test]
651 fn jitter_factor_from_bits_always_in_half_to_one() {
652 for bits in [
653 0_u64,
654 1,
655 7,
656 42,
657 1 << 20,
658 1 << 50,
659 u64::MAX / 4,
660 u64::MAX / 2,
661 u64::MAX,
662 ] {
663 let f = jitter_factor_from_bits(bits);
664 assert!(
665 (0.5..=1.0).contains(&f),
666 "factor({bits}) = {f} out of [0.5, 1.0]"
667 );
668 }
669 }
670
671 #[test]
680 fn apply_jitter_normal_factor() {
681 assert_eq!(
683 apply_jitter(Duration::from_secs(2), 0.5),
684 Duration::from_secs(1)
685 );
686 assert_eq!(
688 apply_jitter(Duration::from_secs(4), 0.75),
689 Duration::from_secs(3)
690 );
691 assert_eq!(
693 apply_jitter(Duration::from_secs(5), 1.0),
694 Duration::from_secs(5)
695 );
696 }
697
698 #[test]
702 fn apply_jitter_zero_factor_returns_zero() {
703 assert_eq!(
704 apply_jitter(Duration::from_secs(5), 0.0),
705 Duration::ZERO,
706 "factor=0.0 must produce Duration::ZERO via from_secs_f64 path"
707 );
708 }
709
710 #[test]
714 fn apply_jitter_negative_factor_returns_backoff() {
715 assert_eq!(
716 apply_jitter(Duration::from_secs(3), -0.5),
717 Duration::from_secs(3),
718 "negative factor must short-circuit to backoff"
719 );
720 }
721
722 #[test]
731 fn apply_jitter_infinite_factor_returns_backoff() {
732 assert_eq!(
733 apply_jitter(Duration::from_secs(2), f64::INFINITY),
734 Duration::from_secs(2),
735 "infinite factor must short-circuit to backoff"
736 );
737 }
738
739 #[test]
740 fn apply_jitter_nan_factor_returns_backoff() {
741 assert_eq!(
742 apply_jitter(Duration::from_secs(4), f64::NAN),
743 Duration::from_secs(4),
744 "NaN factor must short-circuit to backoff"
745 );
746 }
747
748 use std::collections::HashMap;
751 use std::future::Future;
752 use std::pin::Pin;
753 use std::sync::Arc;
754 use std::sync::atomic::{AtomicUsize, Ordering};
755
756 use crate::streaming::EventStream;
757
758 #[tokio::test(start_paused = true)]
775 async fn the_documented_backoff_arithmetic_is_what_it_does() {
776 let inner = FailNTransport::new(usize::MAX, serde_json::json!(null));
777 let calls = Arc::clone(&inner.call_count);
778 let retry = RetryTransport::new(
779 Box::new(inner),
780 RetryPolicy {
781 max_retries: 3,
782 initial_backoff: Duration::from_millis(500),
783 max_backoff: Duration::from_secs(30),
784 backoff_multiplier: 1.0,
787 },
788 );
789
790 let started = tokio::time::Instant::now();
791 let result = retry
795 .send_request("GetTask", serde_json::json!({}), &HashMap::new())
796 .await;
797 let elapsed = started.elapsed();
798
799 assert!(
800 result.is_err(),
801 "every attempt fails, so the call must fail"
802 );
803 assert_eq!(
804 calls.load(Ordering::SeqCst),
805 4,
806 "max_retries 3 must mean 4 attempts — the formula's (max_retries + 1)"
807 );
808
809 assert!(
820 elapsed >= Duration::from_millis(750) && elapsed < Duration::from_millis(1_500),
821 "three fully-jittered 500ms gaps must total in [750ms, 1500ms), got {elapsed:?}"
822 );
823 }
824
825 struct FailNTransport {
827 failures_remaining: Arc<AtomicUsize>,
828 success_response: serde_json::Value,
829 call_count: Arc<AtomicUsize>,
830 }
831
832 impl FailNTransport {
833 fn new(fail_count: usize, response: serde_json::Value) -> Self {
834 Self {
835 failures_remaining: Arc::new(AtomicUsize::new(fail_count)),
836 success_response: response,
837 call_count: Arc::new(AtomicUsize::new(0)),
838 }
839 }
840 }
841
842 impl crate::transport::Transport for FailNTransport {
843 fn send_request<'a>(
844 &'a self,
845 _method: &'a str,
846 _params: serde_json::Value,
847 _extra_headers: &'a HashMap<String, String>,
848 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
849 self.call_count.fetch_add(1, Ordering::SeqCst);
850 let remaining = self.failures_remaining.fetch_sub(1, Ordering::SeqCst);
851 let resp = self.success_response.clone();
852 Box::pin(async move {
853 if remaining > 0 {
854 Err(ClientError::Timeout("transient".into()))
855 } else {
856 Ok(resp)
857 }
858 })
859 }
860
861 fn send_streaming_request<'a>(
862 &'a self,
863 _method: &'a str,
864 _params: serde_json::Value,
865 _extra_headers: &'a HashMap<String, String>,
866 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
867 self.call_count.fetch_add(1, Ordering::SeqCst);
868 let remaining = self.failures_remaining.fetch_sub(1, Ordering::SeqCst);
869 Box::pin(async move {
870 if remaining > 0 {
871 Err(ClientError::Timeout("transient".into()))
872 } else {
873 Err(ClientError::Transport("streaming not mocked".into()))
874 }
875 })
876 }
877 }
878
879 struct NonRetryableErrorTransport {
881 call_count: Arc<AtomicUsize>,
882 }
883
884 impl NonRetryableErrorTransport {
885 fn new() -> Self {
886 Self {
887 call_count: Arc::new(AtomicUsize::new(0)),
888 }
889 }
890 }
891
892 impl crate::transport::Transport for NonRetryableErrorTransport {
893 fn send_request<'a>(
894 &'a self,
895 _method: &'a str,
896 _params: serde_json::Value,
897 _extra_headers: &'a HashMap<String, String>,
898 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
899 self.call_count.fetch_add(1, Ordering::SeqCst);
900 Box::pin(async move { Err(ClientError::InvalidEndpoint("bad url".into())) })
901 }
902
903 fn send_streaming_request<'a>(
904 &'a self,
905 _method: &'a str,
906 _params: serde_json::Value,
907 _extra_headers: &'a HashMap<String, String>,
908 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
909 self.call_count.fetch_add(1, Ordering::SeqCst);
910 Box::pin(async move { Err(ClientError::InvalidEndpoint("bad url".into())) })
911 }
912 }
913
914 #[tokio::test]
915 async fn retry_transport_retries_on_transient_error() {
916 let inner = FailNTransport::new(2, serde_json::json!({"ok": true}));
917 let call_count = Arc::clone(&inner.call_count);
918 let transport = RetryTransport::new(
919 Box::new(inner),
920 RetryPolicy::default()
921 .with_initial_backoff(Duration::from_millis(1))
922 .with_max_retries(3),
923 );
924
925 let headers = HashMap::new();
926 let result = transport
927 .send_request("GetTask", serde_json::Value::Null, &headers)
928 .await;
929 assert!(result.is_ok(), "should succeed after retries");
930 assert_eq!(
931 call_count.load(Ordering::SeqCst),
932 3,
933 "should have made 3 attempts (2 failures + 1 success)"
934 );
935 }
936
937 #[tokio::test]
938 async fn retry_transport_gives_up_after_max_retries() {
939 let inner = FailNTransport::new(10, serde_json::json!({"ok": true}));
941 let call_count = Arc::clone(&inner.call_count);
942 let transport = RetryTransport::new(
943 Box::new(inner),
944 RetryPolicy::default()
945 .with_initial_backoff(Duration::from_millis(1))
946 .with_max_retries(2),
947 );
948
949 let headers = HashMap::new();
950 let result = transport
951 .send_request("GetTask", serde_json::Value::Null, &headers)
952 .await;
953 assert!(result.is_err(), "should fail after exhausting retries");
954 assert_eq!(
955 call_count.load(Ordering::SeqCst),
956 3,
957 "should have made 3 attempts (initial + 2 retries)"
958 );
959 }
960
961 #[tokio::test]
962 async fn retry_transport_no_retry_on_non_retryable() {
963 let inner = NonRetryableErrorTransport::new();
964 let call_count = Arc::clone(&inner.call_count);
965 let transport = RetryTransport::new(
966 Box::new(inner),
967 RetryPolicy::default()
968 .with_initial_backoff(Duration::from_millis(1))
969 .with_max_retries(3),
970 );
971
972 let headers = HashMap::new();
973 let result = transport
974 .send_request("GetTask", serde_json::Value::Null, &headers)
975 .await;
976 assert!(result.is_err());
977 assert!(matches!(
978 result.unwrap_err(),
979 ClientError::InvalidEndpoint(_)
980 ));
981 assert_eq!(
982 call_count.load(Ordering::SeqCst),
983 1,
984 "non-retryable error should not be retried"
985 );
986 }
987
988 #[tokio::test]
989 async fn retry_transport_streaming_retries() {
990 let inner = FailNTransport::new(1, serde_json::json!(null));
991 let call_count = Arc::clone(&inner.call_count);
992 let transport = RetryTransport::new(
993 Box::new(inner),
994 RetryPolicy::default()
995 .with_initial_backoff(Duration::from_millis(1))
996 .with_max_retries(2),
997 );
998
999 let headers = HashMap::new();
1000 let result = transport
1001 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1002 .await;
1003 assert!(result.is_err());
1006 assert_eq!(
1007 call_count.load(Ordering::SeqCst),
1008 2,
1009 "should have retried once for streaming"
1010 );
1011 }
1012
1013 #[tokio::test]
1014 async fn retry_transport_streaming_no_retry_on_non_retryable() {
1015 let inner = NonRetryableErrorTransport::new();
1016 let call_count = Arc::clone(&inner.call_count);
1017 let transport = RetryTransport::new(
1018 Box::new(inner),
1019 RetryPolicy::default()
1020 .with_initial_backoff(Duration::from_millis(1))
1021 .with_max_retries(3),
1022 );
1023
1024 let headers = HashMap::new();
1025 let result = transport
1026 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1027 .await;
1028 assert!(matches!(
1029 result.unwrap_err(),
1030 ClientError::InvalidEndpoint(_)
1031 ));
1032 assert_eq!(
1033 call_count.load(Ordering::SeqCst),
1034 1,
1035 "non-retryable streaming error should not be retried"
1036 );
1037 }
1038
1039 #[tokio::test]
1042 async fn retry_transport_streaming_succeeds_after_retry() {
1043 use tokio::sync::mpsc;
1044
1045 struct FailThenStreamTransport {
1047 call_count: Arc<AtomicUsize>,
1048 }
1049
1050 impl crate::transport::Transport for FailThenStreamTransport {
1051 fn send_request<'a>(
1052 &'a self,
1053 _method: &'a str,
1054 _params: serde_json::Value,
1055 _extra_headers: &'a HashMap<String, String>,
1056 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1057 {
1058 Box::pin(async move { Ok(serde_json::Value::Null) })
1059 }
1060
1061 fn send_streaming_request<'a>(
1062 &'a self,
1063 _method: &'a str,
1064 _params: serde_json::Value,
1065 _extra_headers: &'a HashMap<String, String>,
1066 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1067 let attempt = self.call_count.fetch_add(1, Ordering::SeqCst);
1068 Box::pin(async move {
1069 if attempt == 0 {
1070 Err(ClientError::Timeout("transient timeout".into()))
1071 } else {
1072 let (tx, rx) = mpsc::channel(8);
1074 drop(tx); Ok(EventStream::new(rx))
1076 }
1077 })
1078 }
1079 }
1080
1081 let call_count = Arc::new(AtomicUsize::new(0));
1082 let inner = FailThenStreamTransport {
1083 call_count: Arc::clone(&call_count),
1084 };
1085 let transport = RetryTransport::new(
1086 Box::new(inner),
1087 RetryPolicy::default()
1088 .with_initial_backoff(Duration::from_millis(1))
1089 .with_max_retries(2),
1090 );
1091
1092 let headers = HashMap::new();
1093 let result = transport
1094 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1095 .await;
1096 assert!(result.is_ok(), "streaming should succeed after retry");
1097 assert_eq!(
1098 call_count.load(Ordering::SeqCst),
1099 2,
1100 "should have made 2 attempts (1 failure + 1 success)"
1101 );
1102 }
1103
1104 #[tokio::test]
1105 async fn retry_transport_streaming_exhausts_retries() {
1106 let inner = FailNTransport::new(10, serde_json::json!(null));
1107 let call_count = Arc::clone(&inner.call_count);
1108 let transport = RetryTransport::new(
1109 Box::new(inner),
1110 RetryPolicy::default()
1111 .with_initial_backoff(Duration::from_millis(1))
1112 .with_max_retries(2),
1113 );
1114
1115 let headers = HashMap::new();
1116 let result = transport
1117 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1118 .await;
1119 assert!(result.is_err());
1120 assert_eq!(
1121 call_count.load(Ordering::SeqCst),
1122 3,
1123 "should make 3 attempts total for streaming"
1124 );
1125 }
1126
1127 #[tokio::test]
1128 async fn retry_transport_succeeds_without_retry_on_first_attempt() {
1129 let inner = FailNTransport::new(0, serde_json::json!({"ok": true}));
1130 let call_count = Arc::clone(&inner.call_count);
1131 let transport = RetryTransport::new(
1132 Box::new(inner),
1133 RetryPolicy::default()
1134 .with_initial_backoff(Duration::from_millis(1))
1135 .with_max_retries(3),
1136 );
1137
1138 let headers = HashMap::new();
1139 let result = transport
1140 .send_request("GetTask", serde_json::Value::Null, &headers)
1141 .await;
1142 assert!(result.is_ok());
1143 assert_eq!(
1144 call_count.load(Ordering::SeqCst),
1145 1,
1146 "should succeed on first try"
1147 );
1148 }
1149
1150 #[tokio::test(start_paused = true)]
1156 async fn no_backoff_before_first_attempt() {
1157 let inner = FailNTransport::new(0, serde_json::json!({"ok": true}));
1158 let transport = RetryTransport::new(
1159 Box::new(inner),
1160 RetryPolicy::default()
1161 .with_initial_backoff(Duration::from_secs(100))
1162 .with_max_retries(1),
1163 );
1164
1165 let start = tokio::time::Instant::now();
1166 let headers = HashMap::new();
1167 let result = transport
1168 .send_request("GetTask", serde_json::Value::Null, &headers)
1169 .await;
1170 assert!(result.is_ok());
1171 assert!(
1172 start.elapsed() < Duration::from_secs(1),
1173 "first attempt must not sleep, elapsed: {:?}",
1174 start.elapsed()
1175 );
1176 }
1177
1178 #[tokio::test(start_paused = true)]
1182 async fn backoff_applied_on_retry() {
1183 let inner = FailNTransport::new(1, serde_json::json!({"ok": true}));
1184 let transport = RetryTransport::new(
1185 Box::new(inner),
1186 RetryPolicy::default()
1187 .with_initial_backoff(Duration::from_secs(100))
1188 .with_max_retries(2),
1189 );
1190
1191 let start = tokio::time::Instant::now();
1192 let headers = HashMap::new();
1193 let result = transport
1194 .send_request("GetTask", serde_json::Value::Null, &headers)
1195 .await;
1196 assert!(result.is_ok());
1197 assert!(
1198 start.elapsed() >= Duration::from_secs(50),
1199 "retry should sleep (jittered backoff), elapsed: {:?}",
1200 start.elapsed()
1201 );
1202 }
1203
1204 #[tokio::test(start_paused = true)]
1206 async fn no_backoff_before_first_streaming_attempt() {
1207 use tokio::sync::mpsc;
1208
1209 struct ImmediateStreamTransport;
1210 impl crate::transport::Transport for ImmediateStreamTransport {
1211 fn send_request<'a>(
1212 &'a self,
1213 _method: &'a str,
1214 _params: serde_json::Value,
1215 _extra_headers: &'a HashMap<String, String>,
1216 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1217 {
1218 Box::pin(async { Ok(serde_json::Value::Null) })
1219 }
1220 fn send_streaming_request<'a>(
1221 &'a self,
1222 _method: &'a str,
1223 _params: serde_json::Value,
1224 _extra_headers: &'a HashMap<String, String>,
1225 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1226 Box::pin(async {
1227 let (tx, rx) = mpsc::channel(1);
1228 drop(tx);
1229 Ok(EventStream::new(rx))
1230 })
1231 }
1232 }
1233
1234 let transport = RetryTransport::new(
1235 Box::new(ImmediateStreamTransport),
1236 RetryPolicy::default()
1237 .with_initial_backoff(Duration::from_secs(100))
1238 .with_max_retries(1),
1239 );
1240
1241 let start = tokio::time::Instant::now();
1242 let headers = HashMap::new();
1243 let result = transport
1244 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1245 .await;
1246 assert!(result.is_ok());
1247 assert!(
1248 start.elapsed() < Duration::from_secs(1),
1249 "first streaming attempt must not sleep, elapsed: {:?}",
1250 start.elapsed()
1251 );
1252 }
1253
1254 #[tokio::test(start_paused = true)]
1256 async fn backoff_applied_on_streaming_retry() {
1257 let inner = FailNTransport::new(1, serde_json::json!(null));
1258 let transport = RetryTransport::new(
1259 Box::new(inner),
1260 RetryPolicy::default()
1261 .with_initial_backoff(Duration::from_secs(100))
1262 .with_max_retries(2),
1263 );
1264
1265 let start = tokio::time::Instant::now();
1266 let headers = HashMap::new();
1267 let _result = transport
1268 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1269 .await;
1270 assert!(
1273 start.elapsed() >= Duration::from_secs(50),
1274 "streaming retry should sleep, elapsed: {:?}",
1275 start.elapsed()
1276 );
1277 }
1278
1279 #[test]
1284 fn cap_backoff_zero_multiplier_returns_zero() {
1285 let max = Duration::from_secs(30);
1286 let result = cap_backoff(Duration::from_secs(5), 0.0, max);
1287 assert_eq!(
1288 result,
1289 Duration::ZERO,
1290 "0 * any = 0, should not clamp to max"
1291 );
1292 }
1293
1294 #[test]
1297 fn method_idempotency_classification() {
1298 for m in [
1299 "GetTask",
1300 "ListTasks",
1301 "CancelTask",
1302 "SubscribeToTask",
1303 "GetExtendedAgentCard",
1304 "GetTaskPushNotificationConfig",
1305 "ListTaskPushNotificationConfigs",
1306 "DeleteTaskPushNotificationConfig",
1307 ] {
1308 assert!(is_idempotent_method(m), "{m} should be idempotent");
1309 }
1310 for m in [
1311 "SendMessage",
1312 "SendStreamingMessage",
1313 "CreateTaskPushNotificationConfig",
1314 "SomethingBrandNew",
1315 ] {
1316 assert!(!is_idempotent_method(m), "{m} must be non-idempotent");
1317 }
1318 }
1319
1320 #[test]
1321 fn only_429_503_are_safe_for_non_idempotent() {
1322 let mk = |status| ClientError::UnexpectedStatus {
1323 status,
1324 body: String::new(),
1325 retry_after: None,
1326 };
1327 assert!(safe_to_retry_non_idempotent(&mk(429)));
1328 assert!(safe_to_retry_non_idempotent(&mk(503)));
1329 assert!(!safe_to_retry_non_idempotent(&mk(502)));
1331 assert!(!safe_to_retry_non_idempotent(&mk(504)));
1332 assert!(!safe_to_retry_non_idempotent(&ClientError::Timeout(
1333 "t".into()
1334 )));
1335 assert!(!safe_to_retry_non_idempotent(&ClientError::HttpClient(
1336 "reset".into()
1337 )));
1338 }
1339
1340 #[tokio::test]
1343 async fn non_idempotent_not_retried_on_timeout() {
1344 let inner = FailNTransport::new(5, serde_json::json!({"ok": true}));
1345 let call_count = Arc::clone(&inner.call_count);
1346 let transport = RetryTransport::new(
1347 Box::new(inner),
1348 RetryPolicy::default()
1349 .with_initial_backoff(Duration::from_millis(1))
1350 .with_max_retries(3),
1351 );
1352 let headers = HashMap::new();
1353 let result = transport
1354 .send_request("SendMessage", serde_json::Value::Null, &headers)
1355 .await;
1356 assert!(result.is_err());
1357 assert_eq!(
1358 call_count.load(Ordering::SeqCst),
1359 1,
1360 "SendMessage must not be re-sent on an ambiguous timeout"
1361 );
1362 }
1363
1364 #[tokio::test]
1367 async fn non_idempotent_retried_on_503() {
1368 struct Fail503 {
1370 remaining: Arc<AtomicUsize>,
1371 calls: Arc<AtomicUsize>,
1372 }
1373 impl crate::transport::Transport for Fail503 {
1374 fn send_request<'a>(
1375 &'a self,
1376 _m: &'a str,
1377 _p: serde_json::Value,
1378 _h: &'a HashMap<String, String>,
1379 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1380 {
1381 self.calls.fetch_add(1, Ordering::SeqCst);
1382 let left = self.remaining.fetch_sub(1, Ordering::SeqCst);
1383 Box::pin(async move {
1384 if left > 0 {
1385 Err(ClientError::UnexpectedStatus {
1386 status: 503,
1387 body: String::new(),
1388 retry_after: None,
1389 })
1390 } else {
1391 Ok(serde_json::json!({"ok": true}))
1392 }
1393 })
1394 }
1395 fn send_streaming_request<'a>(
1396 &'a self,
1397 _m: &'a str,
1398 _p: serde_json::Value,
1399 _h: &'a HashMap<String, String>,
1400 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1401 Box::pin(async { Err(ClientError::Transport("n/a".into())) })
1402 }
1403 }
1404 let calls = Arc::new(AtomicUsize::new(0));
1405 let inner = Fail503 {
1406 remaining: Arc::new(AtomicUsize::new(1)),
1407 calls: Arc::clone(&calls),
1408 };
1409 let transport = RetryTransport::new(
1410 Box::new(inner),
1411 RetryPolicy::default()
1412 .with_initial_backoff(Duration::from_millis(1))
1413 .with_max_retries(3),
1414 );
1415 let headers = HashMap::new();
1416 let result = transport
1417 .send_request("SendMessage", serde_json::Value::Null, &headers)
1418 .await;
1419 assert!(result.is_ok(), "503 is a safe retry for non-idempotent");
1420 assert_eq!(calls.load(Ordering::SeqCst), 2);
1421 }
1422
1423 #[tokio::test(start_paused = true)]
1425 async fn retry_after_header_is_honored() {
1426 struct RetryAfter429 {
1428 calls: Arc<AtomicUsize>,
1429 }
1430 impl crate::transport::Transport for RetryAfter429 {
1431 fn send_request<'a>(
1432 &'a self,
1433 _m: &'a str,
1434 _p: serde_json::Value,
1435 _h: &'a HashMap<String, String>,
1436 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1437 {
1438 let n = self.calls.fetch_add(1, Ordering::SeqCst);
1439 Box::pin(async move {
1440 if n == 0 {
1441 Err(ClientError::UnexpectedStatus {
1442 status: 429,
1443 body: String::new(),
1444 retry_after: Some(Duration::from_secs(20)),
1445 })
1446 } else {
1447 Ok(serde_json::json!({"ok": true}))
1448 }
1449 })
1450 }
1451 fn send_streaming_request<'a>(
1452 &'a self,
1453 _m: &'a str,
1454 _p: serde_json::Value,
1455 _h: &'a HashMap<String, String>,
1456 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1457 Box::pin(async { Err(ClientError::Transport("n/a".into())) })
1458 }
1459 }
1460 let calls = Arc::new(AtomicUsize::new(0));
1461 let transport = RetryTransport::new(
1462 Box::new(RetryAfter429 {
1463 calls: Arc::clone(&calls),
1464 }),
1465 RetryPolicy::default()
1468 .with_initial_backoff(Duration::from_millis(1))
1469 .with_max_backoff(Duration::from_secs(30)),
1470 );
1471 let start = tokio::time::Instant::now();
1472 let headers = HashMap::new();
1473 let result = transport
1474 .send_request("GetTask", serde_json::Value::Null, &headers)
1475 .await;
1476 assert!(result.is_ok());
1477 assert!(
1478 start.elapsed() >= Duration::from_secs(20),
1479 "should have waited the server-requested 20s, waited {:?}",
1480 start.elapsed()
1481 );
1482 }
1483
1484 #[test]
1485 fn retry_after_clamped_to_max_backoff() {
1486 let err = ClientError::UnexpectedStatus {
1487 status: 503,
1488 body: String::new(),
1489 retry_after: Some(Duration::from_secs(9999)),
1490 };
1491 let delay = retry_delay(Some(&err), Duration::from_secs(1), Duration::from_secs(30));
1492 assert_eq!(
1493 delay,
1494 Duration::from_secs(30),
1495 "retry-after must be clamped"
1496 );
1497 }
1498}