1use crate::error::AgentLoopError;
18use rand::RngExt;
19use std::future::Future;
20use std::time::Duration;
21
22const MAX_RETRY_AFTER_SECS: u64 = 60;
25
26#[derive(Debug, Clone)]
35pub struct LlmRetryConfig {
36 pub max_retries: u32,
38 pub initial_backoff: Duration,
40 pub max_backoff: Duration,
42 pub backoff_multiplier: f64,
44 pub jitter_factor: f64,
47 pub max_retry_elapsed: Duration,
50}
51
52impl Default for LlmRetryConfig {
53 fn default() -> Self {
54 Self {
56 max_retries: 2,
57 initial_backoff: Duration::from_secs(1),
58 max_backoff: Duration::from_secs(60),
59 backoff_multiplier: 2.0,
60 jitter_factor: 0.25,
61 max_retry_elapsed: Duration::from_secs(30),
62 }
63 }
64}
65
66impl LlmRetryConfig {
67 pub fn no_retry() -> Self {
69 Self {
70 max_retries: 0,
71 ..Default::default()
72 }
73 }
74
75 pub fn aggressive() -> Self {
77 Self {
78 max_retries: 5,
79 initial_backoff: Duration::from_millis(500),
80 max_backoff: Duration::from_secs(120),
81 backoff_multiplier: 2.0,
82 jitter_factor: 0.25,
83 max_retry_elapsed: Duration::from_secs(120),
84 }
85 }
86
87 pub fn calculate_backoff(&self, attempt: u32) -> Duration {
89 let base_backoff =
90 self.initial_backoff.as_secs_f64() * self.backoff_multiplier.powi(attempt as i32);
91 let capped_backoff = base_backoff.min(self.max_backoff.as_secs_f64());
92
93 let jitter = if self.jitter_factor > 0.0 {
97 let jitter_range = capped_backoff * self.jitter_factor;
98 let jitter_offset = rand::rng().random::<f64>() * 2.0 - 1.0;
102 jitter_range * jitter_offset
103 } else {
104 0.0
105 };
106
107 Duration::from_secs_f64((capped_backoff + jitter).max(0.0))
108 }
109}
110
111pub fn reserve_retry_wait(
118 config: &LlmRetryConfig,
119 started_at: &mut Option<tokio::time::Instant>,
120 wait: Duration,
121) -> Option<Duration> {
122 let started = *started_at.get_or_insert_with(tokio::time::Instant::now);
123 let remaining = config.max_retry_elapsed.checked_sub(started.elapsed())?;
124 (wait < remaining).then_some(wait)
125}
126
127pub fn remaining_retry_time(
129 config: &LlmRetryConfig,
130 started_at: Option<tokio::time::Instant>,
131) -> Option<Duration> {
132 started_at.and_then(|started| config.max_retry_elapsed.checked_sub(started.elapsed()))
133}
134
135#[derive(Debug, Clone, Default)]
137pub struct RateLimitInfo {
138 pub retry_after_secs: Option<u64>,
140 pub requests_remaining: Option<u32>,
142 pub tokens_remaining: Option<u32>,
144 pub requests_reset: Option<String>,
146 pub tokens_reset: Option<String>,
148 pub limit_type: Option<RateLimitType>,
150}
151
152#[derive(Debug, Clone, PartialEq, Eq)]
154pub enum RateLimitType {
155 Requests,
157 InputTokens,
159 OutputTokens,
161 TotalTokens,
163 Unknown,
165}
166
167impl RateLimitInfo {
168 pub fn recommended_wait(&self, config: &LlmRetryConfig, attempt: u32) -> Duration {
171 if let Some(retry_after) = self.retry_after_secs {
172 if retry_after > 0 && retry_after <= MAX_RETRY_AFTER_SECS {
175 return Duration::from_secs(retry_after);
176 }
177 }
178 config.calculate_backoff(attempt)
180 }
181
182 pub fn from_anthropic_headers(headers: &reqwest::header::HeaderMap) -> Self {
184 let mut info = Self::default();
185
186 if let Some(val) = headers.get("retry-after-ms")
189 && let Ok(s) = val.to_str()
190 && let Ok(ms) = s.parse::<u64>()
191 {
192 info.retry_after_secs = Some(ms.div_ceil(1000));
194 }
195
196 if info.retry_after_secs.is_none()
198 && let Some(val) = headers.get("retry-after")
199 && let Ok(s) = val.to_str()
200 {
201 info.retry_after_secs = s.parse().ok();
202 }
203
204 if let Some(val) = headers.get("anthropic-ratelimit-requests-remaining")
206 && let Ok(s) = val.to_str()
207 {
208 info.requests_remaining = s.parse().ok();
209 }
210
211 if let Some(val) = headers.get("anthropic-ratelimit-tokens-remaining")
213 && let Ok(s) = val.to_str()
214 {
215 info.tokens_remaining = s.parse().ok();
216 }
217
218 if let Some(val) = headers.get("anthropic-ratelimit-requests-reset")
220 && let Ok(s) = val.to_str()
221 {
222 info.requests_reset = Some(s.to_string());
223 }
224
225 if let Some(val) = headers.get("anthropic-ratelimit-tokens-reset")
227 && let Ok(s) = val.to_str()
228 {
229 info.tokens_reset = Some(s.to_string());
230 }
231
232 if info.requests_remaining == Some(0) {
234 info.limit_type = Some(RateLimitType::Requests);
235 } else if info.tokens_remaining == Some(0) {
236 info.limit_type = Some(RateLimitType::InputTokens);
237 }
238
239 info
240 }
241
242 pub fn from_openai_headers(headers: &reqwest::header::HeaderMap) -> Self {
244 let mut info = Self::default();
245
246 if let Some(val) = headers.get("retry-after-ms")
248 && let Ok(s) = val.to_str()
249 && let Ok(ms) = s.parse::<u64>()
250 {
251 info.retry_after_secs = Some(ms.div_ceil(1000));
253 }
254
255 if info.retry_after_secs.is_none()
257 && let Some(val) = headers.get("retry-after")
258 && let Ok(s) = val.to_str()
259 {
260 info.retry_after_secs = s.parse().ok();
261 }
262
263 if let Some(val) = headers.get("x-ratelimit-remaining-requests")
265 && let Ok(s) = val.to_str()
266 {
267 info.requests_remaining = s.parse().ok();
268 }
269
270 if let Some(val) = headers.get("x-ratelimit-remaining-tokens")
272 && let Ok(s) = val.to_str()
273 {
274 let val: i64 = s.parse().unwrap_or(-1);
276 if val >= 0 {
277 info.tokens_remaining = Some(val as u32);
278 }
279 }
280
281 if let Some(val) = headers.get("x-ratelimit-reset-requests")
283 && let Ok(s) = val.to_str()
284 {
285 info.requests_reset = Some(s.to_string());
286 if info.retry_after_secs.is_none() {
288 info.retry_after_secs = parse_duration_string(s);
289 }
290 }
291
292 if let Some(val) = headers.get("x-ratelimit-reset-tokens")
294 && let Ok(s) = val.to_str()
295 {
296 info.tokens_reset = Some(s.to_string());
297 }
298
299 if info.requests_remaining == Some(0) {
301 info.limit_type = Some(RateLimitType::Requests);
302 } else if info.tokens_remaining == Some(0) {
303 info.limit_type = Some(RateLimitType::TotalTokens);
304 }
305
306 info
307 }
308}
309
310fn parse_duration_string(s: &str) -> Option<u64> {
312 let s = s.trim();
313 if s.is_empty() {
314 return None;
315 }
316
317 let mut total_secs: u64 = 0;
318 let mut current_num = String::new();
319
320 for c in s.chars() {
321 if c.is_ascii_digit() {
322 current_num.push(c);
323 } else {
324 let num: u64 = current_num.parse().ok()?;
325 current_num.clear();
326
327 match c {
328 'h' => total_secs += num * 3600,
329 'm' => total_secs += num * 60,
330 's' => total_secs += num,
331 _ => return None,
332 }
333 }
334 }
335
336 if total_secs > 0 {
337 Some(total_secs)
338 } else {
339 None
340 }
341}
342
343#[derive(Debug, Clone, Default)]
345pub struct RetryMetadata {
346 pub attempts: u32,
348 pub total_retry_wait: Duration,
350 pub total_retry_elapsed: Duration,
352 pub last_rate_limit_info: Option<RateLimitInfo>,
354}
355
356impl RetryMetadata {
357 pub fn had_retries(&self) -> bool {
359 self.attempts > 0
360 }
361
362 pub fn first_attempt_success() -> Self {
364 Self::default()
365 }
366
367 pub fn record_retry(
369 &mut self,
370 wait_duration: Duration,
371 rate_limit_info: Option<RateLimitInfo>,
372 ) {
373 self.attempts += 1;
374 self.total_retry_wait += wait_duration;
375 if rate_limit_info.is_some() {
376 self.last_rate_limit_info = rate_limit_info;
377 }
378 }
379
380 pub fn absorb(&mut self, other: RetryMetadata) {
382 self.attempts = self.attempts.saturating_add(other.attempts);
383 self.total_retry_wait = self.total_retry_wait.saturating_add(other.total_retry_wait);
384 self.total_retry_elapsed = self
385 .total_retry_elapsed
386 .saturating_add(other.total_retry_elapsed);
387 if other.last_rate_limit_info.is_some() {
388 self.last_rate_limit_info = other.last_rate_limit_info;
389 }
390 }
391}
392
393pub fn is_rate_limit_status(status: reqwest::StatusCode) -> bool {
395 status == reqwest::StatusCode::TOO_MANY_REQUESTS
396}
397
398pub fn is_transient_error(status: reqwest::StatusCode) -> bool {
406 if status == reqwest::StatusCode::REQUEST_TIMEOUT {
408 return true;
409 }
410 if status == reqwest::StatusCode::CONFLICT {
412 return true;
413 }
414 if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
416 return true;
417 }
418 if status.is_server_error() && status != reqwest::StatusCode::NOT_IMPLEMENTED {
420 return true;
421 }
422 false
423}
424
425pub fn is_transient_send_error(err: &reqwest::Error) -> bool {
441 err.is_connect() || err.is_timeout() || err.is_request()
442}
443
444pub fn send_error_message(err: &reqwest::Error, attempts: u32) -> String {
447 if attempts > 0 {
448 format!("Failed to send request: {err} (after {attempts} retries)")
449 } else {
450 format!("Failed to send request: {err}")
451 }
452}
453
454pub fn is_transient_error_message(message: &str) -> bool {
459 if crate::user_facing_error::is_usage_limit_message(message) {
465 return false;
466 }
467
468 let msg = message.trim().to_ascii_lowercase();
469
470 if msg.contains("provider stream stall") {
476 return true;
477 }
478
479 [
480 "server_error",
481 "internal server error",
482 "overloaded",
483 "overloaded_error",
484 "rate limit",
485 "too many requests",
486 "request timeout",
487 "timed out",
488 "service unavailable",
489 "bad gateway",
490 "gateway timeout",
491 "temporarily unavailable",
492 ]
493 .iter()
494 .any(|needle| msg.contains(needle))
495}
496
497pub fn is_transient_stream_error(error: &crate::driver_registry::LlmStreamError) -> bool {
503 if let Some(code) = error.code.as_deref()
504 && let Some(kind) = crate::error::LlmErrorKind::from_provider_code(code)
505 {
506 return matches!(
507 kind,
508 crate::error::LlmErrorKind::RateLimited | crate::error::LlmErrorKind::Unavailable
509 );
510 }
511
512 if let Some(status) = error
513 .status
514 .and_then(|status| reqwest::StatusCode::from_u16(status).ok())
515 {
516 return is_transient_error(status);
517 }
518
519 is_transient_error_message(&error.message)
520}
521
522pub enum SendOutcome {
532 Send(reqwest::Error),
535 Fatal(AgentLoopError),
538}
539
540pub enum RetryDecision {
543 Retry {
546 wait: Duration,
547 rate_limit_info: Option<RateLimitInfo>,
548 },
549 RetryNow,
554 Terminal(AgentLoopError),
556}
557
558pub async fn retry_request<S, SFut, C, CFut, E>(
578 config: &LlmRetryConfig,
579 driver_name: &str,
580 mut send: S,
581 mut classify: C,
582 send_error: E,
583) -> Result<(reqwest::Response, RetryMetadata), AgentLoopError>
584where
585 S: FnMut() -> SFut,
586 SFut: Future<Output = Result<reqwest::Response, SendOutcome>>,
587 C: FnMut(reqwest::Response, u32, bool) -> CFut,
588 CFut: Future<Output = RetryDecision>,
589 E: Fn(&reqwest::Error, u32) -> AgentLoopError,
590{
591 let mut retry_metadata = RetryMetadata::default();
592 let mut retry_started_at = None;
593
594 let response = loop {
595 let send_result = if let Some(remaining) = remaining_retry_time(config, retry_started_at) {
596 match tokio::time::timeout(remaining, send()).await {
597 Ok(result) => result,
598 Err(_) => {
599 return Err(AgentLoopError::llm_kind(
600 crate::error::LlmErrorKind::Unavailable,
601 format!(
602 "{driver_name} retry time budget exhausted after {} retries over {:.1}s",
603 retry_metadata.attempts,
604 config.max_retry_elapsed.as_secs_f64()
605 ),
606 )
607 .with_retry_metadata(&retry_metadata));
608 }
609 }
610 } else {
611 send().await
612 };
613 let response = match send_result {
614 Ok(response) => response,
615 Err(SendOutcome::Fatal(err)) => return Err(err),
616 Err(SendOutcome::Send(e)) => {
617 if is_transient_send_error(&e) && retry_metadata.attempts < config.max_retries {
622 let proposed_wait = config.calculate_backoff(retry_metadata.attempts);
623 let Some(wait_duration) =
624 reserve_retry_wait(config, &mut retry_started_at, proposed_wait)
625 else {
626 return Err(send_error(&e, retry_metadata.attempts)
627 .with_retry_metadata(&retry_metadata));
628 };
629 tracing::warn!(
630 error = %e,
631 driver = driver_name,
632 attempt = retry_metadata.attempts + 1,
633 max_retries = config.max_retries,
634 wait_secs = wait_duration.as_secs_f64(),
635 "transient connection error sending request, retrying"
636 );
637 retry_metadata.record_retry(wait_duration, None);
638 tokio::time::sleep(wait_duration).await;
639 continue;
640 }
641 return Err(
642 send_error(&e, retry_metadata.attempts).with_retry_metadata(&retry_metadata)
643 );
644 }
645 };
646
647 let status = response.status();
648 if status.is_success() {
649 break response;
650 }
651
652 let can_retry = is_transient_error(status) && retry_metadata.attempts < config.max_retries;
653 match classify(response, retry_metadata.attempts, can_retry).await {
654 RetryDecision::Retry {
655 wait,
656 rate_limit_info,
657 } => {
658 let Some(wait) = reserve_retry_wait(config, &mut retry_started_at, wait) else {
659 return Err(AgentLoopError::llm_kind(
660 crate::error::LlmErrorKind::Unavailable,
661 format!(
662 "{driver_name} retry time budget exhausted after {} retries over {:.1}s",
663 retry_metadata.attempts,
664 config.max_retry_elapsed.as_secs_f64()
665 ),
666 )
667 .with_retry_metadata(&retry_metadata));
668 };
669 tracing::warn!(
670 status = %status,
671 driver = driver_name,
672 attempt = retry_metadata.attempts + 1,
673 max_retries = config.max_retries,
674 wait_secs = wait.as_secs_f64(),
675 "rate limit or transient error, retrying"
676 );
677 retry_metadata.record_retry(wait, rate_limit_info);
678 tokio::time::sleep(wait).await;
679 continue;
680 }
681 RetryDecision::RetryNow => continue,
682 RetryDecision::Terminal(err) => {
683 return Err(err.with_retry_metadata(&retry_metadata));
684 }
685 }
686 };
687
688 if retry_metadata.had_retries() {
689 retry_metadata.total_retry_elapsed = retry_started_at
690 .map(|started| started.elapsed())
691 .unwrap_or_default();
692 tracing::info!(
693 driver = driver_name,
694 attempts = retry_metadata.attempts,
695 total_wait_secs = retry_metadata.total_retry_wait.as_secs_f64(),
696 "request succeeded after retries"
697 );
698 }
699
700 Ok((response, retry_metadata))
701}
702
703#[cfg(test)]
708mod tests {
709 use super::*;
710
711 #[test]
712 fn test_default_config_matches_official_sdks() {
713 let config = LlmRetryConfig::default();
715 assert_eq!(config.max_retries, 2); assert_eq!(config.initial_backoff, Duration::from_secs(1));
717 assert_eq!(config.max_backoff, Duration::from_secs(60));
718 assert_eq!(config.backoff_multiplier, 2.0);
719 assert!((config.jitter_factor - 0.25).abs() < 0.001); }
721
722 #[test]
723 fn test_calculate_backoff_exponential() {
724 let config = LlmRetryConfig {
725 initial_backoff: Duration::from_secs(1),
726 max_backoff: Duration::from_secs(60),
727 backoff_multiplier: 2.0,
728 jitter_factor: 0.0, ..Default::default()
730 };
731
732 assert_eq!(config.calculate_backoff(0), Duration::from_secs(1));
734 assert_eq!(config.calculate_backoff(1), Duration::from_secs(2));
736 assert_eq!(config.calculate_backoff(2), Duration::from_secs(4));
738 assert_eq!(config.calculate_backoff(3), Duration::from_secs(8));
740 }
741
742 #[test]
743 fn test_calculate_backoff_capped() {
744 let config = LlmRetryConfig {
745 initial_backoff: Duration::from_secs(10),
746 max_backoff: Duration::from_secs(30),
747 backoff_multiplier: 2.0,
748 jitter_factor: 0.0,
749 ..Default::default()
750 };
751
752 assert_eq!(config.calculate_backoff(0), Duration::from_secs(10));
754 assert_eq!(config.calculate_backoff(1), Duration::from_secs(20));
756 assert_eq!(config.calculate_backoff(2), Duration::from_secs(30));
758 assert_eq!(config.calculate_backoff(3), Duration::from_secs(30));
760 }
761
762 #[test]
765 fn test_backoff_jitter_is_randomized() {
766 let config = LlmRetryConfig {
767 initial_backoff: Duration::from_secs(10),
768 max_backoff: Duration::from_secs(60),
769 backoff_multiplier: 2.0,
770 jitter_factor: 0.25,
771 ..Default::default()
772 };
773 let samples: std::collections::HashSet<u128> = (0..20)
774 .map(|_| config.calculate_backoff(1).as_nanos())
775 .collect();
776 assert!(
777 samples.len() > 1,
778 "jittered backoff should vary across calls, got {} distinct value(s)",
779 samples.len()
780 );
781 for _ in 0..50 {
783 let secs = config.calculate_backoff(1).as_secs_f64();
784 assert!(
785 (15.0..=25.0).contains(&secs),
786 "backoff {secs}s out of range"
787 );
788 }
789 }
790
791 #[test]
792 fn test_parse_duration_string() {
793 assert_eq!(parse_duration_string("1s"), Some(1));
794 assert_eq!(parse_duration_string("30s"), Some(30));
795 assert_eq!(parse_duration_string("1m"), Some(60));
796 assert_eq!(parse_duration_string("6m0s"), Some(360));
797 assert_eq!(parse_duration_string("1h"), Some(3600));
798 assert_eq!(parse_duration_string("1h30m"), Some(5400));
799 assert_eq!(parse_duration_string("1h30m45s"), Some(5445));
800 assert_eq!(parse_duration_string(""), None);
801 assert_eq!(parse_duration_string("invalid"), None);
802 }
803
804 #[test]
805 fn test_rate_limit_info_recommended_wait_with_retry_after() {
806 let config = LlmRetryConfig::default();
807 let info = RateLimitInfo {
808 retry_after_secs: Some(10),
809 ..Default::default()
810 };
811
812 assert_eq!(info.recommended_wait(&config, 0), Duration::from_secs(10));
814 assert_eq!(info.recommended_wait(&config, 5), Duration::from_secs(10));
815 }
816
817 #[test]
818 fn test_rate_limit_info_recommended_wait_capped_at_60s() {
819 let config = LlmRetryConfig {
821 jitter_factor: 0.0, ..Default::default()
823 };
824 let info = RateLimitInfo {
825 retry_after_secs: Some(120), ..Default::default()
827 };
828
829 assert_eq!(info.recommended_wait(&config, 0), Duration::from_secs(1));
831 }
832
833 #[test]
834 fn test_rate_limit_info_recommended_wait_fallback() {
835 let config = LlmRetryConfig {
836 initial_backoff: Duration::from_secs(1),
837 backoff_multiplier: 2.0,
838 jitter_factor: 0.0,
839 ..Default::default()
840 };
841 let info = RateLimitInfo::default(); assert_eq!(info.recommended_wait(&config, 0), Duration::from_secs(1));
845 assert_eq!(info.recommended_wait(&config, 1), Duration::from_secs(2));
846 }
847
848 #[test]
849 fn test_retry_metadata_record() {
850 let mut meta = RetryMetadata::default();
851 assert!(!meta.had_retries());
852 assert_eq!(meta.attempts, 0);
853
854 meta.record_retry(Duration::from_secs(1), None);
855 assert!(meta.had_retries());
856 assert_eq!(meta.attempts, 1);
857 assert_eq!(meta.total_retry_wait, Duration::from_secs(1));
858
859 meta.record_retry(Duration::from_secs(2), None);
860 assert_eq!(meta.attempts, 2);
861 assert_eq!(meta.total_retry_wait, Duration::from_secs(3));
862 }
863
864 #[test]
865 fn test_is_transient_error_matches_official_sdks() {
866 assert!(is_transient_error(reqwest::StatusCode::REQUEST_TIMEOUT)); assert!(is_transient_error(reqwest::StatusCode::CONFLICT)); assert!(is_transient_error(reqwest::StatusCode::TOO_MANY_REQUESTS)); assert!(is_transient_error(
871 reqwest::StatusCode::INTERNAL_SERVER_ERROR
872 )); assert!(is_transient_error(reqwest::StatusCode::BAD_GATEWAY)); assert!(is_transient_error(reqwest::StatusCode::SERVICE_UNAVAILABLE)); assert!(is_transient_error(reqwest::StatusCode::GATEWAY_TIMEOUT)); assert!(!is_transient_error(reqwest::StatusCode::OK));
879 assert!(!is_transient_error(reqwest::StatusCode::BAD_REQUEST)); assert!(!is_transient_error(reqwest::StatusCode::UNAUTHORIZED)); assert!(!is_transient_error(reqwest::StatusCode::FORBIDDEN)); assert!(!is_transient_error(reqwest::StatusCode::NOT_FOUND)); assert!(!is_transient_error(reqwest::StatusCode::NOT_IMPLEMENTED)); }
885
886 #[tokio::test]
890 async fn test_is_transient_send_error_on_connection_refused() {
891 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
894 let addr = listener.local_addr().unwrap();
895 drop(listener);
896
897 let err = reqwest::Client::new()
898 .get(format!("http://{addr}/"))
899 .send()
900 .await
901 .expect_err("request to a closed port should fail");
902
903 assert!(
904 is_transient_send_error(&err),
905 "connection-refused send error should be transient: {err:?}"
906 );
907 }
908
909 #[test]
910 fn test_is_transient_error_message_detects_provider_server_errors() {
911 assert!(is_transient_error_message(
912 "server_error: An error occurred while processing your request."
913 ));
914 assert!(is_transient_error_message("Rate limit exceeded"));
915 assert!(is_transient_error_message(
916 "Service temporarily unavailable"
917 ));
918 }
919
920 #[test]
921 fn test_is_transient_error_message_rejects_non_retryable_messages() {
922 assert!(!is_transient_error_message(
923 "invalid_request_error: bad tool schema"
924 ));
925 assert!(!is_transient_error_message("Model not available: gpt-99"));
926 }
927
928 #[test]
929 fn structured_stream_error_prefers_code_and_status_over_message() {
930 use crate::driver_registry::LlmStreamError;
931
932 assert!(is_transient_stream_error(&LlmStreamError::provider(
933 Some("processing_error"),
934 None,
935 "An error occurred while processing your request.",
936 )));
937 assert!(is_transient_stream_error(&LlmStreamError::provider(
938 None::<String>,
939 Some(503),
940 "opaque failure",
941 )));
942 assert!(!is_transient_stream_error(&LlmStreamError::provider(
943 Some("invalid_request_error"),
944 Some(503),
945 "server unavailable",
946 )));
947 assert!(!is_transient_stream_error(&LlmStreamError::provider(
948 Some("insufficient_quota"),
949 Some(429),
950 "rate limit",
951 )));
952 }
953
954 #[test]
955 fn test_provider_stream_stall_is_transient() {
956 assert!(is_transient_error_message(
959 "provider stream stall: no tokens for 120s"
960 ));
961 use crate::driver_registry::LlmStreamError;
964 assert!(is_transient_stream_error(&LlmStreamError::new(
965 "provider stream stall: no tokens for 120s"
966 )));
967 }
968
969 #[test]
970 fn test_is_transient_error_message_treats_usage_limit_as_non_transient() {
971 assert!(!is_transient_error_message(
975 "Codex API error (429 Too Many Requests): {\"error\":{\"type\":\"usage_limit_reached\",\"resets_at\":1783767823}}"
976 ));
977 }
978
979 #[test]
980 fn test_max_retry_after_constant() {
981 assert_eq!(MAX_RETRY_AFTER_SECS, 60);
983 }
984
985 fn fake_response(status: u16, body: &str) -> reqwest::Response {
991 let http_response = http::Response::builder()
992 .status(status)
993 .body(body.to_string())
994 .unwrap();
995 reqwest::Response::from(http_response)
996 }
997
998 fn fast_config(max_retries: u32) -> LlmRetryConfig {
1000 LlmRetryConfig {
1001 max_retries,
1002 initial_backoff: Duration::from_millis(0),
1003 max_backoff: Duration::from_millis(0),
1004 backoff_multiplier: 1.0,
1005 jitter_factor: 0.0,
1006 ..Default::default()
1007 }
1008 }
1009
1010 #[tokio::test]
1011 async fn test_retry_request_success_first_try() {
1012 let config = fast_config(2);
1013 let (resp, meta) = retry_request(
1014 &config,
1015 "TestDriver",
1016 || async { Ok(fake_response(200, "ok")) },
1017 |_resp, _attempt, _can_retry| async {
1018 RetryDecision::Terminal(AgentLoopError::llm("unreachable"))
1019 },
1020 |e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
1021 )
1022 .await
1023 .expect("should succeed");
1024 assert!(resp.status().is_success());
1025 assert_eq!(meta.attempts, 0);
1026 assert!(!meta.had_retries());
1027 }
1028
1029 #[tokio::test]
1030 async fn test_retry_request_retries_then_succeeds() {
1031 let config = fast_config(3);
1032 let calls = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
1033 let calls_send = calls.clone();
1034 let (resp, meta) = retry_request(
1035 &config,
1036 "TestDriver",
1037 move || {
1038 let calls = calls_send.clone();
1039 async move {
1040 let n = calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1041 if n < 2 {
1043 Ok(fake_response(429, "rate limited"))
1044 } else {
1045 Ok(fake_response(200, "ok"))
1046 }
1047 }
1048 },
1049 |_resp, _attempt, can_retry| async move {
1050 assert!(can_retry, "429 within budget should be retryable");
1051 RetryDecision::Retry {
1052 wait: Duration::from_millis(0),
1053 rate_limit_info: None,
1054 }
1055 },
1056 |e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
1057 )
1058 .await
1059 .expect("should eventually succeed");
1060 assert!(resp.status().is_success());
1061 assert_eq!(meta.attempts, 2);
1062 }
1063
1064 #[tokio::test]
1065 async fn test_retry_request_terminal_decision_propagates() {
1066 let config = fast_config(2);
1067 let result = retry_request(
1068 &config,
1069 "TestDriver",
1070 || async { Ok(fake_response(400, "bad request")) },
1071 |_resp, _attempt, _can_retry| async {
1072 RetryDecision::Terminal(AgentLoopError::llm("classified terminal"))
1073 },
1074 |e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
1075 )
1076 .await;
1077 let err = result.expect_err("terminal decision should error");
1078 assert!(err.to_string().contains("classified terminal"));
1079 }
1080
1081 #[tokio::test]
1082 async fn test_retry_request_retry_now_does_not_count_attempt() {
1083 let config = fast_config(2);
1084 let calls = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
1085 let calls_send = calls.clone();
1086 let (resp, meta) = retry_request(
1087 &config,
1088 "TestDriver",
1089 move || {
1090 let calls = calls_send.clone();
1091 async move {
1092 let n = calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1093 if n == 0 {
1094 Ok(fake_response(400, "max_tokens too large"))
1095 } else {
1096 Ok(fake_response(200, "ok"))
1097 }
1098 }
1099 },
1100 {
1101 let mut used_fallback = false;
1102 move |_resp, _attempt, _can_retry| {
1103 let do_fallback = !used_fallback;
1104 used_fallback = true;
1105 async move {
1106 if do_fallback {
1107 RetryDecision::RetryNow
1108 } else {
1109 RetryDecision::Terminal(AgentLoopError::llm("unreachable"))
1110 }
1111 }
1112 }
1113 },
1114 |e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
1115 )
1116 .await
1117 .expect("RetryNow then success");
1118 assert!(resp.status().is_success());
1119 assert_eq!(meta.attempts, 0);
1121 }
1122
1123 #[tokio::test]
1124 async fn test_retry_request_send_error_exhausts() {
1125 let config = fast_config(1);
1128
1129 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
1131 let addr = listener.local_addr().unwrap();
1132 drop(listener);
1133 let make_err = || async {
1134 reqwest::Client::new()
1135 .get(format!("http://{addr}/"))
1136 .send()
1137 .await
1138 .expect_err("closed port")
1139 };
1140
1141 let result = retry_request(
1142 &config,
1143 "TestDriver",
1144 move || async move { Err(SendOutcome::Send(make_err().await)) },
1145 |_resp, _attempt, _can_retry| async {
1146 RetryDecision::Terminal(AgentLoopError::llm("unreachable"))
1147 },
1148 |e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
1149 )
1150 .await;
1151 let err = result.expect_err("send errors should exhaust to terminal");
1152 assert!(err.to_string().contains("after 1 retries"), "got: {err}");
1154 }
1155
1156 #[tokio::test]
1157 async fn test_retry_request_fatal_send_propagates_immediately() {
1158 let config = fast_config(3);
1159 let result = retry_request(
1160 &config,
1161 "TestDriver",
1162 || async { Err(SendOutcome::Fatal(AgentLoopError::llm("auth failed"))) },
1163 |_resp, _attempt, _can_retry| async {
1164 RetryDecision::Terminal(AgentLoopError::llm("unreachable"))
1165 },
1166 |e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
1167 )
1168 .await;
1169 let err = result.expect_err("fatal send should propagate");
1170 assert!(err.to_string().contains("auth failed"));
1171 }
1172}