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)]
85pub struct RetryPolicy {
86 pub max_retries: u32,
88 pub initial_backoff: Duration,
90 pub max_backoff: Duration,
92 pub backoff_multiplier: f64,
94}
95
96impl Default for RetryPolicy {
97 fn default() -> Self {
98 Self {
99 max_retries: 3,
100 initial_backoff: Duration::from_millis(500),
101 max_backoff: Duration::from_secs(30),
102 backoff_multiplier: 2.0,
103 }
104 }
105}
106
107impl RetryPolicy {
108 #[must_use]
110 pub const fn with_max_retries(mut self, max_retries: u32) -> Self {
111 self.max_retries = max_retries;
112 self
113 }
114
115 #[must_use]
117 pub const fn with_initial_backoff(mut self, backoff: Duration) -> Self {
118 self.initial_backoff = backoff;
119 self
120 }
121
122 #[must_use]
124 pub const fn with_max_backoff(mut self, max: Duration) -> Self {
125 self.max_backoff = max;
126 self
127 }
128
129 #[must_use]
131 pub const fn with_backoff_multiplier(mut self, multiplier: f64) -> Self {
132 self.backoff_multiplier = multiplier;
133 self
134 }
135}
136
137impl ClientError {
140 #[must_use]
147 pub const fn is_retryable(&self) -> bool {
148 match self {
149 Self::Http(_) | Self::HttpClient(_) | Self::Timeout(_) => true,
150 Self::UnexpectedStatus { status, .. } => {
151 matches!(status, 429 | 502 | 503 | 504)
152 }
153 Self::Serialization(_)
155 | Self::Protocol(_)
156 | Self::Transport(_)
157 | Self::InvalidEndpoint(_)
158 | Self::AuthRequired { .. }
159 | Self::ProtocolBindingMismatch(_) => false,
160 }
161 }
162}
163
164fn is_idempotent_method(method: &str) -> bool {
176 matches!(
177 method,
178 "GetTask"
179 | "ListTasks"
180 | "CancelTask"
181 | "SubscribeToTask"
182 | "GetExtendedAgentCard"
183 | "GetTaskPushNotificationConfig"
184 | "ListTaskPushNotificationConfigs"
185 | "DeleteTaskPushNotificationConfig"
186 )
187}
188
189const fn safe_to_retry_non_idempotent(error: &ClientError) -> bool {
199 matches!(
200 error,
201 ClientError::UnexpectedStatus {
202 status: 429 | 503,
203 ..
204 }
205 )
206}
207
208fn retry_delay(
211 last_err: Option<&ClientError>,
212 backoff: Duration,
213 max_backoff: Duration,
214) -> Duration {
215 last_err
216 .and_then(ClientError::retry_after)
217 .map_or_else(|| jittered(backoff), |after| after.min(max_backoff))
218}
219
220pub(crate) struct RetryTransport {
225 inner: Box<dyn Transport>,
226 policy: RetryPolicy,
227}
228
229impl RetryTransport {
230 pub(crate) fn new(inner: Box<dyn Transport>, policy: RetryPolicy) -> Self {
232 Self { inner, policy }
233 }
234}
235
236impl Transport for RetryTransport {
237 fn send_request<'a>(
238 &'a self,
239 method: &'a str,
240 params: serde_json::Value,
241 extra_headers: &'a HashMap<String, String>,
242 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
243 Box::pin(async move {
244 let mut last_err: Option<ClientError> = None;
245 let mut backoff = self.policy.initial_backoff;
246 let idempotent = is_idempotent_method(method);
247
248 let serialized = serde_json::to_vec(¶ms).map_err(ClientError::Serialization)?;
251
252 for attempt in 0..=self.policy.max_retries {
253 if attempt > 0 {
254 let delay = retry_delay(last_err.as_ref(), backoff, self.policy.max_backoff);
255 trace_info!(method, attempt, ?delay, "retrying after backoff");
256 tokio::time::sleep(delay).await;
257 backoff = cap_backoff(
258 backoff,
259 self.policy.backoff_multiplier,
260 self.policy.max_backoff,
261 );
262 }
263
264 let attempt_params: serde_json::Value =
265 serde_json::from_slice(&serialized).map_err(ClientError::Serialization)?;
266
267 match self
268 .inner
269 .send_request(method, attempt_params, extra_headers)
270 .await
271 {
272 Ok(result) => return Ok(result),
273 Err(e)
280 if e.is_retryable() && (idempotent || safe_to_retry_non_idempotent(&e)) =>
281 {
282 trace_warn!(method, attempt, error = %e, "transient error, will retry");
283 last_err = Some(e);
284 }
285 Err(e) => return Err(e),
286 }
287 }
288
289 Err(last_err.expect("at least one attempt was made"))
290 })
291 }
292
293 fn send_streaming_request<'a>(
294 &'a self,
295 method: &'a str,
296 params: serde_json::Value,
297 extra_headers: &'a HashMap<String, String>,
298 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
299 Box::pin(async move {
300 let mut last_err: Option<ClientError> = None;
301 let mut backoff = self.policy.initial_backoff;
302 let idempotent = is_idempotent_method(method);
303
304 let serialized = serde_json::to_vec(¶ms).map_err(ClientError::Serialization)?;
307
308 for attempt in 0..=self.policy.max_retries {
309 if attempt > 0 {
310 let delay = retry_delay(last_err.as_ref(), backoff, self.policy.max_backoff);
311 trace_info!(
312 method,
313 attempt,
314 ?delay,
315 "retrying stream connect after backoff"
316 );
317 tokio::time::sleep(delay).await;
318 backoff = cap_backoff(
319 backoff,
320 self.policy.backoff_multiplier,
321 self.policy.max_backoff,
322 );
323 }
324
325 let attempt_params: serde_json::Value =
326 serde_json::from_slice(&serialized).map_err(ClientError::Serialization)?;
327
328 match self
329 .inner
330 .send_streaming_request(method, attempt_params, extra_headers)
331 .await
332 {
333 Ok(stream) => return Ok(stream),
334 Err(e)
338 if e.is_retryable() && (idempotent || safe_to_retry_non_idempotent(&e)) =>
339 {
340 trace_warn!(method, attempt, error = %e, "transient error, will retry");
341 last_err = Some(e);
342 }
343 Err(e) => return Err(e),
344 }
345 }
346
347 Err(last_err.expect("at least one attempt was made"))
348 })
349 }
350}
351
352fn cap_backoff(current: Duration, multiplier: f64, max: Duration) -> Duration {
359 let next_secs = current.as_secs_f64() * multiplier;
360 Duration::try_from_secs_f64(next_secs).map_or(max, |next| {
365 std::cmp::min(next, max)
370 })
371}
372
373#[allow(clippy::cast_precision_loss)] fn jitter_factor_from_bits(random_bits: u64) -> f64 {
380 (random_bits as f64 / u64::MAX as f64).mul_add(0.5, 0.5)
381}
382
383fn apply_jitter(backoff: Duration, factor: f64) -> Duration {
388 let jittered_secs = backoff.as_secs_f64() * factor;
389 Duration::try_from_secs_f64(jittered_secs).unwrap_or(backoff)
395}
396
397fn jittered(backoff: Duration) -> Duration {
404 use std::hash::{BuildHasher, Hasher};
405 let mut hasher = std::collections::hash_map::RandomState::new().build_hasher();
406 hasher.write_u128(backoff.as_nanos());
408 let factor = jitter_factor_from_bits(hasher.finish());
409 apply_jitter(backoff, factor)
410}
411
412#[cfg(test)]
415mod tests {
416 use super::*;
417
418 #[test]
419 fn http_errors_are_retryable() {
420 let e = ClientError::HttpClient("connection refused".into());
421 assert!(e.is_retryable());
422 }
423
424 #[test]
425 fn timeout_is_retryable() {
426 let e = ClientError::Timeout("request timed out".into());
427 assert!(e.is_retryable());
428 }
429
430 #[test]
431 fn status_503_is_retryable() {
432 let e = ClientError::UnexpectedStatus {
433 status: 503,
434 body: "Service Unavailable".into(),
435 retry_after: None,
436 };
437 assert!(e.is_retryable());
438 }
439
440 #[test]
441 fn status_429_is_retryable() {
442 let e = ClientError::UnexpectedStatus {
443 status: 429,
444 body: "Too Many Requests".into(),
445 retry_after: None,
446 };
447 assert!(e.is_retryable());
448 }
449
450 #[test]
451 fn status_404_is_not_retryable() {
452 let e = ClientError::UnexpectedStatus {
453 status: 404,
454 body: "Not Found".into(),
455 retry_after: None,
456 };
457 assert!(!e.is_retryable());
458 }
459
460 #[test]
461 fn serialization_error_is_not_retryable() {
462 let e = ClientError::Serialization(serde_json::from_str::<String>("not json").unwrap_err());
463 assert!(!e.is_retryable());
464 }
465
466 #[test]
467 fn protocol_error_is_not_retryable() {
468 let e = ClientError::Protocol(a2a_protocol_types::A2aError::task_not_found("t1"));
469 assert!(!e.is_retryable());
470 }
471
472 #[test]
473 fn default_retry_policy() {
474 let p = RetryPolicy::default();
475 assert_eq!(p.max_retries, 3);
476 assert_eq!(p.initial_backoff, Duration::from_millis(500));
477 assert_eq!(p.max_backoff, Duration::from_secs(30));
478 assert!((p.backoff_multiplier - 2.0).abs() < f64::EPSILON);
479 }
480
481 #[test]
482 fn cap_backoff_works() {
483 let result = cap_backoff(Duration::from_secs(1), 2.0, Duration::from_secs(5));
484 assert_eq!(result, Duration::from_secs(2));
485
486 let result = cap_backoff(Duration::from_secs(4), 2.0, Duration::from_secs(5));
487 assert_eq!(result, Duration::from_secs(5));
488 }
489
490 #[test]
491 fn status_502_is_retryable() {
492 let e = ClientError::UnexpectedStatus {
493 status: 502,
494 body: "Bad Gateway".into(),
495 retry_after: None,
496 };
497 assert!(e.is_retryable());
498 }
499
500 #[test]
501 fn status_504_is_retryable() {
502 let e = ClientError::UnexpectedStatus {
503 status: 504,
504 body: "Gateway Timeout".into(),
505 retry_after: None,
506 };
507 assert!(e.is_retryable());
508 }
509
510 #[test]
512 fn status_boundary_not_retryable() {
513 for status in [428, 430, 500, 501, 505] {
514 let e = ClientError::UnexpectedStatus {
515 status,
516 body: String::new(),
517 retry_after: None,
518 };
519 assert!(!e.is_retryable(), "status {status} should not be retryable");
520 }
521 }
522
523 #[test]
524 fn retry_policy_builder_methods() {
525 let p = RetryPolicy::default()
526 .with_max_retries(5)
527 .with_initial_backoff(Duration::from_secs(1))
528 .with_max_backoff(Duration::from_secs(60))
529 .with_backoff_multiplier(3.0);
530 assert_eq!(p.max_retries, 5);
531 assert_eq!(p.initial_backoff, Duration::from_secs(1));
532 assert_eq!(p.max_backoff, Duration::from_secs(60));
533 assert!((p.backoff_multiplier - 3.0).abs() < f64::EPSILON);
534 }
535
536 #[test]
537 fn cap_backoff_exact_boundary() {
538 let result = cap_backoff(Duration::from_secs(5), 1.0, Duration::from_secs(5));
540 assert_eq!(result, Duration::from_secs(5));
541
542 let result = cap_backoff(Duration::from_millis(1), 2.0, Duration::from_secs(5));
544 assert_eq!(result, Duration::from_millis(2));
545 }
546
547 #[test]
548 fn cap_backoff_infinity_returns_max() {
549 let max = Duration::from_secs(30);
551 let result = cap_backoff(Duration::from_secs(u64::MAX / 2), f64::MAX, max);
552 assert_eq!(result, max, "infinity should clamp to max");
553 }
554
555 #[test]
556 fn cap_backoff_finite_overflow_returns_max() {
557 let max = Duration::from_secs(30);
562 let result = cap_backoff(Duration::from_secs(10_000_000_000_000_000_000), 10.0, max);
563 assert_eq!(
564 result, max,
565 "finite-but-overflowing backoff should clamp to max"
566 );
567 }
568
569 #[test]
571 fn jittered_backoff_in_expected_range() {
572 let backoff = Duration::from_secs(2);
573 for _ in 0..100 {
575 let result = jittered(backoff);
576 assert!(
577 result >= Duration::from_secs(1),
578 "jittered backoff should be >= backoff/2, got {result:?}"
579 );
580 assert!(
581 result <= backoff,
582 "jittered backoff should be <= backoff, got {result:?}"
583 );
584 }
585 }
586
587 #[test]
589 fn jittered_zero_backoff() {
590 let result = jittered(Duration::ZERO);
591 assert_eq!(result, Duration::ZERO);
592 }
593
594 #[test]
595 fn cap_backoff_nan_returns_max() {
596 let max = Duration::from_secs(30);
597 let result = cap_backoff(Duration::from_secs(0), f64::NAN, max);
598 assert_eq!(result, max, "NaN should clamp to max");
599 }
600
601 #[test]
606 fn jitter_factor_from_bits_zero() {
607 let f = jitter_factor_from_bits(0);
608 assert!(
609 (f - 0.5).abs() < f64::EPSILON,
610 "factor(0) should be 0.5, got {f}"
611 );
612 }
613
614 #[test]
616 fn jitter_factor_from_bits_midpoint() {
617 let f = jitter_factor_from_bits(u64::MAX / 2);
618 assert!(
620 (0.74..=0.76).contains(&f),
621 "factor(u64::MAX/2) should be ~0.75, got {f}"
622 );
623 }
624
625 #[test]
627 fn jitter_factor_from_bits_max() {
628 let f = jitter_factor_from_bits(u64::MAX);
629 assert!(
632 (0.9..=1.0).contains(&f),
633 "factor(u64::MAX) should be ~1.0, got {f}"
634 );
635 }
636
637 #[test]
641 fn jitter_factor_from_bits_always_in_half_to_one() {
642 for bits in [
643 0_u64,
644 1,
645 7,
646 42,
647 1 << 20,
648 1 << 50,
649 u64::MAX / 4,
650 u64::MAX / 2,
651 u64::MAX,
652 ] {
653 let f = jitter_factor_from_bits(bits);
654 assert!(
655 (0.5..=1.0).contains(&f),
656 "factor({bits}) = {f} out of [0.5, 1.0]"
657 );
658 }
659 }
660
661 #[test]
670 fn apply_jitter_normal_factor() {
671 assert_eq!(
673 apply_jitter(Duration::from_secs(2), 0.5),
674 Duration::from_secs(1)
675 );
676 assert_eq!(
678 apply_jitter(Duration::from_secs(4), 0.75),
679 Duration::from_secs(3)
680 );
681 assert_eq!(
683 apply_jitter(Duration::from_secs(5), 1.0),
684 Duration::from_secs(5)
685 );
686 }
687
688 #[test]
692 fn apply_jitter_zero_factor_returns_zero() {
693 assert_eq!(
694 apply_jitter(Duration::from_secs(5), 0.0),
695 Duration::ZERO,
696 "factor=0.0 must produce Duration::ZERO via from_secs_f64 path"
697 );
698 }
699
700 #[test]
704 fn apply_jitter_negative_factor_returns_backoff() {
705 assert_eq!(
706 apply_jitter(Duration::from_secs(3), -0.5),
707 Duration::from_secs(3),
708 "negative factor must short-circuit to backoff"
709 );
710 }
711
712 #[test]
721 fn apply_jitter_infinite_factor_returns_backoff() {
722 assert_eq!(
723 apply_jitter(Duration::from_secs(2), f64::INFINITY),
724 Duration::from_secs(2),
725 "infinite factor must short-circuit to backoff"
726 );
727 }
728
729 #[test]
730 fn apply_jitter_nan_factor_returns_backoff() {
731 assert_eq!(
732 apply_jitter(Duration::from_secs(4), f64::NAN),
733 Duration::from_secs(4),
734 "NaN factor must short-circuit to backoff"
735 );
736 }
737
738 use std::collections::HashMap;
741 use std::future::Future;
742 use std::pin::Pin;
743 use std::sync::atomic::{AtomicUsize, Ordering};
744 use std::sync::Arc;
745
746 use crate::streaming::EventStream;
747
748 #[tokio::test(start_paused = true)]
765 async fn the_documented_backoff_arithmetic_is_what_it_does() {
766 let inner = FailNTransport::new(usize::MAX, serde_json::json!(null));
767 let calls = Arc::clone(&inner.call_count);
768 let retry = RetryTransport::new(
769 Box::new(inner),
770 RetryPolicy {
771 max_retries: 3,
772 initial_backoff: Duration::from_millis(500),
773 max_backoff: Duration::from_secs(30),
774 backoff_multiplier: 1.0,
777 },
778 );
779
780 let started = tokio::time::Instant::now();
781 let result = retry
785 .send_request("GetTask", serde_json::json!({}), &HashMap::new())
786 .await;
787 let elapsed = started.elapsed();
788
789 assert!(
790 result.is_err(),
791 "every attempt fails, so the call must fail"
792 );
793 assert_eq!(
794 calls.load(Ordering::SeqCst),
795 4,
796 "max_retries 3 must mean 4 attempts — the formula's (max_retries + 1)"
797 );
798
799 assert!(
810 elapsed >= Duration::from_millis(750) && elapsed < Duration::from_millis(1_500),
811 "three fully-jittered 500ms gaps must total in [750ms, 1500ms), got {elapsed:?}"
812 );
813 }
814
815 struct FailNTransport {
817 failures_remaining: Arc<AtomicUsize>,
818 success_response: serde_json::Value,
819 call_count: Arc<AtomicUsize>,
820 }
821
822 impl FailNTransport {
823 fn new(fail_count: usize, response: serde_json::Value) -> Self {
824 Self {
825 failures_remaining: Arc::new(AtomicUsize::new(fail_count)),
826 success_response: response,
827 call_count: Arc::new(AtomicUsize::new(0)),
828 }
829 }
830 }
831
832 impl crate::transport::Transport for FailNTransport {
833 fn send_request<'a>(
834 &'a self,
835 _method: &'a str,
836 _params: serde_json::Value,
837 _extra_headers: &'a HashMap<String, String>,
838 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
839 self.call_count.fetch_add(1, Ordering::SeqCst);
840 let remaining = self.failures_remaining.fetch_sub(1, Ordering::SeqCst);
841 let resp = self.success_response.clone();
842 Box::pin(async move {
843 if remaining > 0 {
844 Err(ClientError::Timeout("transient".into()))
845 } else {
846 Ok(resp)
847 }
848 })
849 }
850
851 fn send_streaming_request<'a>(
852 &'a self,
853 _method: &'a str,
854 _params: serde_json::Value,
855 _extra_headers: &'a HashMap<String, String>,
856 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
857 self.call_count.fetch_add(1, Ordering::SeqCst);
858 let remaining = self.failures_remaining.fetch_sub(1, Ordering::SeqCst);
859 Box::pin(async move {
860 if remaining > 0 {
861 Err(ClientError::Timeout("transient".into()))
862 } else {
863 Err(ClientError::Transport("streaming not mocked".into()))
864 }
865 })
866 }
867 }
868
869 struct NonRetryableErrorTransport {
871 call_count: Arc<AtomicUsize>,
872 }
873
874 impl NonRetryableErrorTransport {
875 fn new() -> Self {
876 Self {
877 call_count: Arc::new(AtomicUsize::new(0)),
878 }
879 }
880 }
881
882 impl crate::transport::Transport for NonRetryableErrorTransport {
883 fn send_request<'a>(
884 &'a self,
885 _method: &'a str,
886 _params: serde_json::Value,
887 _extra_headers: &'a HashMap<String, String>,
888 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
889 self.call_count.fetch_add(1, Ordering::SeqCst);
890 Box::pin(async move { Err(ClientError::InvalidEndpoint("bad url".into())) })
891 }
892
893 fn send_streaming_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<EventStream>> + 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
904 #[tokio::test]
905 async fn retry_transport_retries_on_transient_error() {
906 let inner = FailNTransport::new(2, serde_json::json!({"ok": true}));
907 let call_count = Arc::clone(&inner.call_count);
908 let transport = RetryTransport::new(
909 Box::new(inner),
910 RetryPolicy::default()
911 .with_initial_backoff(Duration::from_millis(1))
912 .with_max_retries(3),
913 );
914
915 let headers = HashMap::new();
916 let result = transport
917 .send_request("GetTask", serde_json::Value::Null, &headers)
918 .await;
919 assert!(result.is_ok(), "should succeed after retries");
920 assert_eq!(
921 call_count.load(Ordering::SeqCst),
922 3,
923 "should have made 3 attempts (2 failures + 1 success)"
924 );
925 }
926
927 #[tokio::test]
928 async fn retry_transport_gives_up_after_max_retries() {
929 let inner = FailNTransport::new(10, serde_json::json!({"ok": true}));
931 let call_count = Arc::clone(&inner.call_count);
932 let transport = RetryTransport::new(
933 Box::new(inner),
934 RetryPolicy::default()
935 .with_initial_backoff(Duration::from_millis(1))
936 .with_max_retries(2),
937 );
938
939 let headers = HashMap::new();
940 let result = transport
941 .send_request("GetTask", serde_json::Value::Null, &headers)
942 .await;
943 assert!(result.is_err(), "should fail after exhausting retries");
944 assert_eq!(
945 call_count.load(Ordering::SeqCst),
946 3,
947 "should have made 3 attempts (initial + 2 retries)"
948 );
949 }
950
951 #[tokio::test]
952 async fn retry_transport_no_retry_on_non_retryable() {
953 let inner = NonRetryableErrorTransport::new();
954 let call_count = Arc::clone(&inner.call_count);
955 let transport = RetryTransport::new(
956 Box::new(inner),
957 RetryPolicy::default()
958 .with_initial_backoff(Duration::from_millis(1))
959 .with_max_retries(3),
960 );
961
962 let headers = HashMap::new();
963 let result = transport
964 .send_request("GetTask", serde_json::Value::Null, &headers)
965 .await;
966 assert!(result.is_err());
967 assert!(matches!(
968 result.unwrap_err(),
969 ClientError::InvalidEndpoint(_)
970 ));
971 assert_eq!(
972 call_count.load(Ordering::SeqCst),
973 1,
974 "non-retryable error should not be retried"
975 );
976 }
977
978 #[tokio::test]
979 async fn retry_transport_streaming_retries() {
980 let inner = FailNTransport::new(1, serde_json::json!(null));
981 let call_count = Arc::clone(&inner.call_count);
982 let transport = RetryTransport::new(
983 Box::new(inner),
984 RetryPolicy::default()
985 .with_initial_backoff(Duration::from_millis(1))
986 .with_max_retries(2),
987 );
988
989 let headers = HashMap::new();
990 let result = transport
991 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
992 .await;
993 assert!(result.is_err());
996 assert_eq!(
997 call_count.load(Ordering::SeqCst),
998 2,
999 "should have retried once for streaming"
1000 );
1001 }
1002
1003 #[tokio::test]
1004 async fn retry_transport_streaming_no_retry_on_non_retryable() {
1005 let inner = NonRetryableErrorTransport::new();
1006 let call_count = Arc::clone(&inner.call_count);
1007 let transport = RetryTransport::new(
1008 Box::new(inner),
1009 RetryPolicy::default()
1010 .with_initial_backoff(Duration::from_millis(1))
1011 .with_max_retries(3),
1012 );
1013
1014 let headers = HashMap::new();
1015 let result = transport
1016 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1017 .await;
1018 assert!(matches!(
1019 result.unwrap_err(),
1020 ClientError::InvalidEndpoint(_)
1021 ));
1022 assert_eq!(
1023 call_count.load(Ordering::SeqCst),
1024 1,
1025 "non-retryable streaming error should not be retried"
1026 );
1027 }
1028
1029 #[tokio::test]
1032 async fn retry_transport_streaming_succeeds_after_retry() {
1033 use tokio::sync::mpsc;
1034
1035 struct FailThenStreamTransport {
1037 call_count: Arc<AtomicUsize>,
1038 }
1039
1040 impl crate::transport::Transport for FailThenStreamTransport {
1041 fn send_request<'a>(
1042 &'a self,
1043 _method: &'a str,
1044 _params: serde_json::Value,
1045 _extra_headers: &'a HashMap<String, String>,
1046 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1047 {
1048 Box::pin(async move { Ok(serde_json::Value::Null) })
1049 }
1050
1051 fn send_streaming_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<EventStream>> + Send + 'a>> {
1057 let attempt = self.call_count.fetch_add(1, Ordering::SeqCst);
1058 Box::pin(async move {
1059 if attempt == 0 {
1060 Err(ClientError::Timeout("transient timeout".into()))
1061 } else {
1062 let (tx, rx) = mpsc::channel(8);
1064 drop(tx); Ok(EventStream::new(rx))
1066 }
1067 })
1068 }
1069 }
1070
1071 let call_count = Arc::new(AtomicUsize::new(0));
1072 let inner = FailThenStreamTransport {
1073 call_count: Arc::clone(&call_count),
1074 };
1075 let transport = RetryTransport::new(
1076 Box::new(inner),
1077 RetryPolicy::default()
1078 .with_initial_backoff(Duration::from_millis(1))
1079 .with_max_retries(2),
1080 );
1081
1082 let headers = HashMap::new();
1083 let result = transport
1084 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1085 .await;
1086 assert!(result.is_ok(), "streaming should succeed after retry");
1087 assert_eq!(
1088 call_count.load(Ordering::SeqCst),
1089 2,
1090 "should have made 2 attempts (1 failure + 1 success)"
1091 );
1092 }
1093
1094 #[tokio::test]
1095 async fn retry_transport_streaming_exhausts_retries() {
1096 let inner = FailNTransport::new(10, serde_json::json!(null));
1097 let call_count = Arc::clone(&inner.call_count);
1098 let transport = RetryTransport::new(
1099 Box::new(inner),
1100 RetryPolicy::default()
1101 .with_initial_backoff(Duration::from_millis(1))
1102 .with_max_retries(2),
1103 );
1104
1105 let headers = HashMap::new();
1106 let result = transport
1107 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1108 .await;
1109 assert!(result.is_err());
1110 assert_eq!(
1111 call_count.load(Ordering::SeqCst),
1112 3,
1113 "should make 3 attempts total for streaming"
1114 );
1115 }
1116
1117 #[tokio::test]
1118 async fn retry_transport_succeeds_without_retry_on_first_attempt() {
1119 let inner = FailNTransport::new(0, serde_json::json!({"ok": true}));
1120 let call_count = Arc::clone(&inner.call_count);
1121 let transport = RetryTransport::new(
1122 Box::new(inner),
1123 RetryPolicy::default()
1124 .with_initial_backoff(Duration::from_millis(1))
1125 .with_max_retries(3),
1126 );
1127
1128 let headers = HashMap::new();
1129 let result = transport
1130 .send_request("GetTask", serde_json::Value::Null, &headers)
1131 .await;
1132 assert!(result.is_ok());
1133 assert_eq!(
1134 call_count.load(Ordering::SeqCst),
1135 1,
1136 "should succeed on first try"
1137 );
1138 }
1139
1140 #[tokio::test(start_paused = true)]
1146 async fn no_backoff_before_first_attempt() {
1147 let inner = FailNTransport::new(0, serde_json::json!({"ok": true}));
1148 let transport = RetryTransport::new(
1149 Box::new(inner),
1150 RetryPolicy::default()
1151 .with_initial_backoff(Duration::from_secs(100))
1152 .with_max_retries(1),
1153 );
1154
1155 let start = tokio::time::Instant::now();
1156 let headers = HashMap::new();
1157 let result = transport
1158 .send_request("GetTask", serde_json::Value::Null, &headers)
1159 .await;
1160 assert!(result.is_ok());
1161 assert!(
1162 start.elapsed() < Duration::from_secs(1),
1163 "first attempt must not sleep, elapsed: {:?}",
1164 start.elapsed()
1165 );
1166 }
1167
1168 #[tokio::test(start_paused = true)]
1172 async fn backoff_applied_on_retry() {
1173 let inner = FailNTransport::new(1, serde_json::json!({"ok": true}));
1174 let transport = RetryTransport::new(
1175 Box::new(inner),
1176 RetryPolicy::default()
1177 .with_initial_backoff(Duration::from_secs(100))
1178 .with_max_retries(2),
1179 );
1180
1181 let start = tokio::time::Instant::now();
1182 let headers = HashMap::new();
1183 let result = transport
1184 .send_request("GetTask", serde_json::Value::Null, &headers)
1185 .await;
1186 assert!(result.is_ok());
1187 assert!(
1188 start.elapsed() >= Duration::from_secs(50),
1189 "retry should sleep (jittered backoff), elapsed: {:?}",
1190 start.elapsed()
1191 );
1192 }
1193
1194 #[tokio::test(start_paused = true)]
1196 async fn no_backoff_before_first_streaming_attempt() {
1197 use tokio::sync::mpsc;
1198
1199 struct ImmediateStreamTransport;
1200 impl crate::transport::Transport for ImmediateStreamTransport {
1201 fn send_request<'a>(
1202 &'a self,
1203 _method: &'a str,
1204 _params: serde_json::Value,
1205 _extra_headers: &'a HashMap<String, String>,
1206 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1207 {
1208 Box::pin(async { Ok(serde_json::Value::Null) })
1209 }
1210 fn send_streaming_request<'a>(
1211 &'a self,
1212 _method: &'a str,
1213 _params: serde_json::Value,
1214 _extra_headers: &'a HashMap<String, String>,
1215 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1216 Box::pin(async {
1217 let (tx, rx) = mpsc::channel(1);
1218 drop(tx);
1219 Ok(EventStream::new(rx))
1220 })
1221 }
1222 }
1223
1224 let transport = RetryTransport::new(
1225 Box::new(ImmediateStreamTransport),
1226 RetryPolicy::default()
1227 .with_initial_backoff(Duration::from_secs(100))
1228 .with_max_retries(1),
1229 );
1230
1231 let start = tokio::time::Instant::now();
1232 let headers = HashMap::new();
1233 let result = transport
1234 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1235 .await;
1236 assert!(result.is_ok());
1237 assert!(
1238 start.elapsed() < Duration::from_secs(1),
1239 "first streaming attempt must not sleep, elapsed: {:?}",
1240 start.elapsed()
1241 );
1242 }
1243
1244 #[tokio::test(start_paused = true)]
1246 async fn backoff_applied_on_streaming_retry() {
1247 let inner = FailNTransport::new(1, serde_json::json!(null));
1248 let transport = RetryTransport::new(
1249 Box::new(inner),
1250 RetryPolicy::default()
1251 .with_initial_backoff(Duration::from_secs(100))
1252 .with_max_retries(2),
1253 );
1254
1255 let start = tokio::time::Instant::now();
1256 let headers = HashMap::new();
1257 let _result = transport
1258 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1259 .await;
1260 assert!(
1263 start.elapsed() >= Duration::from_secs(50),
1264 "streaming retry should sleep, elapsed: {:?}",
1265 start.elapsed()
1266 );
1267 }
1268
1269 #[test]
1274 fn cap_backoff_zero_multiplier_returns_zero() {
1275 let max = Duration::from_secs(30);
1276 let result = cap_backoff(Duration::from_secs(5), 0.0, max);
1277 assert_eq!(
1278 result,
1279 Duration::ZERO,
1280 "0 * any = 0, should not clamp to max"
1281 );
1282 }
1283
1284 #[test]
1287 fn method_idempotency_classification() {
1288 for m in [
1289 "GetTask",
1290 "ListTasks",
1291 "CancelTask",
1292 "SubscribeToTask",
1293 "GetExtendedAgentCard",
1294 "GetTaskPushNotificationConfig",
1295 "ListTaskPushNotificationConfigs",
1296 "DeleteTaskPushNotificationConfig",
1297 ] {
1298 assert!(is_idempotent_method(m), "{m} should be idempotent");
1299 }
1300 for m in [
1301 "SendMessage",
1302 "SendStreamingMessage",
1303 "CreateTaskPushNotificationConfig",
1304 "SomethingBrandNew",
1305 ] {
1306 assert!(!is_idempotent_method(m), "{m} must be non-idempotent");
1307 }
1308 }
1309
1310 #[test]
1311 fn only_429_503_are_safe_for_non_idempotent() {
1312 let mk = |status| ClientError::UnexpectedStatus {
1313 status,
1314 body: String::new(),
1315 retry_after: None,
1316 };
1317 assert!(safe_to_retry_non_idempotent(&mk(429)));
1318 assert!(safe_to_retry_non_idempotent(&mk(503)));
1319 assert!(!safe_to_retry_non_idempotent(&mk(502)));
1321 assert!(!safe_to_retry_non_idempotent(&mk(504)));
1322 assert!(!safe_to_retry_non_idempotent(&ClientError::Timeout(
1323 "t".into()
1324 )));
1325 assert!(!safe_to_retry_non_idempotent(&ClientError::HttpClient(
1326 "reset".into()
1327 )));
1328 }
1329
1330 #[tokio::test]
1333 async fn non_idempotent_not_retried_on_timeout() {
1334 let inner = FailNTransport::new(5, serde_json::json!({"ok": true}));
1335 let call_count = Arc::clone(&inner.call_count);
1336 let transport = RetryTransport::new(
1337 Box::new(inner),
1338 RetryPolicy::default()
1339 .with_initial_backoff(Duration::from_millis(1))
1340 .with_max_retries(3),
1341 );
1342 let headers = HashMap::new();
1343 let result = transport
1344 .send_request("SendMessage", serde_json::Value::Null, &headers)
1345 .await;
1346 assert!(result.is_err());
1347 assert_eq!(
1348 call_count.load(Ordering::SeqCst),
1349 1,
1350 "SendMessage must not be re-sent on an ambiguous timeout"
1351 );
1352 }
1353
1354 #[tokio::test]
1357 async fn non_idempotent_retried_on_503() {
1358 struct Fail503 {
1360 remaining: Arc<AtomicUsize>,
1361 calls: Arc<AtomicUsize>,
1362 }
1363 impl crate::transport::Transport for Fail503 {
1364 fn send_request<'a>(
1365 &'a self,
1366 _m: &'a str,
1367 _p: serde_json::Value,
1368 _h: &'a HashMap<String, String>,
1369 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1370 {
1371 self.calls.fetch_add(1, Ordering::SeqCst);
1372 let left = self.remaining.fetch_sub(1, Ordering::SeqCst);
1373 Box::pin(async move {
1374 if left > 0 {
1375 Err(ClientError::UnexpectedStatus {
1376 status: 503,
1377 body: String::new(),
1378 retry_after: None,
1379 })
1380 } else {
1381 Ok(serde_json::json!({"ok": true}))
1382 }
1383 })
1384 }
1385 fn send_streaming_request<'a>(
1386 &'a self,
1387 _m: &'a str,
1388 _p: serde_json::Value,
1389 _h: &'a HashMap<String, String>,
1390 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1391 Box::pin(async { Err(ClientError::Transport("n/a".into())) })
1392 }
1393 }
1394 let calls = Arc::new(AtomicUsize::new(0));
1395 let inner = Fail503 {
1396 remaining: Arc::new(AtomicUsize::new(1)),
1397 calls: Arc::clone(&calls),
1398 };
1399 let transport = RetryTransport::new(
1400 Box::new(inner),
1401 RetryPolicy::default()
1402 .with_initial_backoff(Duration::from_millis(1))
1403 .with_max_retries(3),
1404 );
1405 let headers = HashMap::new();
1406 let result = transport
1407 .send_request("SendMessage", serde_json::Value::Null, &headers)
1408 .await;
1409 assert!(result.is_ok(), "503 is a safe retry for non-idempotent");
1410 assert_eq!(calls.load(Ordering::SeqCst), 2);
1411 }
1412
1413 #[tokio::test(start_paused = true)]
1415 async fn retry_after_header_is_honored() {
1416 struct RetryAfter429 {
1418 calls: Arc<AtomicUsize>,
1419 }
1420 impl crate::transport::Transport for RetryAfter429 {
1421 fn send_request<'a>(
1422 &'a self,
1423 _m: &'a str,
1424 _p: serde_json::Value,
1425 _h: &'a HashMap<String, String>,
1426 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1427 {
1428 let n = self.calls.fetch_add(1, Ordering::SeqCst);
1429 Box::pin(async move {
1430 if n == 0 {
1431 Err(ClientError::UnexpectedStatus {
1432 status: 429,
1433 body: String::new(),
1434 retry_after: Some(Duration::from_secs(20)),
1435 })
1436 } else {
1437 Ok(serde_json::json!({"ok": true}))
1438 }
1439 })
1440 }
1441 fn send_streaming_request<'a>(
1442 &'a self,
1443 _m: &'a str,
1444 _p: serde_json::Value,
1445 _h: &'a HashMap<String, String>,
1446 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1447 Box::pin(async { Err(ClientError::Transport("n/a".into())) })
1448 }
1449 }
1450 let calls = Arc::new(AtomicUsize::new(0));
1451 let transport = RetryTransport::new(
1452 Box::new(RetryAfter429 {
1453 calls: Arc::clone(&calls),
1454 }),
1455 RetryPolicy::default()
1458 .with_initial_backoff(Duration::from_millis(1))
1459 .with_max_backoff(Duration::from_secs(30)),
1460 );
1461 let start = tokio::time::Instant::now();
1462 let headers = HashMap::new();
1463 let result = transport
1464 .send_request("GetTask", serde_json::Value::Null, &headers)
1465 .await;
1466 assert!(result.is_ok());
1467 assert!(
1468 start.elapsed() >= Duration::from_secs(20),
1469 "should have waited the server-requested 20s, waited {:?}",
1470 start.elapsed()
1471 );
1472 }
1473
1474 #[test]
1475 fn retry_after_clamped_to_max_backoff() {
1476 let err = ClientError::UnexpectedStatus {
1477 status: 503,
1478 body: String::new(),
1479 retry_after: Some(Duration::from_secs(9999)),
1480 };
1481 let delay = retry_delay(Some(&err), Duration::from_secs(1), Duration::from_secs(30));
1482 assert_eq!(
1483 delay,
1484 Duration::from_secs(30),
1485 "retry-after must be clamped"
1486 );
1487 }
1488}