open_agent/

retry.rs

//! Retry utilities with exponential backoff
//!
//! This module provides utilities for retrying operations with configurable
//! backoff strategies. Useful for handling transient failures when communicating
//! with LLM servers.
//!
//! # Examples
//!
//! ```rust,no_run
//! use open_agent::retry::{retry_with_backoff, RetryConfig};
//! use std::time::Duration;
//!
//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
//! let config = RetryConfig::default()
//!     .with_max_attempts(3)
//!     .with_initial_delay(Duration::from_secs(1));
//!
//! let result = retry_with_backoff(config, || async {
//!     // Your async operation here
//!     Ok::<_, open_agent::Error>(42)
//! }).await?;
//! # Ok(())
//! # }
//! ```

use crate::{Error, Result};
use std::future::Future;
use std::time::Duration;
use tokio::time::sleep;

/// Configuration for retry behavior
#[derive(Debug, Clone)]
pub struct RetryConfig {
    /// Maximum number of retry attempts
    pub max_attempts: u32,

    /// Initial delay before first retry
    pub initial_delay: Duration,

    /// Maximum delay between retries
    pub max_delay: Duration,

    /// Multiplier for exponential backoff (e.g., 2.0 doubles the delay each time)
    pub backoff_multiplier: f64,

    /// Add random jitter to prevent thundering herd (0.0 to 1.0)
    pub jitter_factor: f64,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            initial_delay: Duration::from_secs(1),
            max_delay: Duration::from_secs(60),
            backoff_multiplier: 2.0,
            jitter_factor: 0.1,
        }
    }
}

impl RetryConfig {
    /// Create a new retry configuration
    pub fn new() -> Self {
        Self::default()
    }

    /// Set maximum number of attempts
    pub fn with_max_attempts(mut self, attempts: u32) -> Self {
        self.max_attempts = attempts;
        self
    }

    /// Set initial delay
    pub fn with_initial_delay(mut self, delay: Duration) -> Self {
        self.initial_delay = delay;
        self
    }

    /// Set maximum delay
    pub fn with_max_delay(mut self, delay: Duration) -> Self {
        self.max_delay = delay;
        self
    }

    /// Set backoff multiplier
    pub fn with_backoff_multiplier(mut self, multiplier: f64) -> Self {
        self.backoff_multiplier = multiplier;
        self
    }

    /// Set jitter factor (0.0 to 1.0)
    pub fn with_jitter_factor(mut self, jitter: f64) -> Self {
        self.jitter_factor = jitter.clamp(0.0, 1.0);
        self
    }

    /// Calculate delay for a given attempt with exponential backoff and jitter
    fn calculate_delay(&self, attempt: u32) -> Duration {
        let base_delay_ms = self.initial_delay.as_millis() as f64;
        let exponential_delay = base_delay_ms * self.backoff_multiplier.powi(attempt as i32);

        // Cap at max delay
        let capped_delay = exponential_delay.min(self.max_delay.as_millis() as f64);

        // Add jitter, centred on the capped delay
        let jitter_range = capped_delay * self.jitter_factor;
        let jitter = rand::random::<f64>() * jitter_range;
        let jittered_delay = capped_delay + jitter - (jitter_range / 2.0);

        // Re-apply the cap. Jitter is added *after* capping, so without this a `max_delay` of
        // 60s would still produce a 66s sleep at a 0.2 jitter factor — which contradicts what
        // `with_max_delay` promises.
        let final_delay = jittered_delay.clamp(0.0, self.max_delay.as_millis() as f64);

        Duration::from_millis(final_delay as u64)
    }
}

/// Retry an async operation with exponential backoff
///
/// # Arguments
///
/// * `config` - Retry configuration
/// * `operation` - Async function to retry
///
/// # Returns
///
/// The result of the operation if successful, or the last error if all retries failed
///
/// # Examples
///
/// ```rust,no_run
/// use open_agent::retry::{retry_with_backoff, RetryConfig};
/// use open_agent::{Client, AgentOptions};
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = RetryConfig::default().with_max_attempts(3);
/// let options = AgentOptions::builder()
///     .model("qwen3:8b")
///     .base_url("http://localhost:11434/v1")
///     .build()?;
///
/// let result = retry_with_backoff(config, || async {
///     let mut client = Client::new(options.clone())?;
///     client.send("Hello").await?;
///     Ok::<_, open_agent::Error>(())
/// }).await?;
/// # Ok(())
/// # }
/// ```
pub async fn retry_with_backoff<F, Fut, T>(config: RetryConfig, mut operation: F) -> Result<T>
where
    F: FnMut() -> Fut,
    Fut: Future<Output = Result<T>>,
{
    let mut last_error = None;

    for attempt in 0..config.max_attempts {
        match operation().await {
            Ok(result) => return Ok(result),
            Err(err) => {
                last_error = Some(err);

                // Don't sleep after the last attempt
                if attempt < config.max_attempts - 1 {
                    let delay = config.calculate_delay(attempt);
                    sleep(delay).await;
                }
            }
        }
    }

    Err(last_error.unwrap_or_else(|| Error::other("Retry failed with no error")))
}

/// HTTP status codes that indicate a transient failure worth retrying.
///
/// - `408 Request Timeout` and `429 Too Many Requests` are the standard client-side signals
///   to back off and try again; 429 in particular is the single most common reason a request
///   should be retried, and rejecting it defeats the point of conditional retry.
/// - `500`, `502`, `503`, `504` are transient server/gateway failures.
/// - `529` is the de facto "overloaded" status used by several inference providers.
///
/// Other 5xx codes are deliberately excluded: `501 Not Implemented` and `505 HTTP Version Not
/// Supported` describe permanent server capabilities, so retrying them only wastes attempts.
const RETRYABLE_STATUS_CODES: &[u16] = &[408, 429, 500, 502, 503, 504, 529];

/// Determine if an error is retryable
///
/// Returns true for transient errors like network issues, timeouts, rate limiting (429), and
/// 5xx server errors. Returns false for client errors like invalid requests (4xx) or
/// configuration errors.
///
/// API errors are classified on [`Error::status_code`]. An API error raised without a status
/// — via [`Error::api`] rather than [`Error::api_status`] — is treated as non-retryable,
/// matching the conservative default applied to every other unclassified error.
pub fn is_retryable_error(error: &Error) -> bool {
    match error {
        // Transport failures, timeouts, and interrupted streams are all transient by nature.
        Error::Http(_) | Error::Timeout | Error::Stream(_) => true,
        // Everything else is retryable only if it carries a transient HTTP status.
        // `status_code()` is `None` for every non-API variant, so configuration errors,
        // invalid input, tool failures, and serialization failures all fall through to false.
        _ => error
            .status_code()
            .is_some_and(|status| RETRYABLE_STATUS_CODES.contains(&status)),
    }
}

/// Retry an async operation with exponential backoff, only retrying on retryable errors
///
/// This is a smarter version of `retry_with_backoff` that only retries transient errors.
///
/// # Examples
///
/// ```rust,no_run
/// use open_agent::retry::{retry_with_backoff_conditional, RetryConfig};
/// use open_agent::{Client, AgentOptions};
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = RetryConfig::default();
/// let options = AgentOptions::builder()
///     .model("qwen3:8b")
///     .base_url("http://localhost:11434/v1")
///     .build()?;
///
/// let result = retry_with_backoff_conditional(config, || async {
///     let mut client = Client::new(options.clone())?;
///     client.send("Hello").await?;
///     Ok::<_, open_agent::Error>(())
/// }).await?;
/// # Ok(())
/// # }
/// ```
pub async fn retry_with_backoff_conditional<F, Fut, T>(
    config: RetryConfig,
    mut operation: F,
) -> Result<T>
where
    F: FnMut() -> Fut,
    Fut: Future<Output = Result<T>>,
{
    let mut last_error = None;

    for attempt in 0..config.max_attempts {
        match operation().await {
            Ok(result) => return Ok(result),
            Err(err) => {
                // Check if error is retryable
                if !is_retryable_error(&err) {
                    return Err(err);
                }

                last_error = Some(err);

                // Don't sleep after the last attempt
                if attempt < config.max_attempts - 1 {
                    let delay = config.calculate_delay(attempt);
                    sleep(delay).await;
                }
            }
        }
    }

    Err(last_error.unwrap_or_else(|| Error::other("Retry failed with no error")))
}

#[cfg(test)]
mod tests {
    use super::*;

    include!("retry/tests.rs");
}