polyc_agent/

retry.rs

//! Bounded retry for the model call with exponential backoff + jitter.
//!
//! The turn loop's `provider.complete(req)` is the connect/initial-response
//! boundary: it returns `Err` for a refused connection, a 429, or an immediate
//! 5xx *before* any chunk is yielded, so retrying it is safe (nothing has been
//! forwarded to the client yet). A failure *during* streaming — after chunks may
//! already have been forwarded — is the caller's and is NOT retried here.
//!
//! Only retryable kinds (rate-limit / timeout / unavailable) are retried; a
//! terminal error (auth / bad-request) returns immediately. When the provider
//! captured a server `Retry-After`, that wait is honored over the computed
//! backoff.

use std::time::Duration;

use polyc_llm::{Chunk, CompletionRequest, LlmError, LlmErrorKind, LlmProvider};

/// How many times to retry, and the backoff envelope.
#[derive(Debug, Clone, Copy)]
pub struct RetryConfig {
    /// Maximum retries *after* the first attempt (so `max_retries` of 4 ⇒ up to
    /// 5 total calls).
    pub max_retries: u32,
    /// Base delay; the nth retry waits ~`base * 2^n` (jittered, capped).
    pub base_delay: Duration,
    /// Ceiling on a single backoff wait.
    pub max_delay: Duration,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_retries: 4,
            base_delay: Duration::from_millis(500),
            max_delay: Duration::from_secs(30),
        }
    }
}

impl RetryConfig {
    /// Load from the environment, falling back to [`Default`] for any unset or
    /// unparseable value:
    /// - `POLYCHROME_LLM_MAX_RETRIES`
    /// - `POLYCHROME_LLM_RETRY_BASE_MS`
    /// - `POLYCHROME_LLM_RETRY_MAX_MS`
    #[must_use]
    pub fn from_env() -> Self {
        let d = Self::default();
        Self {
            max_retries: env_parse("POLYCHROME_LLM_MAX_RETRIES").unwrap_or(d.max_retries),
            base_delay: env_parse("POLYCHROME_LLM_RETRY_BASE_MS")
                .map_or(d.base_delay, Duration::from_millis),
            max_delay: env_parse("POLYCHROME_LLM_RETRY_MAX_MS")
                .map_or(d.max_delay, Duration::from_millis),
        }
    }
}

fn env_parse<T: std::str::FromStr>(key: &str) -> Option<T> {
    std::env::var(key).ok()?.parse().ok()
}

/// Whether an error kind warrants a retry.
const fn is_retryable(kind: LlmErrorKind) -> bool {
    matches!(
        kind,
        LlmErrorKind::RateLimit | LlmErrorKind::Timeout | LlmErrorKind::Unavailable
    )
}

/// Exponential backoff with equal jitter.
///
/// Half the delay is fixed and half is scaled by `jitter_frac` ∈ [0, 1), so the
/// wait lands in `[0.5, 1.0] × base × 2^attempt` (capped). Pure, so callers and
/// tests control the jitter.
#[must_use]
pub fn backoff_delay(attempt: u32, base: Duration, cap: Duration, jitter_frac: f64) -> Duration {
    // `2^attempt`, saturating so a large attempt can't panic on shift overflow.
    let factor = 1u32.checked_shl(attempt.min(16)).unwrap_or(u32::MAX);
    let exp = base.saturating_mul(factor).min(cap);
    // Equal jitter: 0.5 + 0.5*frac, written as a fused multiply-add.
    let scale = 0.5_f64.mul_add(jitter_frac.clamp(0.0, 1.0), 0.5);
    exp.mul_f64(scale)
}

/// Cheap entropy for jitter — the exact value is irrelevant (it only spreads
/// retries to avoid a thundering herd), so a clock read suffices.
fn jitter_frac() -> f64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_or(0, |d| d.subsec_nanos());
    f64::from(nanos % 1_000_000) / 1_000_000.0
}

/// Call `provider.complete`, retrying retryable failures up to `cfg.max_retries`
/// with exponential backoff + jitter, honoring a server `Retry-After`
/// ([`LlmError::retry_after`]) when present.
///
/// # Errors
///
/// Returns the provider's error once it is terminal (non-retryable) or the retry
/// budget is exhausted.
pub async fn complete_with_retry<P>(
    provider: &P,
    req: CompletionRequest,
    cfg: &RetryConfig,
) -> Result<futures::stream::BoxStream<'static, Result<Chunk, P::Error>>, P::Error>
where
    P: LlmProvider + ?Sized,
{
    let mut attempt = 0u32;
    loop {
        match provider.complete(req.clone()).await {
            Ok(stream) => return Ok(stream),
            Err(err) => {
                let kind = err.kind();
                if !is_retryable(kind) || attempt >= cfg.max_retries {
                    return Err(err);
                }
                // Honor a server `Retry-After` (capped), but fall back to
                // computed backoff for a zero/absent value — a literal
                // `Retry-After: 0` must NOT collapse the spacing into a tight,
                // sleep-free retry loop against an already-rate-limited upstream.
                let delay = err.retry_after().filter(|d| !d.is_zero()).map_or_else(
                    || backoff_delay(attempt, cfg.base_delay, cfg.max_delay, jitter_frac()),
                    |d| d.min(cfg.max_delay),
                );
                attempt += 1;
                tracing::warn!(
                    attempt,
                    ?kind,
                    delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX),
                    "model call failed; retrying"
                );
                tokio::time::sleep(delay).await;
            }
        }
    }
}

#[cfg(test)]
#[allow(clippy::pedantic, clippy::nursery, missing_docs)]
mod tests {
    use std::sync::atomic::{AtomicUsize, Ordering};

    use async_trait::async_trait;
    use futures::StreamExt as _;
    use polyc_llm::error::DummyError;
    use polyc_llm::{Chunk, CompletionRequest, LlmProvider, StopReason};

    use super::*;

    /// Fails the first `fail_n` calls with `err`, then streams a one-chunk turn.
    struct FlakyProvider {
        calls: AtomicUsize,
        fail_n: usize,
        err: fn() -> DummyError,
    }

    #[async_trait]
    impl LlmProvider for FlakyProvider {
        type Error = DummyError;
        async fn complete(
            &self,
            _req: CompletionRequest,
        ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
        {
            let n = self.calls.fetch_add(1, Ordering::SeqCst);
            if n < self.fail_n {
                return Err((self.err)());
            }
            Ok(futures::stream::iter(vec![Ok(Chunk::Stop(StopReason::EndTurn))]).boxed())
        }
    }

    fn fast_cfg(max_retries: u32) -> RetryConfig {
        RetryConfig {
            max_retries,
            base_delay: Duration::from_millis(0),
            max_delay: Duration::from_millis(0),
        }
    }

    fn unavailable() -> DummyError {
        DummyError::Transport("reset".to_owned())
    }
    fn bad_request() -> DummyError {
        DummyError::Provider {
            status: 400,
            body: "nope".to_owned(),
        }
    }

    #[tokio::test]
    async fn retries_then_succeeds() {
        let p = FlakyProvider {
            calls: AtomicUsize::new(0),
            fail_n: 2,
            err: unavailable,
        };
        let out = complete_with_retry(&p, CompletionRequest::new("m"), &fast_cfg(4)).await;
        assert!(out.is_ok(), "should succeed after 2 retries");
        assert_eq!(p.calls.load(Ordering::SeqCst), 3, "2 failures + 1 success");
    }

    #[tokio::test]
    async fn gives_up_after_budget() {
        let p = FlakyProvider {
            calls: AtomicUsize::new(0),
            fail_n: 99,
            err: unavailable,
        };
        let out = complete_with_retry(&p, CompletionRequest::new("m"), &fast_cfg(3)).await;
        assert!(out.is_err(), "exhausts the budget");
        // initial attempt + 3 retries = 4 calls.
        assert_eq!(p.calls.load(Ordering::SeqCst), 4);
    }

    #[tokio::test]
    async fn terminal_error_is_not_retried() {
        let p = FlakyProvider {
            calls: AtomicUsize::new(0),
            fail_n: 99,
            err: bad_request,
        };
        let out = complete_with_retry(&p, CompletionRequest::new("m"), &fast_cfg(4)).await;
        assert!(out.is_err());
        assert_eq!(p.calls.load(Ordering::SeqCst), 1, "bad-request is terminal");
    }

    #[test]
    fn backoff_grows_and_caps() {
        let base = Duration::from_millis(100);
        let cap = Duration::from_millis(1000);
        // jitter_frac = 1.0 → full delay (no shrink).
        assert_eq!(backoff_delay(0, base, cap, 1.0), Duration::from_millis(100));
        assert_eq!(backoff_delay(1, base, cap, 1.0), Duration::from_millis(200));
        assert_eq!(backoff_delay(2, base, cap, 1.0), Duration::from_millis(400));
        // 100 * 2^4 = 1600 → capped at 1000.
        assert_eq!(
            backoff_delay(4, base, cap, 1.0),
            Duration::from_millis(1000)
        );
        // jitter_frac = 0.0 → half delay (equal-jitter floor).
        assert_eq!(backoff_delay(1, base, cap, 0.0), Duration::from_millis(100));
    }
}