1use std::collections::HashMap;
36use std::future::Future;
37use std::pin::Pin;
38use std::time::Duration;
39
40use a2a_protocol_types::idempotency::{IDEMPOTENCY_METADATA_KEY, validate_key};
41
42use crate::error::{ClientError, ClientResult};
43use crate::streaming::EventStream;
44use crate::transport::Transport;
45
46#[derive(Debug, Clone)]
91#[non_exhaustive]
92pub struct RetryPolicy {
93 pub max_retries: u32,
95 pub initial_backoff: Duration,
97 pub max_backoff: Duration,
99 pub backoff_multiplier: f64,
101}
102
103impl Default for RetryPolicy {
104 fn default() -> Self {
105 Self {
106 max_retries: 3,
107 initial_backoff: Duration::from_millis(500),
108 max_backoff: Duration::from_secs(30),
109 backoff_multiplier: 2.0,
110 }
111 }
112}
113
114impl RetryPolicy {
115 #[must_use]
117 pub const fn with_max_retries(mut self, max_retries: u32) -> Self {
118 self.max_retries = max_retries;
119 self
120 }
121
122 #[must_use]
124 pub const fn with_initial_backoff(mut self, backoff: Duration) -> Self {
125 self.initial_backoff = backoff;
126 self
127 }
128
129 #[must_use]
131 pub const fn with_max_backoff(mut self, max: Duration) -> Self {
132 self.max_backoff = max;
133 self
134 }
135
136 #[must_use]
138 pub const fn with_backoff_multiplier(mut self, multiplier: f64) -> Self {
139 self.backoff_multiplier = multiplier;
140 self
141 }
142}
143
144impl ClientError {
147 #[must_use]
156 pub const fn is_retryable(&self) -> bool {
157 match self {
158 Self::Http(_)
159 | Self::HttpClient(_)
160 | Self::Timeout(_)
161 | Self::TooManyPendingRequests { .. } => true,
162 Self::UnexpectedStatus { status, .. } => {
163 matches!(status, 429 | 502 | 503 | 504)
164 }
165 Self::Serialization(_)
167 | Self::Protocol(_)
168 | Self::Transport(_)
169 | Self::InvalidEndpoint(_)
170 | Self::AuthRequired { .. }
171 | Self::ProtocolBindingMismatch(_) => false,
172 }
173 }
174}
175
176fn is_idempotent_method(method: &str) -> bool {
192 matches!(
193 method,
194 "GetTask"
195 | "ListTasks"
196 | "CancelTask"
197 | "SubscribeToTask"
198 | "GetExtendedAgentCard"
199 | "GetTaskPushNotificationConfig"
200 | "ListTaskPushNotificationConfigs"
201 | "DeleteTaskPushNotificationConfig"
202 )
203}
204
205fn carries_idempotency_key(
238 method: &str,
239 params: &serde_json::Value,
240 peer_honours_idempotency: bool,
241) -> bool {
242 if !peer_honours_idempotency {
243 return false;
244 }
245 if !matches!(method, "SendMessage" | "SendStreamingMessage") {
246 return false;
247 }
248 params
249 .get("message")
250 .and_then(|message| message.get("metadata"))
251 .and_then(|metadata| metadata.get(IDEMPOTENCY_METADATA_KEY))
252 .and_then(serde_json::Value::as_str)
253 .is_some_and(|key| validate_key(key).is_ok())
254}
255
256const fn safe_to_retry_non_idempotent(error: &ClientError) -> bool {
266 matches!(
267 error,
268 ClientError::UnexpectedStatus {
269 status: 429 | 503,
270 ..
271 }
272 )
273}
274
275fn retry_delay(
278 last_err: Option<&ClientError>,
279 backoff: Duration,
280 max_backoff: Duration,
281) -> Duration {
282 last_err
283 .and_then(ClientError::retry_after)
284 .map_or_else(|| jittered(backoff), |after| after.min(max_backoff))
285}
286
287pub(crate) struct RetryTransport {
292 inner: Box<dyn Transport>,
293 policy: RetryPolicy,
294 peer_honours_idempotency: bool,
300}
301
302impl RetryTransport {
303 pub(crate) fn new(inner: Box<dyn Transport>, policy: RetryPolicy) -> Self {
305 Self {
306 inner,
307 policy,
308 peer_honours_idempotency: false,
309 }
310 }
311
312 #[must_use]
314 pub(crate) const fn with_peer_honouring_idempotency(mut self, honours: bool) -> Self {
315 self.peer_honours_idempotency = honours;
316 self
317 }
318}
319
320impl Transport for RetryTransport {
321 fn send_request<'a>(
322 &'a self,
323 method: &'a str,
324 params: serde_json::Value,
325 extra_headers: &'a HashMap<String, String>,
326 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
327 Box::pin(async move {
328 let mut last_err: Option<ClientError> = None;
329 let mut backoff = self.policy.initial_backoff;
330 let idempotent = is_idempotent_method(method)
331 || carries_idempotency_key(method, ¶ms, self.peer_honours_idempotency);
332
333 let serialized = serde_json::to_vec(¶ms).map_err(ClientError::Serialization)?;
336
337 for attempt in 0..=self.policy.max_retries {
338 if attempt > 0 {
339 let delay = retry_delay(last_err.as_ref(), backoff, self.policy.max_backoff);
340 trace_info!(method, attempt, ?delay, "retrying after backoff");
341 tokio::time::sleep(delay).await;
342 backoff = cap_backoff(
343 backoff,
344 self.policy.backoff_multiplier,
345 self.policy.max_backoff,
346 );
347 }
348
349 let attempt_params: serde_json::Value =
350 serde_json::from_slice(&serialized).map_err(ClientError::Serialization)?;
351
352 match self
353 .inner
354 .send_request(method, attempt_params, extra_headers)
355 .await
356 {
357 Ok(result) => return Ok(result),
358 Err(e)
365 if e.is_retryable() && (idempotent || safe_to_retry_non_idempotent(&e)) =>
366 {
367 trace_warn!(method, attempt, error = %e, "transient error, will retry");
368 last_err = Some(e);
369 }
370 Err(e) => return Err(e),
371 }
372 }
373
374 Err(last_err.expect("at least one attempt was made"))
375 })
376 }
377
378 fn send_streaming_request<'a>(
379 &'a self,
380 method: &'a str,
381 params: serde_json::Value,
382 extra_headers: &'a HashMap<String, String>,
383 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
384 Box::pin(async move {
385 let mut last_err: Option<ClientError> = None;
386 let mut backoff = self.policy.initial_backoff;
387 let idempotent = is_idempotent_method(method)
388 || carries_idempotency_key(method, ¶ms, self.peer_honours_idempotency);
389
390 let serialized = serde_json::to_vec(¶ms).map_err(ClientError::Serialization)?;
393
394 for attempt in 0..=self.policy.max_retries {
395 if attempt > 0 {
396 let delay = retry_delay(last_err.as_ref(), backoff, self.policy.max_backoff);
397 trace_info!(
398 method,
399 attempt,
400 ?delay,
401 "retrying stream connect after backoff"
402 );
403 tokio::time::sleep(delay).await;
404 backoff = cap_backoff(
405 backoff,
406 self.policy.backoff_multiplier,
407 self.policy.max_backoff,
408 );
409 }
410
411 let attempt_params: serde_json::Value =
412 serde_json::from_slice(&serialized).map_err(ClientError::Serialization)?;
413
414 match self
415 .inner
416 .send_streaming_request(method, attempt_params, extra_headers)
417 .await
418 {
419 Ok(stream) => return Ok(stream),
420 Err(e)
424 if e.is_retryable() && (idempotent || safe_to_retry_non_idempotent(&e)) =>
425 {
426 trace_warn!(method, attempt, error = %e, "transient error, will retry");
427 last_err = Some(e);
428 }
429 Err(e) => return Err(e),
430 }
431 }
432
433 Err(last_err.expect("at least one attempt was made"))
434 })
435 }
436}
437
438fn cap_backoff(current: Duration, multiplier: f64, max: Duration) -> Duration {
445 let next_secs = current.as_secs_f64() * multiplier;
446 Duration::try_from_secs_f64(next_secs).map_or(max, |next| {
451 std::cmp::min(next, max)
456 })
457}
458
459#[allow(clippy::cast_precision_loss)] fn jitter_factor_from_bits(random_bits: u64) -> f64 {
466 (random_bits as f64 / u64::MAX as f64).mul_add(0.5, 0.5)
467}
468
469fn apply_jitter(backoff: Duration, factor: f64) -> Duration {
474 let jittered_secs = backoff.as_secs_f64() * factor;
475 Duration::try_from_secs_f64(jittered_secs).unwrap_or(backoff)
481}
482
483fn jittered(backoff: Duration) -> Duration {
490 use std::hash::{BuildHasher, Hasher};
491 let mut hasher = std::collections::hash_map::RandomState::new().build_hasher();
492 hasher.write_u128(backoff.as_nanos());
494 let factor = jitter_factor_from_bits(hasher.finish());
495 apply_jitter(backoff, factor)
496}
497
498#[cfg(test)]
501mod tests {
502 use super::*;
503
504 #[test]
505 fn http_errors_are_retryable() {
506 let e = ClientError::HttpClient("connection refused".into());
507 assert!(e.is_retryable());
508 }
509
510 #[test]
511 fn timeout_is_retryable() {
512 let e = ClientError::Timeout("request timed out".into());
513 assert!(e.is_retryable());
514 }
515
516 #[test]
517 fn status_503_is_retryable() {
518 let e = ClientError::UnexpectedStatus {
519 status: 503,
520 body: "Service Unavailable".into(),
521 retry_after: None,
522 };
523 assert!(e.is_retryable());
524 }
525
526 #[test]
527 fn status_429_is_retryable() {
528 let e = ClientError::UnexpectedStatus {
529 status: 429,
530 body: "Too Many Requests".into(),
531 retry_after: None,
532 };
533 assert!(e.is_retryable());
534 }
535
536 #[test]
537 fn status_404_is_not_retryable() {
538 let e = ClientError::UnexpectedStatus {
539 status: 404,
540 body: "Not Found".into(),
541 retry_after: None,
542 };
543 assert!(!e.is_retryable());
544 }
545
546 #[test]
547 fn serialization_error_is_not_retryable() {
548 let e = ClientError::Serialization(serde_json::from_str::<String>("not json").unwrap_err());
549 assert!(!e.is_retryable());
550 }
551
552 #[test]
553 fn protocol_error_is_not_retryable() {
554 let e = ClientError::Protocol(a2a_protocol_types::A2aError::task_not_found("t1"));
555 assert!(!e.is_retryable());
556 }
557
558 #[test]
559 fn default_retry_policy() {
560 let p = RetryPolicy::default();
561 assert_eq!(p.max_retries, 3);
562 assert_eq!(p.initial_backoff, Duration::from_millis(500));
563 assert_eq!(p.max_backoff, Duration::from_secs(30));
564 assert!((p.backoff_multiplier - 2.0).abs() < f64::EPSILON);
565 }
566
567 #[test]
568 fn cap_backoff_works() {
569 let result = cap_backoff(Duration::from_secs(1), 2.0, Duration::from_secs(5));
570 assert_eq!(result, Duration::from_secs(2));
571
572 let result = cap_backoff(Duration::from_secs(4), 2.0, Duration::from_secs(5));
573 assert_eq!(result, Duration::from_secs(5));
574 }
575
576 #[test]
577 fn status_502_is_retryable() {
578 let e = ClientError::UnexpectedStatus {
579 status: 502,
580 body: "Bad Gateway".into(),
581 retry_after: None,
582 };
583 assert!(e.is_retryable());
584 }
585
586 #[test]
587 fn status_504_is_retryable() {
588 let e = ClientError::UnexpectedStatus {
589 status: 504,
590 body: "Gateway Timeout".into(),
591 retry_after: None,
592 };
593 assert!(e.is_retryable());
594 }
595
596 #[test]
598 fn status_boundary_not_retryable() {
599 for status in [428, 430, 500, 501, 505] {
600 let e = ClientError::UnexpectedStatus {
601 status,
602 body: String::new(),
603 retry_after: None,
604 };
605 assert!(!e.is_retryable(), "status {status} should not be retryable");
606 }
607 }
608
609 #[test]
610 fn retry_policy_builder_methods() {
611 let p = RetryPolicy::default()
612 .with_max_retries(5)
613 .with_initial_backoff(Duration::from_secs(1))
614 .with_max_backoff(Duration::from_secs(60))
615 .with_backoff_multiplier(3.0);
616 assert_eq!(p.max_retries, 5);
617 assert_eq!(p.initial_backoff, Duration::from_secs(1));
618 assert_eq!(p.max_backoff, Duration::from_secs(60));
619 assert!((p.backoff_multiplier - 3.0).abs() < f64::EPSILON);
620 }
621
622 #[test]
623 fn cap_backoff_exact_boundary() {
624 let result = cap_backoff(Duration::from_secs(5), 1.0, Duration::from_secs(5));
626 assert_eq!(result, Duration::from_secs(5));
627
628 let result = cap_backoff(Duration::from_millis(1), 2.0, Duration::from_secs(5));
630 assert_eq!(result, Duration::from_millis(2));
631 }
632
633 #[test]
634 fn cap_backoff_infinity_returns_max() {
635 let max = Duration::from_secs(30);
637 let result = cap_backoff(Duration::from_secs(u64::MAX / 2), f64::MAX, max);
638 assert_eq!(result, max, "infinity should clamp to max");
639 }
640
641 #[test]
642 fn cap_backoff_finite_overflow_returns_max() {
643 let max = Duration::from_secs(30);
648 let result = cap_backoff(Duration::from_secs(10_000_000_000_000_000_000), 10.0, max);
649 assert_eq!(
650 result, max,
651 "finite-but-overflowing backoff should clamp to max"
652 );
653 }
654
655 #[test]
657 fn jittered_backoff_in_expected_range() {
658 let backoff = Duration::from_secs(2);
659 for _ in 0..100 {
661 let result = jittered(backoff);
662 assert!(
663 result >= Duration::from_secs(1),
664 "jittered backoff should be >= backoff/2, got {result:?}"
665 );
666 assert!(
667 result <= backoff,
668 "jittered backoff should be <= backoff, got {result:?}"
669 );
670 }
671 }
672
673 #[test]
675 fn jittered_zero_backoff() {
676 let result = jittered(Duration::ZERO);
677 assert_eq!(result, Duration::ZERO);
678 }
679
680 #[test]
681 fn cap_backoff_nan_returns_max() {
682 let max = Duration::from_secs(30);
683 let result = cap_backoff(Duration::from_secs(0), f64::NAN, max);
684 assert_eq!(result, max, "NaN should clamp to max");
685 }
686
687 #[test]
692 fn jitter_factor_from_bits_zero() {
693 let f = jitter_factor_from_bits(0);
694 assert!(
695 (f - 0.5).abs() < f64::EPSILON,
696 "factor(0) should be 0.5, got {f}"
697 );
698 }
699
700 #[test]
702 fn jitter_factor_from_bits_midpoint() {
703 let f = jitter_factor_from_bits(u64::MAX / 2);
704 assert!(
706 (0.74..=0.76).contains(&f),
707 "factor(u64::MAX/2) should be ~0.75, got {f}"
708 );
709 }
710
711 #[test]
713 fn jitter_factor_from_bits_max() {
714 let f = jitter_factor_from_bits(u64::MAX);
715 assert!(
718 (0.9..=1.0).contains(&f),
719 "factor(u64::MAX) should be ~1.0, got {f}"
720 );
721 }
722
723 #[test]
727 fn jitter_factor_from_bits_always_in_half_to_one() {
728 for bits in [
729 0_u64,
730 1,
731 7,
732 42,
733 1 << 20,
734 1 << 50,
735 u64::MAX / 4,
736 u64::MAX / 2,
737 u64::MAX,
738 ] {
739 let f = jitter_factor_from_bits(bits);
740 assert!(
741 (0.5..=1.0).contains(&f),
742 "factor({bits}) = {f} out of [0.5, 1.0]"
743 );
744 }
745 }
746
747 #[test]
756 fn apply_jitter_normal_factor() {
757 assert_eq!(
759 apply_jitter(Duration::from_secs(2), 0.5),
760 Duration::from_secs(1)
761 );
762 assert_eq!(
764 apply_jitter(Duration::from_secs(4), 0.75),
765 Duration::from_secs(3)
766 );
767 assert_eq!(
769 apply_jitter(Duration::from_secs(5), 1.0),
770 Duration::from_secs(5)
771 );
772 }
773
774 #[test]
778 fn apply_jitter_zero_factor_returns_zero() {
779 assert_eq!(
780 apply_jitter(Duration::from_secs(5), 0.0),
781 Duration::ZERO,
782 "factor=0.0 must produce Duration::ZERO via from_secs_f64 path"
783 );
784 }
785
786 #[test]
790 fn apply_jitter_negative_factor_returns_backoff() {
791 assert_eq!(
792 apply_jitter(Duration::from_secs(3), -0.5),
793 Duration::from_secs(3),
794 "negative factor must short-circuit to backoff"
795 );
796 }
797
798 #[test]
807 fn apply_jitter_infinite_factor_returns_backoff() {
808 assert_eq!(
809 apply_jitter(Duration::from_secs(2), f64::INFINITY),
810 Duration::from_secs(2),
811 "infinite factor must short-circuit to backoff"
812 );
813 }
814
815 #[test]
816 fn apply_jitter_nan_factor_returns_backoff() {
817 assert_eq!(
818 apply_jitter(Duration::from_secs(4), f64::NAN),
819 Duration::from_secs(4),
820 "NaN factor must short-circuit to backoff"
821 );
822 }
823
824 use std::collections::HashMap;
827 use std::future::Future;
828 use std::pin::Pin;
829 use std::sync::Arc;
830 use std::sync::atomic::{AtomicUsize, Ordering};
831
832 use crate::streaming::EventStream;
833
834 #[tokio::test(start_paused = true)]
851 async fn the_documented_backoff_arithmetic_is_what_it_does() {
852 let inner = FailNTransport::new(usize::MAX, serde_json::json!(null));
853 let calls = Arc::clone(&inner.call_count);
854 let retry = RetryTransport::new(
855 Box::new(inner),
856 RetryPolicy {
857 max_retries: 3,
858 initial_backoff: Duration::from_millis(500),
859 max_backoff: Duration::from_secs(30),
860 backoff_multiplier: 1.0,
863 },
864 );
865
866 let started = tokio::time::Instant::now();
867 let result = retry
871 .send_request("GetTask", serde_json::json!({}), &HashMap::new())
872 .await;
873 let elapsed = started.elapsed();
874
875 assert!(
876 result.is_err(),
877 "every attempt fails, so the call must fail"
878 );
879 assert_eq!(
880 calls.load(Ordering::SeqCst),
881 4,
882 "max_retries 3 must mean 4 attempts — the formula's (max_retries + 1)"
883 );
884
885 assert!(
896 elapsed >= Duration::from_millis(750) && elapsed < Duration::from_millis(1_500),
897 "three fully-jittered 500ms gaps must total in [750ms, 1500ms), got {elapsed:?}"
898 );
899 }
900
901 struct FailNTransport {
903 failures_remaining: Arc<AtomicUsize>,
904 success_response: serde_json::Value,
905 call_count: Arc<AtomicUsize>,
906 }
907
908 impl FailNTransport {
909 fn new(fail_count: usize, response: serde_json::Value) -> Self {
910 Self {
911 failures_remaining: Arc::new(AtomicUsize::new(fail_count)),
912 success_response: response,
913 call_count: Arc::new(AtomicUsize::new(0)),
914 }
915 }
916 }
917
918 impl crate::transport::Transport for FailNTransport {
919 fn send_request<'a>(
920 &'a self,
921 _method: &'a str,
922 _params: serde_json::Value,
923 _extra_headers: &'a HashMap<String, String>,
924 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
925 self.call_count.fetch_add(1, Ordering::SeqCst);
926 let remaining = self.failures_remaining.fetch_sub(1, Ordering::SeqCst);
927 let resp = self.success_response.clone();
928 Box::pin(async move {
929 if remaining > 0 {
930 Err(ClientError::Timeout("transient".into()))
931 } else {
932 Ok(resp)
933 }
934 })
935 }
936
937 fn send_streaming_request<'a>(
938 &'a self,
939 _method: &'a str,
940 _params: serde_json::Value,
941 _extra_headers: &'a HashMap<String, String>,
942 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
943 self.call_count.fetch_add(1, Ordering::SeqCst);
944 let remaining = self.failures_remaining.fetch_sub(1, Ordering::SeqCst);
945 Box::pin(async move {
946 if remaining > 0 {
947 Err(ClientError::Timeout("transient".into()))
948 } else {
949 Err(ClientError::Transport("streaming not mocked".into()))
950 }
951 })
952 }
953 }
954
955 struct NonRetryableErrorTransport {
957 call_count: Arc<AtomicUsize>,
958 }
959
960 impl NonRetryableErrorTransport {
961 fn new() -> Self {
962 Self {
963 call_count: Arc::new(AtomicUsize::new(0)),
964 }
965 }
966 }
967
968 impl crate::transport::Transport for NonRetryableErrorTransport {
969 fn send_request<'a>(
970 &'a self,
971 _method: &'a str,
972 _params: serde_json::Value,
973 _extra_headers: &'a HashMap<String, String>,
974 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
975 self.call_count.fetch_add(1, Ordering::SeqCst);
976 Box::pin(async move { Err(ClientError::InvalidEndpoint("bad url".into())) })
977 }
978
979 fn send_streaming_request<'a>(
980 &'a self,
981 _method: &'a str,
982 _params: serde_json::Value,
983 _extra_headers: &'a HashMap<String, String>,
984 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
985 self.call_count.fetch_add(1, Ordering::SeqCst);
986 Box::pin(async move { Err(ClientError::InvalidEndpoint("bad url".into())) })
987 }
988 }
989
990 #[tokio::test]
991 async fn retry_transport_retries_on_transient_error() {
992 let inner = FailNTransport::new(2, serde_json::json!({"ok": true}));
993 let call_count = Arc::clone(&inner.call_count);
994 let transport = RetryTransport::new(
995 Box::new(inner),
996 RetryPolicy::default()
997 .with_initial_backoff(Duration::from_millis(1))
998 .with_max_retries(3),
999 );
1000
1001 let headers = HashMap::new();
1002 let result = transport
1003 .send_request("GetTask", serde_json::Value::Null, &headers)
1004 .await;
1005 assert!(result.is_ok(), "should succeed after retries");
1006 assert_eq!(
1007 call_count.load(Ordering::SeqCst),
1008 3,
1009 "should have made 3 attempts (2 failures + 1 success)"
1010 );
1011 }
1012
1013 #[tokio::test]
1014 async fn retry_transport_gives_up_after_max_retries() {
1015 let inner = FailNTransport::new(10, serde_json::json!({"ok": true}));
1017 let call_count = Arc::clone(&inner.call_count);
1018 let transport = RetryTransport::new(
1019 Box::new(inner),
1020 RetryPolicy::default()
1021 .with_initial_backoff(Duration::from_millis(1))
1022 .with_max_retries(2),
1023 );
1024
1025 let headers = HashMap::new();
1026 let result = transport
1027 .send_request("GetTask", serde_json::Value::Null, &headers)
1028 .await;
1029 assert!(result.is_err(), "should fail after exhausting retries");
1030 assert_eq!(
1031 call_count.load(Ordering::SeqCst),
1032 3,
1033 "should have made 3 attempts (initial + 2 retries)"
1034 );
1035 }
1036
1037 #[tokio::test]
1038 async fn retry_transport_no_retry_on_non_retryable() {
1039 let inner = NonRetryableErrorTransport::new();
1040 let call_count = Arc::clone(&inner.call_count);
1041 let transport = RetryTransport::new(
1042 Box::new(inner),
1043 RetryPolicy::default()
1044 .with_initial_backoff(Duration::from_millis(1))
1045 .with_max_retries(3),
1046 );
1047
1048 let headers = HashMap::new();
1049 let result = transport
1050 .send_request("GetTask", serde_json::Value::Null, &headers)
1051 .await;
1052 assert!(result.is_err());
1053 assert!(matches!(
1054 result.unwrap_err(),
1055 ClientError::InvalidEndpoint(_)
1056 ));
1057 assert_eq!(
1058 call_count.load(Ordering::SeqCst),
1059 1,
1060 "non-retryable error should not be retried"
1061 );
1062 }
1063
1064 #[tokio::test]
1065 async fn retry_transport_streaming_retries() {
1066 let inner = FailNTransport::new(1, serde_json::json!(null));
1067 let call_count = Arc::clone(&inner.call_count);
1068 let transport = RetryTransport::new(
1069 Box::new(inner),
1070 RetryPolicy::default()
1071 .with_initial_backoff(Duration::from_millis(1))
1072 .with_max_retries(2),
1073 );
1074
1075 let headers = HashMap::new();
1076 let result = transport
1077 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1078 .await;
1079 assert!(result.is_err());
1082 assert_eq!(
1083 call_count.load(Ordering::SeqCst),
1084 2,
1085 "should have retried once for streaming"
1086 );
1087 }
1088
1089 #[tokio::test]
1090 async fn retry_transport_streaming_no_retry_on_non_retryable() {
1091 let inner = NonRetryableErrorTransport::new();
1092 let call_count = Arc::clone(&inner.call_count);
1093 let transport = RetryTransport::new(
1094 Box::new(inner),
1095 RetryPolicy::default()
1096 .with_initial_backoff(Duration::from_millis(1))
1097 .with_max_retries(3),
1098 );
1099
1100 let headers = HashMap::new();
1101 let result = transport
1102 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1103 .await;
1104 assert!(matches!(
1105 result.unwrap_err(),
1106 ClientError::InvalidEndpoint(_)
1107 ));
1108 assert_eq!(
1109 call_count.load(Ordering::SeqCst),
1110 1,
1111 "non-retryable streaming error should not be retried"
1112 );
1113 }
1114
1115 #[tokio::test]
1118 async fn retry_transport_streaming_succeeds_after_retry() {
1119 use tokio::sync::mpsc;
1120
1121 struct FailThenStreamTransport {
1123 call_count: Arc<AtomicUsize>,
1124 }
1125
1126 impl crate::transport::Transport for FailThenStreamTransport {
1127 fn send_request<'a>(
1128 &'a self,
1129 _method: &'a str,
1130 _params: serde_json::Value,
1131 _extra_headers: &'a HashMap<String, String>,
1132 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1133 {
1134 Box::pin(async move { Ok(serde_json::Value::Null) })
1135 }
1136
1137 fn send_streaming_request<'a>(
1138 &'a self,
1139 _method: &'a str,
1140 _params: serde_json::Value,
1141 _extra_headers: &'a HashMap<String, String>,
1142 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1143 let attempt = self.call_count.fetch_add(1, Ordering::SeqCst);
1144 Box::pin(async move {
1145 if attempt == 0 {
1146 Err(ClientError::Timeout("transient timeout".into()))
1147 } else {
1148 let (tx, rx) = mpsc::channel(8);
1150 drop(tx); Ok(EventStream::new(rx))
1152 }
1153 })
1154 }
1155 }
1156
1157 let call_count = Arc::new(AtomicUsize::new(0));
1158 let inner = FailThenStreamTransport {
1159 call_count: Arc::clone(&call_count),
1160 };
1161 let transport = RetryTransport::new(
1162 Box::new(inner),
1163 RetryPolicy::default()
1164 .with_initial_backoff(Duration::from_millis(1))
1165 .with_max_retries(2),
1166 );
1167
1168 let headers = HashMap::new();
1169 let result = transport
1170 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1171 .await;
1172 assert!(result.is_ok(), "streaming should succeed after retry");
1173 assert_eq!(
1174 call_count.load(Ordering::SeqCst),
1175 2,
1176 "should have made 2 attempts (1 failure + 1 success)"
1177 );
1178 }
1179
1180 #[tokio::test]
1181 async fn retry_transport_streaming_exhausts_retries() {
1182 let inner = FailNTransport::new(10, serde_json::json!(null));
1183 let call_count = Arc::clone(&inner.call_count);
1184 let transport = RetryTransport::new(
1185 Box::new(inner),
1186 RetryPolicy::default()
1187 .with_initial_backoff(Duration::from_millis(1))
1188 .with_max_retries(2),
1189 );
1190
1191 let headers = HashMap::new();
1192 let result = transport
1193 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1194 .await;
1195 assert!(result.is_err());
1196 assert_eq!(
1197 call_count.load(Ordering::SeqCst),
1198 3,
1199 "should make 3 attempts total for streaming"
1200 );
1201 }
1202
1203 #[tokio::test]
1204 async fn retry_transport_succeeds_without_retry_on_first_attempt() {
1205 let inner = FailNTransport::new(0, serde_json::json!({"ok": true}));
1206 let call_count = Arc::clone(&inner.call_count);
1207 let transport = RetryTransport::new(
1208 Box::new(inner),
1209 RetryPolicy::default()
1210 .with_initial_backoff(Duration::from_millis(1))
1211 .with_max_retries(3),
1212 );
1213
1214 let headers = HashMap::new();
1215 let result = transport
1216 .send_request("GetTask", serde_json::Value::Null, &headers)
1217 .await;
1218 assert!(result.is_ok());
1219 assert_eq!(
1220 call_count.load(Ordering::SeqCst),
1221 1,
1222 "should succeed on first try"
1223 );
1224 }
1225
1226 #[tokio::test(start_paused = true)]
1232 async fn no_backoff_before_first_attempt() {
1233 let inner = FailNTransport::new(0, serde_json::json!({"ok": true}));
1234 let transport = RetryTransport::new(
1235 Box::new(inner),
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_request("GetTask", serde_json::Value::Null, &headers)
1245 .await;
1246 assert!(result.is_ok());
1247 assert!(
1248 start.elapsed() < Duration::from_secs(1),
1249 "first attempt must not sleep, elapsed: {:?}",
1250 start.elapsed()
1251 );
1252 }
1253
1254 #[tokio::test(start_paused = true)]
1258 async fn backoff_applied_on_retry() {
1259 let inner = FailNTransport::new(1, serde_json::json!({"ok": true}));
1260 let transport = RetryTransport::new(
1261 Box::new(inner),
1262 RetryPolicy::default()
1263 .with_initial_backoff(Duration::from_secs(100))
1264 .with_max_retries(2),
1265 );
1266
1267 let start = tokio::time::Instant::now();
1268 let headers = HashMap::new();
1269 let result = transport
1270 .send_request("GetTask", serde_json::Value::Null, &headers)
1271 .await;
1272 assert!(result.is_ok());
1273 assert!(
1274 start.elapsed() >= Duration::from_secs(50),
1275 "retry should sleep (jittered backoff), elapsed: {:?}",
1276 start.elapsed()
1277 );
1278 }
1279
1280 #[tokio::test(start_paused = true)]
1282 async fn no_backoff_before_first_streaming_attempt() {
1283 use tokio::sync::mpsc;
1284
1285 struct ImmediateStreamTransport;
1286 impl crate::transport::Transport for ImmediateStreamTransport {
1287 fn send_request<'a>(
1288 &'a self,
1289 _method: &'a str,
1290 _params: serde_json::Value,
1291 _extra_headers: &'a HashMap<String, String>,
1292 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1293 {
1294 Box::pin(async { Ok(serde_json::Value::Null) })
1295 }
1296 fn send_streaming_request<'a>(
1297 &'a self,
1298 _method: &'a str,
1299 _params: serde_json::Value,
1300 _extra_headers: &'a HashMap<String, String>,
1301 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1302 Box::pin(async {
1303 let (tx, rx) = mpsc::channel(1);
1304 drop(tx);
1305 Ok(EventStream::new(rx))
1306 })
1307 }
1308 }
1309
1310 let transport = RetryTransport::new(
1311 Box::new(ImmediateStreamTransport),
1312 RetryPolicy::default()
1313 .with_initial_backoff(Duration::from_secs(100))
1314 .with_max_retries(1),
1315 );
1316
1317 let start = tokio::time::Instant::now();
1318 let headers = HashMap::new();
1319 let result = transport
1320 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1321 .await;
1322 assert!(result.is_ok());
1323 assert!(
1324 start.elapsed() < Duration::from_secs(1),
1325 "first streaming attempt must not sleep, elapsed: {:?}",
1326 start.elapsed()
1327 );
1328 }
1329
1330 #[tokio::test(start_paused = true)]
1332 async fn backoff_applied_on_streaming_retry() {
1333 let inner = FailNTransport::new(1, serde_json::json!(null));
1334 let transport = RetryTransport::new(
1335 Box::new(inner),
1336 RetryPolicy::default()
1337 .with_initial_backoff(Duration::from_secs(100))
1338 .with_max_retries(2),
1339 );
1340
1341 let start = tokio::time::Instant::now();
1342 let headers = HashMap::new();
1343 let _result = transport
1344 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1345 .await;
1346 assert!(
1349 start.elapsed() >= Duration::from_secs(50),
1350 "streaming retry should sleep, elapsed: {:?}",
1351 start.elapsed()
1352 );
1353 }
1354
1355 #[test]
1360 fn cap_backoff_zero_multiplier_returns_zero() {
1361 let max = Duration::from_secs(30);
1362 let result = cap_backoff(Duration::from_secs(5), 0.0, max);
1363 assert_eq!(
1364 result,
1365 Duration::ZERO,
1366 "0 * any = 0, should not clamp to max"
1367 );
1368 }
1369
1370 #[test]
1373 fn method_idempotency_classification() {
1374 for m in [
1375 "GetTask",
1376 "ListTasks",
1377 "CancelTask",
1378 "SubscribeToTask",
1379 "GetExtendedAgentCard",
1380 "GetTaskPushNotificationConfig",
1381 "ListTaskPushNotificationConfigs",
1382 "DeleteTaskPushNotificationConfig",
1383 ] {
1384 assert!(is_idempotent_method(m), "{m} should be idempotent");
1385 }
1386 for m in [
1387 "SendMessage",
1388 "SendStreamingMessage",
1389 "CreateTaskPushNotificationConfig",
1390 "SomethingBrandNew",
1391 ] {
1392 assert!(!is_idempotent_method(m), "{m} must be non-idempotent");
1393 }
1394 }
1395
1396 #[test]
1397 fn only_429_503_are_safe_for_non_idempotent() {
1398 let mk = |status| ClientError::UnexpectedStatus {
1399 status,
1400 body: String::new(),
1401 retry_after: None,
1402 };
1403 assert!(safe_to_retry_non_idempotent(&mk(429)));
1404 assert!(safe_to_retry_non_idempotent(&mk(503)));
1405 assert!(!safe_to_retry_non_idempotent(&mk(502)));
1407 assert!(!safe_to_retry_non_idempotent(&mk(504)));
1408 assert!(!safe_to_retry_non_idempotent(&ClientError::Timeout(
1409 "t".into()
1410 )));
1411 assert!(!safe_to_retry_non_idempotent(&ClientError::HttpClient(
1412 "reset".into()
1413 )));
1414 }
1415
1416 #[tokio::test]
1419 async fn non_idempotent_not_retried_on_timeout() {
1420 let inner = FailNTransport::new(5, serde_json::json!({"ok": true}));
1421 let call_count = Arc::clone(&inner.call_count);
1422 let transport = RetryTransport::new(
1423 Box::new(inner),
1424 RetryPolicy::default()
1425 .with_initial_backoff(Duration::from_millis(1))
1426 .with_max_retries(3),
1427 );
1428 let headers = HashMap::new();
1429 let result = transport
1430 .send_request("SendMessage", serde_json::Value::Null, &headers)
1431 .await;
1432 assert!(result.is_err());
1433 assert_eq!(
1434 call_count.load(Ordering::SeqCst),
1435 1,
1436 "SendMessage must not be re-sent on an ambiguous timeout"
1437 );
1438 }
1439
1440 fn keyed_send_params(key: &str) -> serde_json::Value {
1444 serde_json::json!({
1445 "message": {
1446 "id": "msg-1",
1447 "role": "user",
1448 "parts": [{"kind": "text", "text": "hi"}],
1449 "metadata": { IDEMPOTENCY_METADATA_KEY: key },
1450 }
1451 })
1452 }
1453
1454 const GOOD_KEY: &str = "8f14e45fceea167a5a36dedd4bea2543";
1455
1456 #[test]
1457 fn a_keyed_send_is_recognised_as_safe_to_retry() {
1458 for method in ["SendMessage", "SendStreamingMessage"] {
1459 assert!(
1460 carries_idempotency_key(method, &keyed_send_params(GOOD_KEY), true),
1461 "{method} with a key should be retryable"
1462 );
1463 }
1464 }
1465
1466 #[test]
1467 fn an_unkeyed_send_is_not() {
1468 let bare = serde_json::json!({"message": {"id": "m", "role": "user", "parts": []}});
1469 assert!(!carries_idempotency_key("SendMessage", &bare, true));
1470 assert!(!carries_idempotency_key(
1471 "SendMessage",
1472 &serde_json::Value::Null,
1473 true
1474 ));
1475 assert!(!carries_idempotency_key(
1476 "SendMessage",
1477 &serde_json::json!({"message": {"metadata": {"trace": "abc"}}}),
1478 true
1479 ));
1480 }
1481
1482 #[test]
1483 fn a_key_on_a_method_that_is_not_a_send_changes_nothing() {
1484 assert!(!carries_idempotency_key(
1487 "CreateTaskPushNotificationConfig",
1488 &keyed_send_params(GOOD_KEY),
1489 true
1490 ));
1491 }
1492
1493 #[test]
1494 fn a_malformed_key_does_not_make_a_send_retryable() {
1495 for bad in ["short", "", "has a space in it here"] {
1499 assert!(
1500 !carries_idempotency_key("SendMessage", &keyed_send_params(bad), true),
1501 "{bad:?} should not count as a key"
1502 );
1503 }
1504 assert!(!carries_idempotency_key(
1505 "SendMessage",
1506 &serde_json::json!({"message": {"metadata": {IDEMPOTENCY_METADATA_KEY: 42}}}),
1507 true
1508 ));
1509 }
1510
1511 #[tokio::test]
1515 async fn a_keyed_send_is_retried_on_an_ambiguous_timeout() {
1516 let inner = FailNTransport::new(2, serde_json::json!({"ok": true}));
1517 let call_count = Arc::clone(&inner.call_count);
1518 let transport = RetryTransport::new(
1519 Box::new(inner),
1520 RetryPolicy::default()
1521 .with_initial_backoff(Duration::from_millis(1))
1522 .with_max_retries(3),
1523 )
1524 .with_peer_honouring_idempotency(true);
1525 let headers = HashMap::new();
1526 let result = transport
1527 .send_request("SendMessage", keyed_send_params(GOOD_KEY), &headers)
1528 .await;
1529
1530 assert!(
1531 result.is_ok(),
1532 "the retry should have succeeded: {result:?}"
1533 );
1534 assert_eq!(
1535 call_count.load(Ordering::SeqCst),
1536 3,
1537 "two ambiguous failures then a success"
1538 );
1539 }
1540
1541 #[tokio::test]
1544 async fn an_unkeyed_send_is_still_not_retried_on_a_timeout() {
1545 let inner = FailNTransport::new(2, serde_json::json!({"ok": true}));
1546 let call_count = Arc::clone(&inner.call_count);
1547 let transport = RetryTransport::new(
1548 Box::new(inner),
1549 RetryPolicy::default()
1550 .with_initial_backoff(Duration::from_millis(1))
1551 .with_max_retries(3),
1552 );
1553 let headers = HashMap::new();
1554 let bare = serde_json::json!({"message": {"id": "m", "role": "user", "parts": []}});
1555 let result = transport.send_request("SendMessage", bare, &headers).await;
1556
1557 assert!(result.is_err());
1558 assert_eq!(call_count.load(Ordering::SeqCst), 1);
1559 }
1560
1561 #[tokio::test]
1564 async fn non_idempotent_retried_on_503() {
1565 struct Fail503 {
1567 remaining: Arc<AtomicUsize>,
1568 calls: Arc<AtomicUsize>,
1569 }
1570 impl crate::transport::Transport for Fail503 {
1571 fn send_request<'a>(
1572 &'a self,
1573 _m: &'a str,
1574 _p: serde_json::Value,
1575 _h: &'a HashMap<String, String>,
1576 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1577 {
1578 self.calls.fetch_add(1, Ordering::SeqCst);
1579 let left = self.remaining.fetch_sub(1, Ordering::SeqCst);
1580 Box::pin(async move {
1581 if left > 0 {
1582 Err(ClientError::UnexpectedStatus {
1583 status: 503,
1584 body: String::new(),
1585 retry_after: None,
1586 })
1587 } else {
1588 Ok(serde_json::json!({"ok": true}))
1589 }
1590 })
1591 }
1592 fn send_streaming_request<'a>(
1593 &'a self,
1594 _m: &'a str,
1595 _p: serde_json::Value,
1596 _h: &'a HashMap<String, String>,
1597 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1598 Box::pin(async { Err(ClientError::Transport("n/a".into())) })
1599 }
1600 }
1601 let calls = Arc::new(AtomicUsize::new(0));
1602 let inner = Fail503 {
1603 remaining: Arc::new(AtomicUsize::new(1)),
1604 calls: Arc::clone(&calls),
1605 };
1606 let transport = RetryTransport::new(
1607 Box::new(inner),
1608 RetryPolicy::default()
1609 .with_initial_backoff(Duration::from_millis(1))
1610 .with_max_retries(3),
1611 );
1612 let headers = HashMap::new();
1613 let result = transport
1614 .send_request("SendMessage", serde_json::Value::Null, &headers)
1615 .await;
1616 assert!(result.is_ok(), "503 is a safe retry for non-idempotent");
1617 assert_eq!(calls.load(Ordering::SeqCst), 2);
1618 }
1619
1620 #[tokio::test(start_paused = true)]
1622 async fn retry_after_header_is_honored() {
1623 struct RetryAfter429 {
1625 calls: Arc<AtomicUsize>,
1626 }
1627 impl crate::transport::Transport for RetryAfter429 {
1628 fn send_request<'a>(
1629 &'a self,
1630 _m: &'a str,
1631 _p: serde_json::Value,
1632 _h: &'a HashMap<String, String>,
1633 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1634 {
1635 let n = self.calls.fetch_add(1, Ordering::SeqCst);
1636 Box::pin(async move {
1637 if n == 0 {
1638 Err(ClientError::UnexpectedStatus {
1639 status: 429,
1640 body: String::new(),
1641 retry_after: Some(Duration::from_secs(20)),
1642 })
1643 } else {
1644 Ok(serde_json::json!({"ok": true}))
1645 }
1646 })
1647 }
1648 fn send_streaming_request<'a>(
1649 &'a self,
1650 _m: &'a str,
1651 _p: serde_json::Value,
1652 _h: &'a HashMap<String, String>,
1653 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1654 Box::pin(async { Err(ClientError::Transport("n/a".into())) })
1655 }
1656 }
1657 let calls = Arc::new(AtomicUsize::new(0));
1658 let transport = RetryTransport::new(
1659 Box::new(RetryAfter429 {
1660 calls: Arc::clone(&calls),
1661 }),
1662 RetryPolicy::default()
1665 .with_initial_backoff(Duration::from_millis(1))
1666 .with_max_backoff(Duration::from_secs(30)),
1667 );
1668 let start = tokio::time::Instant::now();
1669 let headers = HashMap::new();
1670 let result = transport
1671 .send_request("GetTask", serde_json::Value::Null, &headers)
1672 .await;
1673 assert!(result.is_ok());
1674 assert!(
1675 start.elapsed() >= Duration::from_secs(20),
1676 "should have waited the server-requested 20s, waited {:?}",
1677 start.elapsed()
1678 );
1679 }
1680
1681 #[test]
1682 fn retry_after_clamped_to_max_backoff() {
1683 let err = ClientError::UnexpectedStatus {
1684 status: 503,
1685 body: String::new(),
1686 retry_after: Some(Duration::from_secs(9999)),
1687 };
1688 let delay = retry_delay(Some(&err), Duration::from_secs(1), Duration::from_secs(30));
1689 assert_eq!(
1690 delay,
1691 Duration::from_secs(30),
1692 "retry-after must be clamped"
1693 );
1694 }
1695 #[tokio::test]
1704 async fn a_keyed_send_is_not_retried_against_a_peer_that_never_advertised_it() {
1705 let inner = FailNTransport::new(2, serde_json::json!({"ok": true}));
1706 let call_count = Arc::clone(&inner.call_count);
1707 let transport = RetryTransport::new(
1710 Box::new(inner),
1711 RetryPolicy::default()
1712 .with_initial_backoff(Duration::from_millis(1))
1713 .with_max_retries(3),
1714 );
1715 let headers = HashMap::new();
1716 let result = transport
1717 .send_request("SendMessage", keyed_send_params(GOOD_KEY), &headers)
1718 .await;
1719
1720 assert!(result.is_err(), "the ambiguous failure must surface");
1721 assert_eq!(
1722 call_count.load(Ordering::SeqCst),
1723 1,
1724 "a key the peer never claimed to honour must not make the send retryable"
1725 );
1726 }
1727
1728 #[test]
1731 fn the_predicate_requires_both_the_key_and_the_advertisement() {
1732 assert!(
1733 carries_idempotency_key("SendMessage", &keyed_send_params(GOOD_KEY), true),
1734 "a valid key against an advertising peer is retryable"
1735 );
1736 assert!(
1737 !carries_idempotency_key("SendMessage", &keyed_send_params(GOOD_KEY), false),
1738 "the identical send against a peer that did not advertise is not"
1739 );
1740 }
1741}