use crate::error::AgentLoopError;
use rand::RngExt;
use std::future::Future;
use std::time::Duration;
const MAX_RETRY_AFTER_SECS: u64 = 60;
#[derive(Debug, Clone)]
pub struct LlmRetryConfig {
pub max_retries: u32,
pub initial_backoff: Duration,
pub max_backoff: Duration,
pub backoff_multiplier: f64,
pub jitter_factor: f64,
}
impl Default for LlmRetryConfig {
fn default() -> Self {
Self {
max_retries: 2,
initial_backoff: Duration::from_secs(1),
max_backoff: Duration::from_secs(60),
backoff_multiplier: 2.0,
jitter_factor: 0.25,
}
}
}
impl LlmRetryConfig {
pub fn no_retry() -> Self {
Self {
max_retries: 0,
..Default::default()
}
}
pub fn aggressive() -> Self {
Self {
max_retries: 5,
initial_backoff: Duration::from_millis(500),
max_backoff: Duration::from_secs(120),
backoff_multiplier: 2.0,
jitter_factor: 0.25,
}
}
pub fn calculate_backoff(&self, attempt: u32) -> Duration {
let base_backoff =
self.initial_backoff.as_secs_f64() * self.backoff_multiplier.powi(attempt as i32);
let capped_backoff = base_backoff.min(self.max_backoff.as_secs_f64());
let jitter = if self.jitter_factor > 0.0 {
let jitter_range = capped_backoff * self.jitter_factor;
let jitter_offset = rand::rng().random::<f64>() * 2.0 - 1.0;
jitter_range * jitter_offset
} else {
0.0
};
Duration::from_secs_f64((capped_backoff + jitter).max(0.0))
}
}
#[derive(Debug, Clone, Default)]
pub struct RateLimitInfo {
pub retry_after_secs: Option<u64>,
pub requests_remaining: Option<u32>,
pub tokens_remaining: Option<u32>,
pub requests_reset: Option<String>,
pub tokens_reset: Option<String>,
pub limit_type: Option<RateLimitType>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RateLimitType {
Requests,
InputTokens,
OutputTokens,
TotalTokens,
Unknown,
}
impl RateLimitInfo {
pub fn recommended_wait(&self, config: &LlmRetryConfig, attempt: u32) -> Duration {
if let Some(retry_after) = self.retry_after_secs {
if retry_after > 0 && retry_after <= MAX_RETRY_AFTER_SECS {
return Duration::from_secs(retry_after);
}
}
config.calculate_backoff(attempt)
}
pub fn from_anthropic_headers(headers: &reqwest::header::HeaderMap) -> Self {
let mut info = Self::default();
if let Some(val) = headers.get("retry-after-ms")
&& let Ok(s) = val.to_str()
&& let Ok(ms) = s.parse::<u64>()
{
info.retry_after_secs = Some(ms.div_ceil(1000));
}
if info.retry_after_secs.is_none()
&& let Some(val) = headers.get("retry-after")
&& let Ok(s) = val.to_str()
{
info.retry_after_secs = s.parse().ok();
}
if let Some(val) = headers.get("anthropic-ratelimit-requests-remaining")
&& let Ok(s) = val.to_str()
{
info.requests_remaining = s.parse().ok();
}
if let Some(val) = headers.get("anthropic-ratelimit-tokens-remaining")
&& let Ok(s) = val.to_str()
{
info.tokens_remaining = s.parse().ok();
}
if let Some(val) = headers.get("anthropic-ratelimit-requests-reset")
&& let Ok(s) = val.to_str()
{
info.requests_reset = Some(s.to_string());
}
if let Some(val) = headers.get("anthropic-ratelimit-tokens-reset")
&& let Ok(s) = val.to_str()
{
info.tokens_reset = Some(s.to_string());
}
if info.requests_remaining == Some(0) {
info.limit_type = Some(RateLimitType::Requests);
} else if info.tokens_remaining == Some(0) {
info.limit_type = Some(RateLimitType::InputTokens);
}
info
}
pub fn from_openai_headers(headers: &reqwest::header::HeaderMap) -> Self {
let mut info = Self::default();
if let Some(val) = headers.get("retry-after-ms")
&& let Ok(s) = val.to_str()
&& let Ok(ms) = s.parse::<u64>()
{
info.retry_after_secs = Some(ms.div_ceil(1000));
}
if info.retry_after_secs.is_none()
&& let Some(val) = headers.get("retry-after")
&& let Ok(s) = val.to_str()
{
info.retry_after_secs = s.parse().ok();
}
if let Some(val) = headers.get("x-ratelimit-remaining-requests")
&& let Ok(s) = val.to_str()
{
info.requests_remaining = s.parse().ok();
}
if let Some(val) = headers.get("x-ratelimit-remaining-tokens")
&& let Ok(s) = val.to_str()
{
let val: i64 = s.parse().unwrap_or(-1);
if val >= 0 {
info.tokens_remaining = Some(val as u32);
}
}
if let Some(val) = headers.get("x-ratelimit-reset-requests")
&& let Ok(s) = val.to_str()
{
info.requests_reset = Some(s.to_string());
if info.retry_after_secs.is_none() {
info.retry_after_secs = parse_duration_string(s);
}
}
if let Some(val) = headers.get("x-ratelimit-reset-tokens")
&& let Ok(s) = val.to_str()
{
info.tokens_reset = Some(s.to_string());
}
if info.requests_remaining == Some(0) {
info.limit_type = Some(RateLimitType::Requests);
} else if info.tokens_remaining == Some(0) {
info.limit_type = Some(RateLimitType::TotalTokens);
}
info
}
}
fn parse_duration_string(s: &str) -> Option<u64> {
let s = s.trim();
if s.is_empty() {
return None;
}
let mut total_secs: u64 = 0;
let mut current_num = String::new();
for c in s.chars() {
if c.is_ascii_digit() {
current_num.push(c);
} else {
let num: u64 = current_num.parse().ok()?;
current_num.clear();
match c {
'h' => total_secs += num * 3600,
'm' => total_secs += num * 60,
's' => total_secs += num,
_ => return None,
}
}
}
if total_secs > 0 {
Some(total_secs)
} else {
None
}
}
#[derive(Debug, Clone, Default)]
pub struct RetryMetadata {
pub attempts: u32,
pub total_retry_wait: Duration,
pub last_rate_limit_info: Option<RateLimitInfo>,
}
impl RetryMetadata {
pub fn had_retries(&self) -> bool {
self.attempts > 0
}
pub fn first_attempt_success() -> Self {
Self::default()
}
pub fn record_retry(
&mut self,
wait_duration: Duration,
rate_limit_info: Option<RateLimitInfo>,
) {
self.attempts += 1;
self.total_retry_wait += wait_duration;
if rate_limit_info.is_some() {
self.last_rate_limit_info = rate_limit_info;
}
}
}
pub fn is_rate_limit_status(status: reqwest::StatusCode) -> bool {
status == reqwest::StatusCode::TOO_MANY_REQUESTS
}
pub fn is_transient_error(status: reqwest::StatusCode) -> bool {
if status == reqwest::StatusCode::REQUEST_TIMEOUT {
return true;
}
if status == reqwest::StatusCode::CONFLICT {
return true;
}
if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
return true;
}
if status.is_server_error() && status != reqwest::StatusCode::NOT_IMPLEMENTED {
return true;
}
false
}
pub fn is_transient_send_error(err: &reqwest::Error) -> bool {
err.is_connect() || err.is_timeout() || err.is_request()
}
pub fn send_error_message(err: &reqwest::Error, attempts: u32) -> String {
if attempts > 0 {
format!("Failed to send request: {err} (after {attempts} retries)")
} else {
format!("Failed to send request: {err}")
}
}
pub fn is_transient_error_message(message: &str) -> bool {
let msg = message.trim().to_ascii_lowercase();
[
"server_error",
"internal server error",
"overloaded",
"overloaded_error",
"rate limit",
"too many requests",
"request timeout",
"timed out",
"service unavailable",
"bad gateway",
"gateway timeout",
"temporarily unavailable",
]
.iter()
.any(|needle| msg.contains(needle))
}
pub fn is_transient_stream_error(error: &crate::driver_registry::LlmStreamError) -> bool {
if let Some(code) = error.code.as_deref()
&& let Some(kind) = crate::error::LlmErrorKind::from_provider_code(code)
{
return matches!(
kind,
crate::error::LlmErrorKind::RateLimited | crate::error::LlmErrorKind::Unavailable
);
}
if let Some(status) = error
.status
.and_then(|status| reqwest::StatusCode::from_u16(status).ok())
{
return is_transient_error(status);
}
is_transient_error_message(&error.message)
}
pub enum SendOutcome {
Send(reqwest::Error),
Fatal(AgentLoopError),
}
pub enum RetryDecision {
Retry {
wait: Duration,
rate_limit_info: Option<RateLimitInfo>,
},
RetryNow,
Terminal(AgentLoopError),
}
pub async fn retry_request<S, SFut, C, CFut, E>(
config: &LlmRetryConfig,
driver_name: &str,
mut send: S,
mut classify: C,
send_error: E,
) -> Result<(reqwest::Response, RetryMetadata), AgentLoopError>
where
S: FnMut() -> SFut,
SFut: Future<Output = Result<reqwest::Response, SendOutcome>>,
C: FnMut(reqwest::Response, u32, bool) -> CFut,
CFut: Future<Output = RetryDecision>,
E: Fn(&reqwest::Error, u32) -> AgentLoopError,
{
let mut retry_metadata = RetryMetadata::default();
let response = loop {
let response = match send().await {
Ok(response) => response,
Err(SendOutcome::Fatal(err)) => return Err(err),
Err(SendOutcome::Send(e)) => {
if is_transient_send_error(&e) && retry_metadata.attempts < config.max_retries {
let wait_duration = config.calculate_backoff(retry_metadata.attempts);
tracing::warn!(
error = %e,
driver = driver_name,
attempt = retry_metadata.attempts + 1,
max_retries = config.max_retries,
wait_secs = wait_duration.as_secs_f64(),
"transient connection error sending request, retrying"
);
retry_metadata.record_retry(wait_duration, None);
tokio::time::sleep(wait_duration).await;
continue;
}
return Err(send_error(&e, retry_metadata.attempts));
}
};
let status = response.status();
if status.is_success() {
break response;
}
let can_retry = is_transient_error(status) && retry_metadata.attempts < config.max_retries;
match classify(response, retry_metadata.attempts, can_retry).await {
RetryDecision::Retry {
wait,
rate_limit_info,
} => {
tracing::warn!(
status = %status,
driver = driver_name,
attempt = retry_metadata.attempts + 1,
max_retries = config.max_retries,
wait_secs = wait.as_secs_f64(),
"rate limit or transient error, retrying"
);
retry_metadata.record_retry(wait, rate_limit_info);
tokio::time::sleep(wait).await;
continue;
}
RetryDecision::RetryNow => continue,
RetryDecision::Terminal(err) => return Err(err),
}
};
if retry_metadata.had_retries() {
tracing::info!(
driver = driver_name,
attempts = retry_metadata.attempts,
total_wait_secs = retry_metadata.total_retry_wait.as_secs_f64(),
"request succeeded after retries"
);
}
Ok((response, retry_metadata))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_config_matches_official_sdks() {
let config = LlmRetryConfig::default();
assert_eq!(config.max_retries, 2); assert_eq!(config.initial_backoff, Duration::from_secs(1));
assert_eq!(config.max_backoff, Duration::from_secs(60));
assert_eq!(config.backoff_multiplier, 2.0);
assert!((config.jitter_factor - 0.25).abs() < 0.001); }
#[test]
fn test_calculate_backoff_exponential() {
let config = LlmRetryConfig {
initial_backoff: Duration::from_secs(1),
max_backoff: Duration::from_secs(60),
backoff_multiplier: 2.0,
jitter_factor: 0.0, ..Default::default()
};
assert_eq!(config.calculate_backoff(0), Duration::from_secs(1));
assert_eq!(config.calculate_backoff(1), Duration::from_secs(2));
assert_eq!(config.calculate_backoff(2), Duration::from_secs(4));
assert_eq!(config.calculate_backoff(3), Duration::from_secs(8));
}
#[test]
fn test_calculate_backoff_capped() {
let config = LlmRetryConfig {
initial_backoff: Duration::from_secs(10),
max_backoff: Duration::from_secs(30),
backoff_multiplier: 2.0,
jitter_factor: 0.0,
..Default::default()
};
assert_eq!(config.calculate_backoff(0), Duration::from_secs(10));
assert_eq!(config.calculate_backoff(1), Duration::from_secs(20));
assert_eq!(config.calculate_backoff(2), Duration::from_secs(30));
assert_eq!(config.calculate_backoff(3), Duration::from_secs(30));
}
#[test]
fn test_backoff_jitter_is_randomized() {
let config = LlmRetryConfig {
initial_backoff: Duration::from_secs(10),
max_backoff: Duration::from_secs(60),
backoff_multiplier: 2.0,
jitter_factor: 0.25,
..Default::default()
};
let samples: std::collections::HashSet<u128> = (0..20)
.map(|_| config.calculate_backoff(1).as_nanos())
.collect();
assert!(
samples.len() > 1,
"jittered backoff should vary across calls, got {} distinct value(s)",
samples.len()
);
for _ in 0..50 {
let secs = config.calculate_backoff(1).as_secs_f64();
assert!(
(15.0..=25.0).contains(&secs),
"backoff {secs}s out of range"
);
}
}
#[test]
fn test_parse_duration_string() {
assert_eq!(parse_duration_string("1s"), Some(1));
assert_eq!(parse_duration_string("30s"), Some(30));
assert_eq!(parse_duration_string("1m"), Some(60));
assert_eq!(parse_duration_string("6m0s"), Some(360));
assert_eq!(parse_duration_string("1h"), Some(3600));
assert_eq!(parse_duration_string("1h30m"), Some(5400));
assert_eq!(parse_duration_string("1h30m45s"), Some(5445));
assert_eq!(parse_duration_string(""), None);
assert_eq!(parse_duration_string("invalid"), None);
}
#[test]
fn test_rate_limit_info_recommended_wait_with_retry_after() {
let config = LlmRetryConfig::default();
let info = RateLimitInfo {
retry_after_secs: Some(10),
..Default::default()
};
assert_eq!(info.recommended_wait(&config, 0), Duration::from_secs(10));
assert_eq!(info.recommended_wait(&config, 5), Duration::from_secs(10));
}
#[test]
fn test_rate_limit_info_recommended_wait_capped_at_60s() {
let config = LlmRetryConfig {
jitter_factor: 0.0, ..Default::default()
};
let info = RateLimitInfo {
retry_after_secs: Some(120), ..Default::default()
};
assert_eq!(info.recommended_wait(&config, 0), Duration::from_secs(1));
}
#[test]
fn test_rate_limit_info_recommended_wait_fallback() {
let config = LlmRetryConfig {
initial_backoff: Duration::from_secs(1),
backoff_multiplier: 2.0,
jitter_factor: 0.0,
..Default::default()
};
let info = RateLimitInfo::default();
assert_eq!(info.recommended_wait(&config, 0), Duration::from_secs(1));
assert_eq!(info.recommended_wait(&config, 1), Duration::from_secs(2));
}
#[test]
fn test_retry_metadata_record() {
let mut meta = RetryMetadata::default();
assert!(!meta.had_retries());
assert_eq!(meta.attempts, 0);
meta.record_retry(Duration::from_secs(1), None);
assert!(meta.had_retries());
assert_eq!(meta.attempts, 1);
assert_eq!(meta.total_retry_wait, Duration::from_secs(1));
meta.record_retry(Duration::from_secs(2), None);
assert_eq!(meta.attempts, 2);
assert_eq!(meta.total_retry_wait, Duration::from_secs(3));
}
#[test]
fn test_is_transient_error_matches_official_sdks() {
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(
reqwest::StatusCode::INTERNAL_SERVER_ERROR
)); 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));
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)); }
#[tokio::test]
async fn test_is_transient_send_error_on_connection_refused() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
drop(listener);
let err = reqwest::Client::new()
.get(format!("http://{addr}/"))
.send()
.await
.expect_err("request to a closed port should fail");
assert!(
is_transient_send_error(&err),
"connection-refused send error should be transient: {err:?}"
);
}
#[test]
fn test_is_transient_error_message_detects_provider_server_errors() {
assert!(is_transient_error_message(
"server_error: An error occurred while processing your request."
));
assert!(is_transient_error_message("Rate limit exceeded"));
assert!(is_transient_error_message(
"Service temporarily unavailable"
));
}
#[test]
fn test_is_transient_error_message_rejects_non_retryable_messages() {
assert!(!is_transient_error_message(
"invalid_request_error: bad tool schema"
));
assert!(!is_transient_error_message("Model not available: gpt-99"));
}
#[test]
fn structured_stream_error_prefers_code_and_status_over_message() {
use crate::driver_registry::LlmStreamError;
assert!(is_transient_stream_error(&LlmStreamError::provider(
Some("processing_error"),
None,
"An error occurred while processing your request.",
)));
assert!(is_transient_stream_error(&LlmStreamError::provider(
None::<String>,
Some(503),
"opaque failure",
)));
assert!(!is_transient_stream_error(&LlmStreamError::provider(
Some("invalid_request_error"),
Some(503),
"server unavailable",
)));
assert!(!is_transient_stream_error(&LlmStreamError::provider(
Some("insufficient_quota"),
Some(429),
"rate limit",
)));
}
#[test]
fn test_max_retry_after_constant() {
assert_eq!(MAX_RETRY_AFTER_SECS, 60);
}
fn fake_response(status: u16, body: &str) -> reqwest::Response {
let http_response = http::Response::builder()
.status(status)
.body(body.to_string())
.unwrap();
reqwest::Response::from(http_response)
}
fn fast_config(max_retries: u32) -> LlmRetryConfig {
LlmRetryConfig {
max_retries,
initial_backoff: Duration::from_millis(0),
max_backoff: Duration::from_millis(0),
backoff_multiplier: 1.0,
jitter_factor: 0.0,
}
}
#[tokio::test]
async fn test_retry_request_success_first_try() {
let config = fast_config(2);
let (resp, meta) = retry_request(
&config,
"TestDriver",
|| async { Ok(fake_response(200, "ok")) },
|_resp, _attempt, _can_retry| async {
RetryDecision::Terminal(AgentLoopError::llm("unreachable"))
},
|e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
)
.await
.expect("should succeed");
assert!(resp.status().is_success());
assert_eq!(meta.attempts, 0);
assert!(!meta.had_retries());
}
#[tokio::test]
async fn test_retry_request_retries_then_succeeds() {
let config = fast_config(3);
let calls = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
let calls_send = calls.clone();
let (resp, meta) = retry_request(
&config,
"TestDriver",
move || {
let calls = calls_send.clone();
async move {
let n = calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if n < 2 {
Ok(fake_response(429, "rate limited"))
} else {
Ok(fake_response(200, "ok"))
}
}
},
|_resp, _attempt, can_retry| async move {
assert!(can_retry, "429 within budget should be retryable");
RetryDecision::Retry {
wait: Duration::from_millis(0),
rate_limit_info: None,
}
},
|e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
)
.await
.expect("should eventually succeed");
assert!(resp.status().is_success());
assert_eq!(meta.attempts, 2);
}
#[tokio::test]
async fn test_retry_request_terminal_decision_propagates() {
let config = fast_config(2);
let result = retry_request(
&config,
"TestDriver",
|| async { Ok(fake_response(400, "bad request")) },
|_resp, _attempt, _can_retry| async {
RetryDecision::Terminal(AgentLoopError::llm("classified terminal"))
},
|e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
)
.await;
let err = result.expect_err("terminal decision should error");
assert!(err.to_string().contains("classified terminal"));
}
#[tokio::test]
async fn test_retry_request_retry_now_does_not_count_attempt() {
let config = fast_config(2);
let calls = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
let calls_send = calls.clone();
let (resp, meta) = retry_request(
&config,
"TestDriver",
move || {
let calls = calls_send.clone();
async move {
let n = calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if n == 0 {
Ok(fake_response(400, "max_tokens too large"))
} else {
Ok(fake_response(200, "ok"))
}
}
},
{
let mut used_fallback = false;
move |_resp, _attempt, _can_retry| {
let do_fallback = !used_fallback;
used_fallback = true;
async move {
if do_fallback {
RetryDecision::RetryNow
} else {
RetryDecision::Terminal(AgentLoopError::llm("unreachable"))
}
}
}
},
|e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
)
.await
.expect("RetryNow then success");
assert!(resp.status().is_success());
assert_eq!(meta.attempts, 0);
}
#[tokio::test]
async fn test_retry_request_send_error_exhausts() {
let config = fast_config(1);
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
drop(listener);
let make_err = || async {
reqwest::Client::new()
.get(format!("http://{addr}/"))
.send()
.await
.expect_err("closed port")
};
let result = retry_request(
&config,
"TestDriver",
move || async move { Err(SendOutcome::Send(make_err().await)) },
|_resp, _attempt, _can_retry| async {
RetryDecision::Terminal(AgentLoopError::llm("unreachable"))
},
|e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
)
.await;
let err = result.expect_err("send errors should exhaust to terminal");
assert!(err.to_string().contains("after 1 retries"), "got: {err}");
}
#[tokio::test]
async fn test_retry_request_fatal_send_propagates_immediately() {
let config = fast_config(3);
let result = retry_request(
&config,
"TestDriver",
|| async { Err(SendOutcome::Fatal(AgentLoopError::llm("auth failed"))) },
|_resp, _attempt, _can_retry| async {
RetryDecision::Terminal(AgentLoopError::llm("unreachable"))
},
|e, attempts| AgentLoopError::llm(send_error_message(e, attempts)),
)
.await;
let err = result.expect_err("fatal send should propagate");
assert!(err.to_string().contains("auth failed"));
}
}