1use std::collections::HashMap;
37use std::future::Future;
38use std::pin::Pin;
39use std::time::Duration;
40
41use a2a_protocol_types::idempotency::{IDEMPOTENCY_METADATA_KEY, validate_key};
42
43use crate::error::{ClientError, ClientResult};
44use crate::streaming::EventStream;
45use crate::transport::Transport;
46
47#[derive(Debug, Clone)]
92#[non_exhaustive]
93pub struct RetryPolicy {
94 pub max_retries: u32,
96 pub initial_backoff: Duration,
98 pub max_backoff: Duration,
100 pub backoff_multiplier: f64,
102}
103
104impl Default for RetryPolicy {
105 fn default() -> Self {
106 Self {
107 max_retries: 3,
108 initial_backoff: Duration::from_millis(500),
109 max_backoff: Duration::from_secs(30),
110 backoff_multiplier: 2.0,
111 }
112 }
113}
114
115impl RetryPolicy {
116 #[must_use]
118 pub const fn with_max_retries(mut self, max_retries: u32) -> Self {
119 self.max_retries = max_retries;
120 self
121 }
122
123 #[must_use]
125 pub const fn with_initial_backoff(mut self, backoff: Duration) -> Self {
126 self.initial_backoff = backoff;
127 self
128 }
129
130 #[must_use]
132 pub const fn with_max_backoff(mut self, max: Duration) -> Self {
133 self.max_backoff = max;
134 self
135 }
136
137 #[must_use]
139 pub const fn with_backoff_multiplier(mut self, multiplier: f64) -> Self {
140 self.backoff_multiplier = multiplier;
141 self
142 }
143}
144
145impl ClientError {
148 #[must_use]
159 pub const fn is_retryable(&self) -> bool {
160 match self {
161 Self::Http(_)
162 | Self::HttpClient(_)
163 | Self::Timeout(_)
164 | Self::IncompleteStream { .. }
165 | Self::TooManyPendingRequests { .. } => true,
166 Self::UnexpectedStatus { status, .. } => {
167 matches!(status, 429 | 502 | 503 | 504)
168 }
169 Self::Serialization(_)
171 | Self::Protocol(_)
172 | Self::Transport(_)
173 | Self::InvalidEndpoint(_)
174 | Self::AuthRequired { .. }
175 | Self::ProtocolBindingMismatch(_) => false,
176 }
177 }
178}
179
180fn is_idempotent_method(method: &str) -> bool {
196 matches!(
197 method,
198 "GetTask"
199 | "ListTasks"
200 | "CancelTask"
201 | "SubscribeToTask"
202 | "GetExtendedAgentCard"
203 | "GetTaskPushNotificationConfig"
204 | "ListTaskPushNotificationConfigs"
205 | "DeleteTaskPushNotificationConfig"
206 )
207}
208
209fn carries_idempotency_key(
242 method: &str,
243 params: &serde_json::Value,
244 peer_honours_idempotency: bool,
245) -> bool {
246 if !peer_honours_idempotency {
247 return false;
248 }
249 if !matches!(method, "SendMessage" | "SendStreamingMessage") {
250 return false;
251 }
252 params
253 .get("message")
254 .and_then(|message| message.get("metadata"))
255 .and_then(|metadata| metadata.get(IDEMPOTENCY_METADATA_KEY))
256 .and_then(serde_json::Value::as_str)
257 .is_some_and(|key| validate_key(key).is_ok())
258}
259
260const fn safe_to_retry_non_idempotent(error: &ClientError) -> bool {
270 matches!(
271 error,
272 ClientError::UnexpectedStatus {
273 status: 429 | 503,
274 ..
275 }
276 )
277}
278
279fn retry_delay(
282 last_err: Option<&ClientError>,
283 backoff: Duration,
284 max_backoff: Duration,
285) -> Duration {
286 last_err
287 .and_then(ClientError::retry_after)
288 .map_or_else(|| jittered(backoff), |after| after.min(max_backoff))
289}
290
291pub(crate) struct RetryTransport {
296 inner: Box<dyn Transport>,
297 policy: RetryPolicy,
298 peer_honours_idempotency: bool,
304}
305
306impl RetryTransport {
307 pub(crate) fn new(inner: Box<dyn Transport>, policy: RetryPolicy) -> Self {
309 Self {
310 inner,
311 policy,
312 peer_honours_idempotency: false,
313 }
314 }
315
316 #[must_use]
318 pub(crate) const fn with_peer_honouring_idempotency(mut self, honours: bool) -> Self {
319 self.peer_honours_idempotency = honours;
320 self
321 }
322}
323
324impl Transport for RetryTransport {
325 fn send_request<'a>(
326 &'a self,
327 method: &'a str,
328 params: serde_json::Value,
329 extra_headers: &'a HashMap<String, String>,
330 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
331 Box::pin(async move {
332 let mut last_err: Option<ClientError> = None;
333 let mut backoff = self.policy.initial_backoff;
334 let idempotent = is_idempotent_method(method)
335 || carries_idempotency_key(method, ¶ms, self.peer_honours_idempotency);
336
337 let serialized = serde_json::to_vec(¶ms).map_err(ClientError::Serialization)?;
340
341 for attempt in 0..=self.policy.max_retries {
342 if attempt > 0 {
343 let delay = retry_delay(last_err.as_ref(), backoff, self.policy.max_backoff);
344 trace_info!(method, attempt, ?delay, "retrying after backoff");
345 tokio::time::sleep(delay).await;
346 backoff = cap_backoff(
347 backoff,
348 self.policy.backoff_multiplier,
349 self.policy.max_backoff,
350 );
351 }
352
353 let attempt_params: serde_json::Value =
354 serde_json::from_slice(&serialized).map_err(ClientError::Serialization)?;
355
356 match self
357 .inner
358 .send_request(method, attempt_params, extra_headers)
359 .await
360 {
361 Ok(result) => return Ok(result),
362 Err(e)
369 if e.is_retryable() && (idempotent || safe_to_retry_non_idempotent(&e)) =>
370 {
371 trace_warn!(method, attempt, error = %e, "transient error, will retry");
372 last_err = Some(e);
373 }
374 Err(e) => return Err(e),
375 }
376 }
377
378 Err(last_err.expect("at least one attempt was made"))
379 })
380 }
381
382 fn send_streaming_request<'a>(
383 &'a self,
384 method: &'a str,
385 params: serde_json::Value,
386 extra_headers: &'a HashMap<String, String>,
387 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
388 Box::pin(async move {
389 let mut last_err: Option<ClientError> = None;
390 let mut backoff = self.policy.initial_backoff;
391 let idempotent = is_idempotent_method(method)
392 || carries_idempotency_key(method, ¶ms, self.peer_honours_idempotency);
393
394 let serialized = serde_json::to_vec(¶ms).map_err(ClientError::Serialization)?;
397
398 for attempt in 0..=self.policy.max_retries {
399 if attempt > 0 {
400 let delay = retry_delay(last_err.as_ref(), backoff, self.policy.max_backoff);
401 trace_info!(
402 method,
403 attempt,
404 ?delay,
405 "retrying stream connect after backoff"
406 );
407 tokio::time::sleep(delay).await;
408 backoff = cap_backoff(
409 backoff,
410 self.policy.backoff_multiplier,
411 self.policy.max_backoff,
412 );
413 }
414
415 let attempt_params: serde_json::Value =
416 serde_json::from_slice(&serialized).map_err(ClientError::Serialization)?;
417
418 match self
419 .inner
420 .send_streaming_request(method, attempt_params, extra_headers)
421 .await
422 {
423 Ok(stream) => return Ok(stream),
424 Err(e)
428 if e.is_retryable() && (idempotent || safe_to_retry_non_idempotent(&e)) =>
429 {
430 trace_warn!(method, attempt, error = %e, "transient error, will retry");
431 last_err = Some(e);
432 }
433 Err(e) => return Err(e),
434 }
435 }
436
437 Err(last_err.expect("at least one attempt was made"))
438 })
439 }
440}
441
442fn cap_backoff(current: Duration, multiplier: f64, max: Duration) -> Duration {
449 let next_secs = current.as_secs_f64() * multiplier;
450 Duration::try_from_secs_f64(next_secs).map_or(max, |next| {
455 std::cmp::min(next, max)
460 })
461}
462
463#[allow(clippy::cast_precision_loss)] fn jitter_factor_from_bits(random_bits: u64) -> f64 {
470 (random_bits as f64 / u64::MAX as f64).mul_add(0.5, 0.5)
471}
472
473fn apply_jitter(backoff: Duration, factor: f64) -> Duration {
478 let jittered_secs = backoff.as_secs_f64() * factor;
479 Duration::try_from_secs_f64(jittered_secs).unwrap_or(backoff)
485}
486
487fn jittered(backoff: Duration) -> Duration {
494 use std::hash::{BuildHasher, Hasher};
495 let mut hasher = std::collections::hash_map::RandomState::new().build_hasher();
496 hasher.write_u128(backoff.as_nanos());
498 let factor = jitter_factor_from_bits(hasher.finish());
499 apply_jitter(backoff, factor)
500}
501
502#[cfg(test)]
505mod tests {
506 use super::*;
507
508 #[test]
509 fn http_errors_are_retryable() {
510 let e = ClientError::HttpClient("connection refused".into());
511 assert!(e.is_retryable());
512 }
513
514 #[test]
515 fn timeout_is_retryable() {
516 let e = ClientError::Timeout("request timed out".into());
517 assert!(e.is_retryable());
518 }
519
520 #[test]
522 fn incomplete_stream_is_retryable() {
523 let e = ClientError::IncompleteStream {
524 last_event_id: Some("7".into()),
525 detail: "cut off".into(),
526 };
527 assert!(e.is_retryable());
528 assert_eq!(
529 e.to_string(),
530 "stream ended before its final event: cut off"
531 );
532 }
533
534 #[test]
535 fn status_503_is_retryable() {
536 let e = ClientError::UnexpectedStatus {
537 status: 503,
538 body: "Service Unavailable".into(),
539 retry_after: None,
540 };
541 assert!(e.is_retryable());
542 }
543
544 #[test]
545 fn status_429_is_retryable() {
546 let e = ClientError::UnexpectedStatus {
547 status: 429,
548 body: "Too Many Requests".into(),
549 retry_after: None,
550 };
551 assert!(e.is_retryable());
552 }
553
554 #[test]
555 fn status_404_is_not_retryable() {
556 let e = ClientError::UnexpectedStatus {
557 status: 404,
558 body: "Not Found".into(),
559 retry_after: None,
560 };
561 assert!(!e.is_retryable());
562 }
563
564 #[test]
565 fn serialization_error_is_not_retryable() {
566 let e = ClientError::Serialization(serde_json::from_str::<String>("not json").unwrap_err());
567 assert!(!e.is_retryable());
568 }
569
570 #[test]
571 fn protocol_error_is_not_retryable() {
572 let e = ClientError::Protocol(a2a_protocol_types::A2aError::task_not_found("t1"));
573 assert!(!e.is_retryable());
574 }
575
576 #[test]
577 fn default_retry_policy() {
578 let p = RetryPolicy::default();
579 assert_eq!(p.max_retries, 3);
580 assert_eq!(p.initial_backoff, Duration::from_millis(500));
581 assert_eq!(p.max_backoff, Duration::from_secs(30));
582 assert!((p.backoff_multiplier - 2.0).abs() < f64::EPSILON);
583 }
584
585 #[test]
586 fn cap_backoff_works() {
587 let result = cap_backoff(Duration::from_secs(1), 2.0, Duration::from_secs(5));
588 assert_eq!(result, Duration::from_secs(2));
589
590 let result = cap_backoff(Duration::from_secs(4), 2.0, Duration::from_secs(5));
591 assert_eq!(result, Duration::from_secs(5));
592 }
593
594 #[test]
595 fn status_502_is_retryable() {
596 let e = ClientError::UnexpectedStatus {
597 status: 502,
598 body: "Bad Gateway".into(),
599 retry_after: None,
600 };
601 assert!(e.is_retryable());
602 }
603
604 #[test]
605 fn status_504_is_retryable() {
606 let e = ClientError::UnexpectedStatus {
607 status: 504,
608 body: "Gateway Timeout".into(),
609 retry_after: None,
610 };
611 assert!(e.is_retryable());
612 }
613
614 #[test]
616 fn status_boundary_not_retryable() {
617 for status in [428, 430, 500, 501, 505] {
618 let e = ClientError::UnexpectedStatus {
619 status,
620 body: String::new(),
621 retry_after: None,
622 };
623 assert!(!e.is_retryable(), "status {status} should not be retryable");
624 }
625 }
626
627 #[test]
628 fn retry_policy_builder_methods() {
629 let p = RetryPolicy::default()
630 .with_max_retries(5)
631 .with_initial_backoff(Duration::from_secs(1))
632 .with_max_backoff(Duration::from_secs(60))
633 .with_backoff_multiplier(3.0);
634 assert_eq!(p.max_retries, 5);
635 assert_eq!(p.initial_backoff, Duration::from_secs(1));
636 assert_eq!(p.max_backoff, Duration::from_secs(60));
637 assert!((p.backoff_multiplier - 3.0).abs() < f64::EPSILON);
638 }
639
640 #[test]
641 fn cap_backoff_exact_boundary() {
642 let result = cap_backoff(Duration::from_secs(5), 1.0, Duration::from_secs(5));
644 assert_eq!(result, Duration::from_secs(5));
645
646 let result = cap_backoff(Duration::from_millis(1), 2.0, Duration::from_secs(5));
648 assert_eq!(result, Duration::from_millis(2));
649 }
650
651 #[test]
652 fn cap_backoff_infinity_returns_max() {
653 let max = Duration::from_secs(30);
655 let result = cap_backoff(Duration::from_secs(u64::MAX / 2), f64::MAX, max);
656 assert_eq!(result, max, "infinity should clamp to max");
657 }
658
659 #[test]
660 fn cap_backoff_finite_overflow_returns_max() {
661 let max = Duration::from_secs(30);
666 let result = cap_backoff(Duration::from_secs(10_000_000_000_000_000_000), 10.0, max);
667 assert_eq!(
668 result, max,
669 "finite-but-overflowing backoff should clamp to max"
670 );
671 }
672
673 #[test]
675 fn jittered_backoff_in_expected_range() {
676 let backoff = Duration::from_secs(2);
677 for _ in 0..100 {
679 let result = jittered(backoff);
680 assert!(
681 result >= Duration::from_secs(1),
682 "jittered backoff should be >= backoff/2, got {result:?}"
683 );
684 assert!(
685 result <= backoff,
686 "jittered backoff should be <= backoff, got {result:?}"
687 );
688 }
689 }
690
691 #[test]
693 fn jittered_zero_backoff() {
694 let result = jittered(Duration::ZERO);
695 assert_eq!(result, Duration::ZERO);
696 }
697
698 #[test]
699 fn cap_backoff_nan_returns_max() {
700 let max = Duration::from_secs(30);
701 let result = cap_backoff(Duration::from_secs(0), f64::NAN, max);
702 assert_eq!(result, max, "NaN should clamp to max");
703 }
704
705 #[test]
710 fn jitter_factor_from_bits_zero() {
711 let f = jitter_factor_from_bits(0);
712 assert!(
713 (f - 0.5).abs() < f64::EPSILON,
714 "factor(0) should be 0.5, got {f}"
715 );
716 }
717
718 #[test]
720 fn jitter_factor_from_bits_midpoint() {
721 let f = jitter_factor_from_bits(u64::MAX / 2);
722 assert!(
724 (0.74..=0.76).contains(&f),
725 "factor(u64::MAX/2) should be ~0.75, got {f}"
726 );
727 }
728
729 #[test]
731 fn jitter_factor_from_bits_max() {
732 let f = jitter_factor_from_bits(u64::MAX);
733 assert!(
736 (0.9..=1.0).contains(&f),
737 "factor(u64::MAX) should be ~1.0, got {f}"
738 );
739 }
740
741 #[test]
745 fn jitter_factor_from_bits_always_in_half_to_one() {
746 for bits in [
747 0_u64,
748 1,
749 7,
750 42,
751 1 << 20,
752 1 << 50,
753 u64::MAX / 4,
754 u64::MAX / 2,
755 u64::MAX,
756 ] {
757 let f = jitter_factor_from_bits(bits);
758 assert!(
759 (0.5..=1.0).contains(&f),
760 "factor({bits}) = {f} out of [0.5, 1.0]"
761 );
762 }
763 }
764
765 #[test]
774 fn apply_jitter_normal_factor() {
775 assert_eq!(
777 apply_jitter(Duration::from_secs(2), 0.5),
778 Duration::from_secs(1)
779 );
780 assert_eq!(
782 apply_jitter(Duration::from_secs(4), 0.75),
783 Duration::from_secs(3)
784 );
785 assert_eq!(
787 apply_jitter(Duration::from_secs(5), 1.0),
788 Duration::from_secs(5)
789 );
790 }
791
792 #[test]
796 fn apply_jitter_zero_factor_returns_zero() {
797 assert_eq!(
798 apply_jitter(Duration::from_secs(5), 0.0),
799 Duration::ZERO,
800 "factor=0.0 must produce Duration::ZERO via from_secs_f64 path"
801 );
802 }
803
804 #[test]
808 fn apply_jitter_negative_factor_returns_backoff() {
809 assert_eq!(
810 apply_jitter(Duration::from_secs(3), -0.5),
811 Duration::from_secs(3),
812 "negative factor must short-circuit to backoff"
813 );
814 }
815
816 #[test]
825 fn apply_jitter_infinite_factor_returns_backoff() {
826 assert_eq!(
827 apply_jitter(Duration::from_secs(2), f64::INFINITY),
828 Duration::from_secs(2),
829 "infinite factor must short-circuit to backoff"
830 );
831 }
832
833 #[test]
834 fn apply_jitter_nan_factor_returns_backoff() {
835 assert_eq!(
836 apply_jitter(Duration::from_secs(4), f64::NAN),
837 Duration::from_secs(4),
838 "NaN factor must short-circuit to backoff"
839 );
840 }
841
842 use std::collections::HashMap;
845 use std::future::Future;
846 use std::pin::Pin;
847 use std::sync::Arc;
848 use std::sync::atomic::{AtomicUsize, Ordering};
849
850 use crate::streaming::EventStream;
851
852 #[tokio::test(start_paused = true)]
869 async fn the_documented_backoff_arithmetic_is_what_it_does() {
870 let inner = FailNTransport::new(usize::MAX, serde_json::json!(null));
871 let calls = Arc::clone(&inner.call_count);
872 let retry = RetryTransport::new(
873 Box::new(inner),
874 RetryPolicy {
875 max_retries: 3,
876 initial_backoff: Duration::from_millis(500),
877 max_backoff: Duration::from_secs(30),
878 backoff_multiplier: 1.0,
881 },
882 );
883
884 let started = tokio::time::Instant::now();
885 let result = retry
889 .send_request("GetTask", serde_json::json!({}), &HashMap::new())
890 .await;
891 let elapsed = started.elapsed();
892
893 assert!(
894 result.is_err(),
895 "every attempt fails, so the call must fail"
896 );
897 assert_eq!(
898 calls.load(Ordering::SeqCst),
899 4,
900 "max_retries 3 must mean 4 attempts — the formula's (max_retries + 1)"
901 );
902
903 assert!(
914 elapsed >= Duration::from_millis(750) && elapsed < Duration::from_millis(1_500),
915 "three fully-jittered 500ms gaps must total in [750ms, 1500ms), got {elapsed:?}"
916 );
917 }
918
919 struct FailNTransport {
921 failures_remaining: Arc<AtomicUsize>,
922 success_response: serde_json::Value,
923 call_count: Arc<AtomicUsize>,
924 }
925
926 impl FailNTransport {
927 fn new(fail_count: usize, response: serde_json::Value) -> Self {
928 Self {
929 failures_remaining: Arc::new(AtomicUsize::new(fail_count)),
930 success_response: response,
931 call_count: Arc::new(AtomicUsize::new(0)),
932 }
933 }
934 }
935
936 impl crate::transport::Transport for FailNTransport {
937 fn send_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<serde_json::Value>> + Send + 'a>> {
943 self.call_count.fetch_add(1, Ordering::SeqCst);
944 let remaining = self.failures_remaining.fetch_sub(1, Ordering::SeqCst);
945 let resp = self.success_response.clone();
946 Box::pin(async move {
947 if remaining > 0 {
948 Err(ClientError::Timeout("transient".into()))
949 } else {
950 Ok(resp)
951 }
952 })
953 }
954
955 fn send_streaming_request<'a>(
956 &'a self,
957 _method: &'a str,
958 _params: serde_json::Value,
959 _extra_headers: &'a HashMap<String, String>,
960 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
961 self.call_count.fetch_add(1, Ordering::SeqCst);
962 let remaining = self.failures_remaining.fetch_sub(1, Ordering::SeqCst);
963 Box::pin(async move {
964 if remaining > 0 {
965 Err(ClientError::Timeout("transient".into()))
966 } else {
967 Err(ClientError::Transport("streaming not mocked".into()))
968 }
969 })
970 }
971 }
972
973 struct NonRetryableErrorTransport {
975 call_count: Arc<AtomicUsize>,
976 }
977
978 impl NonRetryableErrorTransport {
979 fn new() -> Self {
980 Self {
981 call_count: Arc::new(AtomicUsize::new(0)),
982 }
983 }
984 }
985
986 impl crate::transport::Transport for NonRetryableErrorTransport {
987 fn send_request<'a>(
988 &'a self,
989 _method: &'a str,
990 _params: serde_json::Value,
991 _extra_headers: &'a HashMap<String, String>,
992 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
993 self.call_count.fetch_add(1, Ordering::SeqCst);
994 Box::pin(async move { Err(ClientError::InvalidEndpoint("bad url".into())) })
995 }
996
997 fn send_streaming_request<'a>(
998 &'a self,
999 _method: &'a str,
1000 _params: serde_json::Value,
1001 _extra_headers: &'a HashMap<String, String>,
1002 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1003 self.call_count.fetch_add(1, Ordering::SeqCst);
1004 Box::pin(async move { Err(ClientError::InvalidEndpoint("bad url".into())) })
1005 }
1006 }
1007
1008 #[tokio::test]
1009 async fn retry_transport_retries_on_transient_error() {
1010 let inner = FailNTransport::new(2, serde_json::json!({"ok": true}));
1011 let call_count = Arc::clone(&inner.call_count);
1012 let transport = RetryTransport::new(
1013 Box::new(inner),
1014 RetryPolicy::default()
1015 .with_initial_backoff(Duration::from_millis(1))
1016 .with_max_retries(3),
1017 );
1018
1019 let headers = HashMap::new();
1020 let result = transport
1021 .send_request("GetTask", serde_json::Value::Null, &headers)
1022 .await;
1023 assert!(result.is_ok(), "should succeed after retries");
1024 assert_eq!(
1025 call_count.load(Ordering::SeqCst),
1026 3,
1027 "should have made 3 attempts (2 failures + 1 success)"
1028 );
1029 }
1030
1031 #[tokio::test]
1032 async fn retry_transport_gives_up_after_max_retries() {
1033 let inner = FailNTransport::new(10, serde_json::json!({"ok": true}));
1035 let call_count = Arc::clone(&inner.call_count);
1036 let transport = RetryTransport::new(
1037 Box::new(inner),
1038 RetryPolicy::default()
1039 .with_initial_backoff(Duration::from_millis(1))
1040 .with_max_retries(2),
1041 );
1042
1043 let headers = HashMap::new();
1044 let result = transport
1045 .send_request("GetTask", serde_json::Value::Null, &headers)
1046 .await;
1047 assert!(result.is_err(), "should fail after exhausting retries");
1048 assert_eq!(
1049 call_count.load(Ordering::SeqCst),
1050 3,
1051 "should have made 3 attempts (initial + 2 retries)"
1052 );
1053 }
1054
1055 #[tokio::test]
1056 async fn retry_transport_no_retry_on_non_retryable() {
1057 let inner = NonRetryableErrorTransport::new();
1058 let call_count = Arc::clone(&inner.call_count);
1059 let transport = RetryTransport::new(
1060 Box::new(inner),
1061 RetryPolicy::default()
1062 .with_initial_backoff(Duration::from_millis(1))
1063 .with_max_retries(3),
1064 );
1065
1066 let headers = HashMap::new();
1067 let result = transport
1068 .send_request("GetTask", serde_json::Value::Null, &headers)
1069 .await;
1070 assert!(result.is_err());
1071 assert!(matches!(
1072 result.unwrap_err(),
1073 ClientError::InvalidEndpoint(_)
1074 ));
1075 assert_eq!(
1076 call_count.load(Ordering::SeqCst),
1077 1,
1078 "non-retryable error should not be retried"
1079 );
1080 }
1081
1082 #[tokio::test]
1083 async fn retry_transport_streaming_retries() {
1084 let inner = FailNTransport::new(1, serde_json::json!(null));
1085 let call_count = Arc::clone(&inner.call_count);
1086 let transport = RetryTransport::new(
1087 Box::new(inner),
1088 RetryPolicy::default()
1089 .with_initial_backoff(Duration::from_millis(1))
1090 .with_max_retries(2),
1091 );
1092
1093 let headers = HashMap::new();
1094 let result = transport
1095 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1096 .await;
1097 assert!(result.is_err());
1100 assert_eq!(
1101 call_count.load(Ordering::SeqCst),
1102 2,
1103 "should have retried once for streaming"
1104 );
1105 }
1106
1107 #[tokio::test]
1108 async fn retry_transport_streaming_no_retry_on_non_retryable() {
1109 let inner = NonRetryableErrorTransport::new();
1110 let call_count = Arc::clone(&inner.call_count);
1111 let transport = RetryTransport::new(
1112 Box::new(inner),
1113 RetryPolicy::default()
1114 .with_initial_backoff(Duration::from_millis(1))
1115 .with_max_retries(3),
1116 );
1117
1118 let headers = HashMap::new();
1119 let result = transport
1120 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1121 .await;
1122 assert!(matches!(
1123 result.unwrap_err(),
1124 ClientError::InvalidEndpoint(_)
1125 ));
1126 assert_eq!(
1127 call_count.load(Ordering::SeqCst),
1128 1,
1129 "non-retryable streaming error should not be retried"
1130 );
1131 }
1132
1133 #[tokio::test]
1136 async fn retry_transport_streaming_succeeds_after_retry() {
1137 use tokio::sync::mpsc;
1138
1139 struct FailThenStreamTransport {
1141 call_count: Arc<AtomicUsize>,
1142 }
1143
1144 impl crate::transport::Transport for FailThenStreamTransport {
1145 fn send_request<'a>(
1146 &'a self,
1147 _method: &'a str,
1148 _params: serde_json::Value,
1149 _extra_headers: &'a HashMap<String, String>,
1150 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1151 {
1152 Box::pin(async move { Ok(serde_json::Value::Null) })
1153 }
1154
1155 fn send_streaming_request<'a>(
1156 &'a self,
1157 _method: &'a str,
1158 _params: serde_json::Value,
1159 _extra_headers: &'a HashMap<String, String>,
1160 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1161 let attempt = self.call_count.fetch_add(1, Ordering::SeqCst);
1162 Box::pin(async move {
1163 if attempt == 0 {
1164 Err(ClientError::Timeout("transient timeout".into()))
1165 } else {
1166 let (tx, rx) = mpsc::channel(8);
1168 drop(tx); Ok(EventStream::new(rx))
1170 }
1171 })
1172 }
1173 }
1174
1175 let call_count = Arc::new(AtomicUsize::new(0));
1176 let inner = FailThenStreamTransport {
1177 call_count: Arc::clone(&call_count),
1178 };
1179 let transport = RetryTransport::new(
1180 Box::new(inner),
1181 RetryPolicy::default()
1182 .with_initial_backoff(Duration::from_millis(1))
1183 .with_max_retries(2),
1184 );
1185
1186 let headers = HashMap::new();
1187 let result = transport
1188 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1189 .await;
1190 assert!(result.is_ok(), "streaming should succeed after retry");
1191 assert_eq!(
1192 call_count.load(Ordering::SeqCst),
1193 2,
1194 "should have made 2 attempts (1 failure + 1 success)"
1195 );
1196 }
1197
1198 #[tokio::test]
1199 async fn retry_transport_streaming_exhausts_retries() {
1200 let inner = FailNTransport::new(10, serde_json::json!(null));
1201 let call_count = Arc::clone(&inner.call_count);
1202 let transport = RetryTransport::new(
1203 Box::new(inner),
1204 RetryPolicy::default()
1205 .with_initial_backoff(Duration::from_millis(1))
1206 .with_max_retries(2),
1207 );
1208
1209 let headers = HashMap::new();
1210 let result = transport
1211 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1212 .await;
1213 assert!(result.is_err());
1214 assert_eq!(
1215 call_count.load(Ordering::SeqCst),
1216 3,
1217 "should make 3 attempts total for streaming"
1218 );
1219 }
1220
1221 #[tokio::test]
1222 async fn retry_transport_succeeds_without_retry_on_first_attempt() {
1223 let inner = FailNTransport::new(0, serde_json::json!({"ok": true}));
1224 let call_count = Arc::clone(&inner.call_count);
1225 let transport = RetryTransport::new(
1226 Box::new(inner),
1227 RetryPolicy::default()
1228 .with_initial_backoff(Duration::from_millis(1))
1229 .with_max_retries(3),
1230 );
1231
1232 let headers = HashMap::new();
1233 let result = transport
1234 .send_request("GetTask", serde_json::Value::Null, &headers)
1235 .await;
1236 assert!(result.is_ok());
1237 assert_eq!(
1238 call_count.load(Ordering::SeqCst),
1239 1,
1240 "should succeed on first try"
1241 );
1242 }
1243
1244 #[tokio::test(start_paused = true)]
1250 async fn no_backoff_before_first_attempt() {
1251 let inner = FailNTransport::new(0, serde_json::json!({"ok": true}));
1252 let transport = RetryTransport::new(
1253 Box::new(inner),
1254 RetryPolicy::default()
1255 .with_initial_backoff(Duration::from_secs(100))
1256 .with_max_retries(1),
1257 );
1258
1259 let start = tokio::time::Instant::now();
1260 let headers = HashMap::new();
1261 let result = transport
1262 .send_request("GetTask", serde_json::Value::Null, &headers)
1263 .await;
1264 assert!(result.is_ok());
1265 assert!(
1266 start.elapsed() < Duration::from_secs(1),
1267 "first attempt must not sleep, elapsed: {:?}",
1268 start.elapsed()
1269 );
1270 }
1271
1272 #[tokio::test(start_paused = true)]
1276 async fn backoff_applied_on_retry() {
1277 let inner = FailNTransport::new(1, serde_json::json!({"ok": true}));
1278 let transport = RetryTransport::new(
1279 Box::new(inner),
1280 RetryPolicy::default()
1281 .with_initial_backoff(Duration::from_secs(100))
1282 .with_max_retries(2),
1283 );
1284
1285 let start = tokio::time::Instant::now();
1286 let headers = HashMap::new();
1287 let result = transport
1288 .send_request("GetTask", serde_json::Value::Null, &headers)
1289 .await;
1290 assert!(result.is_ok());
1291 assert!(
1292 start.elapsed() >= Duration::from_secs(50),
1293 "retry should sleep (jittered backoff), elapsed: {:?}",
1294 start.elapsed()
1295 );
1296 }
1297
1298 #[tokio::test(start_paused = true)]
1300 async fn no_backoff_before_first_streaming_attempt() {
1301 use tokio::sync::mpsc;
1302
1303 struct ImmediateStreamTransport;
1304 impl crate::transport::Transport for ImmediateStreamTransport {
1305 fn send_request<'a>(
1306 &'a self,
1307 _method: &'a str,
1308 _params: serde_json::Value,
1309 _extra_headers: &'a HashMap<String, String>,
1310 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1311 {
1312 Box::pin(async { Ok(serde_json::Value::Null) })
1313 }
1314 fn send_streaming_request<'a>(
1315 &'a self,
1316 _method: &'a str,
1317 _params: serde_json::Value,
1318 _extra_headers: &'a HashMap<String, String>,
1319 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1320 Box::pin(async {
1321 let (tx, rx) = mpsc::channel(1);
1322 drop(tx);
1323 Ok(EventStream::new(rx))
1324 })
1325 }
1326 }
1327
1328 let transport = RetryTransport::new(
1329 Box::new(ImmediateStreamTransport),
1330 RetryPolicy::default()
1331 .with_initial_backoff(Duration::from_secs(100))
1332 .with_max_retries(1),
1333 );
1334
1335 let start = tokio::time::Instant::now();
1336 let headers = HashMap::new();
1337 let result = transport
1338 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1339 .await;
1340 assert!(result.is_ok());
1341 assert!(
1342 start.elapsed() < Duration::from_secs(1),
1343 "first streaming attempt must not sleep, elapsed: {:?}",
1344 start.elapsed()
1345 );
1346 }
1347
1348 #[tokio::test(start_paused = true)]
1350 async fn backoff_applied_on_streaming_retry() {
1351 let inner = FailNTransport::new(1, serde_json::json!(null));
1352 let transport = RetryTransport::new(
1353 Box::new(inner),
1354 RetryPolicy::default()
1355 .with_initial_backoff(Duration::from_secs(100))
1356 .with_max_retries(2),
1357 );
1358
1359 let start = tokio::time::Instant::now();
1360 let headers = HashMap::new();
1361 let _result = transport
1362 .send_streaming_request("SubscribeToTask", serde_json::Value::Null, &headers)
1363 .await;
1364 assert!(
1367 start.elapsed() >= Duration::from_secs(50),
1368 "streaming retry should sleep, elapsed: {:?}",
1369 start.elapsed()
1370 );
1371 }
1372
1373 #[test]
1378 fn cap_backoff_zero_multiplier_returns_zero() {
1379 let max = Duration::from_secs(30);
1380 let result = cap_backoff(Duration::from_secs(5), 0.0, max);
1381 assert_eq!(
1382 result,
1383 Duration::ZERO,
1384 "0 * any = 0, should not clamp to max"
1385 );
1386 }
1387
1388 #[test]
1391 fn method_idempotency_classification() {
1392 for m in [
1393 "GetTask",
1394 "ListTasks",
1395 "CancelTask",
1396 "SubscribeToTask",
1397 "GetExtendedAgentCard",
1398 "GetTaskPushNotificationConfig",
1399 "ListTaskPushNotificationConfigs",
1400 "DeleteTaskPushNotificationConfig",
1401 ] {
1402 assert!(is_idempotent_method(m), "{m} should be idempotent");
1403 }
1404 for m in [
1405 "SendMessage",
1406 "SendStreamingMessage",
1407 "CreateTaskPushNotificationConfig",
1408 "SomethingBrandNew",
1409 ] {
1410 assert!(!is_idempotent_method(m), "{m} must be non-idempotent");
1411 }
1412 }
1413
1414 #[test]
1415 fn only_429_503_are_safe_for_non_idempotent() {
1416 let mk = |status| ClientError::UnexpectedStatus {
1417 status,
1418 body: String::new(),
1419 retry_after: None,
1420 };
1421 assert!(safe_to_retry_non_idempotent(&mk(429)));
1422 assert!(safe_to_retry_non_idempotent(&mk(503)));
1423 assert!(!safe_to_retry_non_idempotent(&mk(502)));
1425 assert!(!safe_to_retry_non_idempotent(&mk(504)));
1426 assert!(!safe_to_retry_non_idempotent(&ClientError::Timeout(
1427 "t".into()
1428 )));
1429 assert!(!safe_to_retry_non_idempotent(&ClientError::HttpClient(
1430 "reset".into()
1431 )));
1432 }
1433
1434 #[tokio::test]
1437 async fn non_idempotent_not_retried_on_timeout() {
1438 let inner = FailNTransport::new(5, serde_json::json!({"ok": true}));
1439 let call_count = Arc::clone(&inner.call_count);
1440 let transport = RetryTransport::new(
1441 Box::new(inner),
1442 RetryPolicy::default()
1443 .with_initial_backoff(Duration::from_millis(1))
1444 .with_max_retries(3),
1445 );
1446 let headers = HashMap::new();
1447 let result = transport
1448 .send_request("SendMessage", serde_json::Value::Null, &headers)
1449 .await;
1450 assert!(result.is_err());
1451 assert_eq!(
1452 call_count.load(Ordering::SeqCst),
1453 1,
1454 "SendMessage must not be re-sent on an ambiguous timeout"
1455 );
1456 }
1457
1458 fn keyed_send_params(key: &str) -> serde_json::Value {
1462 serde_json::json!({
1463 "message": {
1464 "id": "msg-1",
1465 "role": "user",
1466 "parts": [{"kind": "text", "text": "hi"}],
1467 "metadata": { IDEMPOTENCY_METADATA_KEY: key },
1468 }
1469 })
1470 }
1471
1472 const GOOD_KEY: &str = "8f14e45fceea167a5a36dedd4bea2543";
1473
1474 #[test]
1475 fn a_keyed_send_is_recognised_as_safe_to_retry() {
1476 for method in ["SendMessage", "SendStreamingMessage"] {
1477 assert!(
1478 carries_idempotency_key(method, &keyed_send_params(GOOD_KEY), true),
1479 "{method} with a key should be retryable"
1480 );
1481 }
1482 }
1483
1484 #[test]
1485 fn an_unkeyed_send_is_not() {
1486 let bare = serde_json::json!({"message": {"id": "m", "role": "user", "parts": []}});
1487 assert!(!carries_idempotency_key("SendMessage", &bare, true));
1488 assert!(!carries_idempotency_key(
1489 "SendMessage",
1490 &serde_json::Value::Null,
1491 true
1492 ));
1493 assert!(!carries_idempotency_key(
1494 "SendMessage",
1495 &serde_json::json!({"message": {"metadata": {"trace": "abc"}}}),
1496 true
1497 ));
1498 }
1499
1500 #[test]
1501 fn a_key_on_a_method_that_is_not_a_send_changes_nothing() {
1502 assert!(!carries_idempotency_key(
1505 "CreateTaskPushNotificationConfig",
1506 &keyed_send_params(GOOD_KEY),
1507 true
1508 ));
1509 }
1510
1511 #[test]
1512 fn a_malformed_key_does_not_make_a_send_retryable() {
1513 for bad in ["short", "", "has a space in it here"] {
1517 assert!(
1518 !carries_idempotency_key("SendMessage", &keyed_send_params(bad), true),
1519 "{bad:?} should not count as a key"
1520 );
1521 }
1522 assert!(!carries_idempotency_key(
1523 "SendMessage",
1524 &serde_json::json!({"message": {"metadata": {IDEMPOTENCY_METADATA_KEY: 42}}}),
1525 true
1526 ));
1527 }
1528
1529 #[tokio::test]
1533 async fn a_keyed_send_is_retried_on_an_ambiguous_timeout() {
1534 let inner = FailNTransport::new(2, serde_json::json!({"ok": true}));
1535 let call_count = Arc::clone(&inner.call_count);
1536 let transport = RetryTransport::new(
1537 Box::new(inner),
1538 RetryPolicy::default()
1539 .with_initial_backoff(Duration::from_millis(1))
1540 .with_max_retries(3),
1541 )
1542 .with_peer_honouring_idempotency(true);
1543 let headers = HashMap::new();
1544 let result = transport
1545 .send_request("SendMessage", keyed_send_params(GOOD_KEY), &headers)
1546 .await;
1547
1548 assert!(
1549 result.is_ok(),
1550 "the retry should have succeeded: {result:?}"
1551 );
1552 assert_eq!(
1553 call_count.load(Ordering::SeqCst),
1554 3,
1555 "two ambiguous failures then a success"
1556 );
1557 }
1558
1559 #[tokio::test]
1562 async fn an_unkeyed_send_is_still_not_retried_on_a_timeout() {
1563 let inner = FailNTransport::new(2, serde_json::json!({"ok": true}));
1564 let call_count = Arc::clone(&inner.call_count);
1565 let transport = RetryTransport::new(
1566 Box::new(inner),
1567 RetryPolicy::default()
1568 .with_initial_backoff(Duration::from_millis(1))
1569 .with_max_retries(3),
1570 );
1571 let headers = HashMap::new();
1572 let bare = serde_json::json!({"message": {"id": "m", "role": "user", "parts": []}});
1573 let result = transport.send_request("SendMessage", bare, &headers).await;
1574
1575 assert!(result.is_err());
1576 assert_eq!(call_count.load(Ordering::SeqCst), 1);
1577 }
1578
1579 #[tokio::test]
1582 async fn non_idempotent_retried_on_503() {
1583 struct Fail503 {
1585 remaining: Arc<AtomicUsize>,
1586 calls: Arc<AtomicUsize>,
1587 }
1588 impl crate::transport::Transport for Fail503 {
1589 fn send_request<'a>(
1590 &'a self,
1591 _m: &'a str,
1592 _p: serde_json::Value,
1593 _h: &'a HashMap<String, String>,
1594 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1595 {
1596 self.calls.fetch_add(1, Ordering::SeqCst);
1597 let left = self.remaining.fetch_sub(1, Ordering::SeqCst);
1598 Box::pin(async move {
1599 if left > 0 {
1600 Err(ClientError::UnexpectedStatus {
1601 status: 503,
1602 body: String::new(),
1603 retry_after: None,
1604 })
1605 } else {
1606 Ok(serde_json::json!({"ok": true}))
1607 }
1608 })
1609 }
1610 fn send_streaming_request<'a>(
1611 &'a self,
1612 _m: &'a str,
1613 _p: serde_json::Value,
1614 _h: &'a HashMap<String, String>,
1615 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1616 Box::pin(async { Err(ClientError::Transport("n/a".into())) })
1617 }
1618 }
1619 let calls = Arc::new(AtomicUsize::new(0));
1620 let inner = Fail503 {
1621 remaining: Arc::new(AtomicUsize::new(1)),
1622 calls: Arc::clone(&calls),
1623 };
1624 let transport = RetryTransport::new(
1625 Box::new(inner),
1626 RetryPolicy::default()
1627 .with_initial_backoff(Duration::from_millis(1))
1628 .with_max_retries(3),
1629 );
1630 let headers = HashMap::new();
1631 let result = transport
1632 .send_request("SendMessage", serde_json::Value::Null, &headers)
1633 .await;
1634 assert!(result.is_ok(), "503 is a safe retry for non-idempotent");
1635 assert_eq!(calls.load(Ordering::SeqCst), 2);
1636 }
1637
1638 #[tokio::test(start_paused = true)]
1640 async fn retry_after_header_is_honored() {
1641 struct RetryAfter429 {
1643 calls: Arc<AtomicUsize>,
1644 }
1645 impl crate::transport::Transport for RetryAfter429 {
1646 fn send_request<'a>(
1647 &'a self,
1648 _m: &'a str,
1649 _p: serde_json::Value,
1650 _h: &'a HashMap<String, String>,
1651 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>>
1652 {
1653 let n = self.calls.fetch_add(1, Ordering::SeqCst);
1654 Box::pin(async move {
1655 if n == 0 {
1656 Err(ClientError::UnexpectedStatus {
1657 status: 429,
1658 body: String::new(),
1659 retry_after: Some(Duration::from_secs(20)),
1660 })
1661 } else {
1662 Ok(serde_json::json!({"ok": true}))
1663 }
1664 })
1665 }
1666 fn send_streaming_request<'a>(
1667 &'a self,
1668 _m: &'a str,
1669 _p: serde_json::Value,
1670 _h: &'a HashMap<String, String>,
1671 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
1672 Box::pin(async { Err(ClientError::Transport("n/a".into())) })
1673 }
1674 }
1675 let calls = Arc::new(AtomicUsize::new(0));
1676 let transport = RetryTransport::new(
1677 Box::new(RetryAfter429 {
1678 calls: Arc::clone(&calls),
1679 }),
1680 RetryPolicy::default()
1683 .with_initial_backoff(Duration::from_millis(1))
1684 .with_max_backoff(Duration::from_secs(30)),
1685 );
1686 let start = tokio::time::Instant::now();
1687 let headers = HashMap::new();
1688 let result = transport
1689 .send_request("GetTask", serde_json::Value::Null, &headers)
1690 .await;
1691 assert!(result.is_ok());
1692 assert!(
1693 start.elapsed() >= Duration::from_secs(20),
1694 "should have waited the server-requested 20s, waited {:?}",
1695 start.elapsed()
1696 );
1697 }
1698
1699 #[test]
1700 fn retry_after_clamped_to_max_backoff() {
1701 let err = ClientError::UnexpectedStatus {
1702 status: 503,
1703 body: String::new(),
1704 retry_after: Some(Duration::from_secs(9999)),
1705 };
1706 let delay = retry_delay(Some(&err), Duration::from_secs(1), Duration::from_secs(30));
1707 assert_eq!(
1708 delay,
1709 Duration::from_secs(30),
1710 "retry-after must be clamped"
1711 );
1712 }
1713 #[tokio::test]
1722 async fn a_keyed_send_is_not_retried_against_a_peer_that_never_advertised_it() {
1723 let inner = FailNTransport::new(2, serde_json::json!({"ok": true}));
1724 let call_count = Arc::clone(&inner.call_count);
1725 let transport = RetryTransport::new(
1728 Box::new(inner),
1729 RetryPolicy::default()
1730 .with_initial_backoff(Duration::from_millis(1))
1731 .with_max_retries(3),
1732 );
1733 let headers = HashMap::new();
1734 let result = transport
1735 .send_request("SendMessage", keyed_send_params(GOOD_KEY), &headers)
1736 .await;
1737
1738 assert!(result.is_err(), "the ambiguous failure must surface");
1739 assert_eq!(
1740 call_count.load(Ordering::SeqCst),
1741 1,
1742 "a key the peer never claimed to honour must not make the send retryable"
1743 );
1744 }
1745
1746 #[test]
1749 fn the_predicate_requires_both_the_key_and_the_advertisement() {
1750 assert!(
1751 carries_idempotency_key("SendMessage", &keyed_send_params(GOOD_KEY), true),
1752 "a valid key against an advertising peer is retryable"
1753 );
1754 assert!(
1755 !carries_idempotency_key("SendMessage", &keyed_send_params(GOOD_KEY), false),
1756 "the identical send against a peer that did not advertise is not"
1757 );
1758 }
1759}