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, SystemTime, UNIX_EPOCH};

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

/// The turn's clock and jitter source — its only non-determinism.
///
/// The retry backoff is the sole place the turn loop reads wall time (for jitter
/// entropy) or waits (for the backoff), so routing both through an injected
/// capability makes a whole turn deterministically replayable: production wires
/// the real clock ([`RealClock`]) and behaves exactly as before, while a test
/// wires a virtual clock with a fixed jitter seed and gets byte-identical output
/// it can step without a wall-clock wait.
#[async_trait]
pub trait Clock: std::fmt::Debug + Send + Sync {
    /// The current wall-clock time.
    fn now(&self) -> SystemTime;

    /// A jitter fraction in `[0, 1)` used to spread retries (equal jitter in
    /// [`backoff_delay`]). The exact value only matters for the spread, so the
    /// default derives it from [`Self::now`] — the same cheap clock read the
    /// retry path used before the seam existed. A deterministic clock overrides
    /// this with a seeded draw so a replay reproduces the same spacing.
    fn jitter_frac(&self) -> f64 {
        let nanos = self
            .now()
            .duration_since(UNIX_EPOCH)
            .map_or(0, |d| d.subsec_nanos());
        f64::from(nanos % 1_000_000) / 1_000_000.0
    }

    /// Wait for `dur` before the caller retries.
    ///
    /// # Cancellation
    ///
    /// Cancellation-safe: dropping the returned future cancels the wait with no
    /// observable effect, exactly like the underlying timer.
    async fn sleep(&self, dur: Duration);
}

/// The production [`Clock`]: real wall time and a real timer.
///
/// [`Clock::now`] reads the system clock and [`Clock::sleep`] awaits a tokio
/// timer, so the retry path spreads and waits exactly as it did before the seam
/// was introduced.
#[derive(Debug, Default, Clone, Copy)]
pub struct RealClock;

#[async_trait]
impl Clock for RealClock {
    fn now(&self) -> SystemTime {
        SystemTime::now()
    }

    async fn sleep(&self, dur: Duration) {
        tokio::time::sleep(dur).await;
    }
}

/// 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),
        }
    }
}

/// Parse an environment variable, treating unset OR unparseable as `None` (the
/// caller falls back to a default) rather than erroring — a malformed override
/// must degrade to the shipped default, never fail startup.
///
/// `pub(crate)` so other per-deployment env/config knobs in this crate (e.g.
/// [`crate::resolve_max_steps`]) share the exact same "unset or unparseable ⇒
/// default" resolution instead of re-deriving it.
pub(crate) 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)
}

/// 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.
///
/// `clock` supplies the jitter entropy and the backoff wait; production passes
/// [`RealClock`], so behavior is unchanged, while a test passes a virtual clock
/// with a fixed seed to make the retry deterministic.
///
/// # 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,
    clock: &dyn Clock,
) -> 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, clock.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"
                );
                clock.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), &RealClock).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), &RealClock).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), &RealClock).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));
    }
}