diff --git a/core/src/agent/provider/gemini.rs b/core/src/agent/provider/gemini.rs index f2a8461..e15755a 100644 --- a/core/src/agent/provider/gemini.rs +++ b/core/src/agent/provider/gemini.rs @@ -77,13 +77,14 @@ impl LlmProvider for GeminiProvider { .map_err(|e| ProviderError::Transport(e.to_string()))?; let status = resp.status(); + let retry_after = parse_retry_after(resp.headers()); let text = resp .text() .await .map_err(|e| ProviderError::Transport(e.to_string()))?; if !status.is_success() { - return Err(classify_http_error(status.as_u16(), &text)); + return Err(classify_http_error(status.as_u16(), &text, retry_after)); } let wire: GeminiResponseBody = serde_json::from_str(&text) @@ -93,8 +94,14 @@ impl LlmProvider for GeminiProvider { } /// Map a non-2xx HTTP response to a typed error, pulling Gemini's -/// `error.message` when present. -fn classify_http_error(status: u16, body: &str) -> ProviderError { +/// `error.message` when present. The retry classifier ([A6]) decides +/// what's worth a second attempt — this just picks the variant; see +/// [`ProviderError::is_retryable`]. +fn classify_http_error( + status: u16, + body: &str, + retry_after: Option, +) -> ProviderError { let message = serde_json::from_str::(body) .ok() .map(|e| e.error.message) @@ -102,10 +109,25 @@ fn classify_http_error(status: u16, body: &str) -> ProviderError { match status { 400 => ProviderError::InvalidRequest(message), 401 | 403 => ProviderError::Auth(message), + 429 => ProviderError::RateLimit { + message, + retry_after, + }, + 500..=599 => ProviderError::Transient(format!("HTTP {status}: {message}")), _ => ProviderError::Provider(format!("HTTP {status}: {message}")), } } +/// Parse a `Retry-After` header value as seconds. Gemini sends an integer +/// seconds count rather than the HTTP-date alternative the spec permits, +/// so we only handle the seconds case — an unparseable header yields +/// `None` and the retry loop falls back to exponential backoff. +fn parse_retry_after(headers: &reqwest::header::HeaderMap) -> Option { + let raw = headers.get(reqwest::header::RETRY_AFTER)?.to_str().ok()?; + let secs: u64 = raw.trim().parse().ok()?; + Some(std::time::Duration::from_secs(secs)) +} + // --------------------------------------------------------------------------- // Neutral -> Gemini wire // --------------------------------------------------------------------------- @@ -511,9 +533,70 @@ mod tests { #[test] fn http_errors_are_classified() { let body = r#"{"error": {"code": 403, "message": "API key not valid", "status": "PERMISSION_DENIED"}}"#; - assert!(matches!(classify_http_error(403, body), ProviderError::Auth(m) if m.contains("not valid"))); - assert!(matches!(classify_http_error(400, body), ProviderError::InvalidRequest(_))); - assert!(matches!(classify_http_error(503, body), ProviderError::Provider(_))); + assert!(matches!( + classify_http_error(403, body, None), + ProviderError::Auth(m) if m.contains("not valid") + )); + assert!(matches!( + classify_http_error(400, body, None), + ProviderError::InvalidRequest(_) + )); + } + + /// 5xx maps to `Transient` so the retry loop ([A6]) backs off and + /// tries again; the message preserves the HTTP status for debugging. + #[test] + fn five_hundreds_classify_as_transient() { + let body = r#"{"error": {"code": 503, "message": "backend overloaded"}}"#; + for status in [500, 502, 503, 504] { + let err = classify_http_error(status, body, None); + assert!( + matches!(&err, ProviderError::Transient(m) if m.contains(&status.to_string())), + "expected Transient for {status}, got {err:?}" + ); + assert!(err.is_retryable()); + } + } + + /// 429 maps to `RateLimit`; a `Retry-After` header threads through + /// to the variant so the retry helper honours it verbatim. + #[test] + fn rate_limited_passes_retry_after_through() { + let body = r#"{"error": {"code": 429, "message": "quota exceeded"}}"#; + let hint = std::time::Duration::from_secs(7); + let err = classify_http_error(429, body, Some(hint)); + let ProviderError::RateLimit { retry_after, message } = err else { + panic!("expected RateLimit"); + }; + assert_eq!(retry_after, Some(hint)); + assert_eq!(message, "quota exceeded"); + } + + /// 4xx codes we don't model explicitly stay in the unclassified + /// `Provider` bucket (not retried). + #[test] + fn unmodeled_4xx_stays_in_provider_bucket() { + let body = r#"{"error": {"code": 404, "message": "model not found"}}"#; + let err = classify_http_error(404, body, None); + assert!(matches!(&err, ProviderError::Provider(m) if m.contains("404"))); + assert!(!err.is_retryable()); + } + + #[test] + fn parse_retry_after_handles_integer_seconds() { + use reqwest::header::{HeaderMap, HeaderValue, RETRY_AFTER}; + let mut h = HeaderMap::new(); + h.insert(RETRY_AFTER, HeaderValue::from_static("12")); + assert_eq!(parse_retry_after(&h), Some(std::time::Duration::from_secs(12))); + + // Non-numeric (e.g. HTTP-date format) yields None — caller falls + // back to exponential backoff. + let mut h = HeaderMap::new(); + h.insert(RETRY_AFTER, HeaderValue::from_static("Wed, 21 Oct 2026 07:28:00 GMT")); + assert!(parse_retry_after(&h).is_none()); + + // Missing header. + assert!(parse_retry_after(&HeaderMap::new()).is_none()); } /// Live smoke test against the real Gemini API. Ignored by default; run with diff --git a/core/src/agent/provider/mod.rs b/core/src/agent/provider/mod.rs index df0fee0..11b0fb1 100644 --- a/core/src/agent/provider/mod.rs +++ b/core/src/agent/provider/mod.rs @@ -23,7 +23,9 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; pub mod gemini; +pub mod retry; pub use gemini::GeminiProvider; +pub use retry::{retry_with_backoff, RetryConfig}; /// Who authored a message in the conversation. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -251,27 +253,78 @@ pub struct ChatResponse { pub finish_reason: FinishReason, } -/// Errors a provider can surface. Intentionally small for `#8`; `#13` -/// extends this with rate-limit/transient classification and backoff. +/// Errors a provider can surface. Classification is what the retry loop +/// (`#13`) branches on — see [`ProviderError::is_retryable`] and +/// [`ProviderError::retry_after`]. #[derive(Debug, thiserror::Error)] pub enum ProviderError { - /// Authentication failed (missing/invalid API key). Not retryable. + /// Authentication failed (missing/invalid API key). Not retryable — + /// retrying will not change the answer; surface to the UI immediately. #[error("provider auth failed: {0}")] Auth(String), - /// The request was rejected as malformed (bad model id, unsupported field). + /// The request was rejected as malformed (bad model id, unsupported + /// field). Not retryable — retrying will not change the answer. #[error("invalid request: {0}")] InvalidRequest(String), - /// Transport / network failure talking to the provider. + /// Transport / network failure talking to the provider. Retryable — + /// a TCP reset, DNS hiccup, or proxy blip should not bubble up as a + /// hard failure to the user the first time. #[error("transport error: {0}")] Transport(String), - /// The provider returned a response we couldn't parse into neutral types. + /// The provider returned a response we couldn't parse into neutral + /// types. Not retryable — the bug is in our decoder or the provider's + /// wire format, not in the network. #[error("could not parse provider response: {0}")] Decode(String), - /// Provider-side error not otherwise classified (5xx, quota, etc.). + /// Provider-side rate limit / quota response (typically HTTP 429). + /// Retryable. `retry_after` carries the server's hint (the + /// `Retry-After` header) when present; the retry loop should honour + /// it instead of its own backoff schedule. + #[error("rate limited: {message}")] + RateLimit { + message: String, + retry_after: Option, + }, + /// Transient provider failure (typically HTTP 5xx, timeouts). + /// Retryable with exponential backoff. + #[error("transient provider error: {0}")] + Transient(String), + /// Provider-side error not otherwise classified — bucket for shapes + /// we don't model explicitly. Not retried by default to avoid + /// hammering a provider with a request it has already rejected for a + /// non-transient reason. #[error("provider error: {0}")] Provider(String), } +impl ProviderError { + /// True if the retry loop should attempt this call again. Drives the + /// `retry_with_backoff` decision tree alongside the configured retry + /// budget. + /// + /// Retryable: [`Self::RateLimit`], [`Self::Transient`], [`Self::Transport`]. + /// Not retryable: [`Self::Auth`], [`Self::InvalidRequest`], + /// [`Self::Decode`], [`Self::Provider`]. + pub fn is_retryable(&self) -> bool { + matches!( + self, + ProviderError::RateLimit { .. } + | ProviderError::Transient(_) + | ProviderError::Transport(_) + ) + } + + /// The provider's explicit "wait at least this long" hint, if any. + /// Currently only [`Self::RateLimit`] carries this — `Retry-After` + /// headers are the only place providers reliably set it. + pub fn retry_after(&self) -> Option { + match self { + ProviderError::RateLimit { retry_after, .. } => *retry_after, + _ => None, + } + } +} + /// The contract every LLM provider implements. One method: take a neutral /// [`ChatRequest`], return a neutral [`ChatResponse`]. The API key is passed /// per call (never stored on the provider) so it can come from the OS @@ -389,6 +442,36 @@ mod tests { } } + #[test] + fn is_retryable_matches_classification_table() { + // Retryable. + assert!(ProviderError::Transport("eof".into()).is_retryable()); + assert!(ProviderError::Transient("5xx".into()).is_retryable()); + assert!(ProviderError::RateLimit { + message: "slow down".into(), + retry_after: None, + } + .is_retryable()); + // Not retryable. + assert!(!ProviderError::Auth("bad key".into()).is_retryable()); + assert!(!ProviderError::InvalidRequest("bad model".into()).is_retryable()); + assert!(!ProviderError::Decode("bad json".into()).is_retryable()); + assert!(!ProviderError::Provider("safety".into()).is_retryable()); + } + + #[test] + fn retry_after_only_set_for_rate_limit() { + let d = std::time::Duration::from_secs(2); + let rl = ProviderError::RateLimit { + message: "wait".into(), + retry_after: Some(d), + }; + assert_eq!(rl.retry_after(), Some(d)); + + assert!(ProviderError::Transient("x".into()).retry_after().is_none()); + assert!(ProviderError::Transport("x".into()).retry_after().is_none()); + } + #[tokio::test] async fn provider_trait_is_object_safe_and_callable() { let provider: std::sync::Arc = std::sync::Arc::new(EchoProvider); diff --git a/core/src/agent/provider/retry.rs b/core/src/agent/provider/retry.rs new file mode 100644 index 0000000..4249ef4 --- /dev/null +++ b/core/src/agent/provider/retry.rs @@ -0,0 +1,284 @@ +//! Exponential-backoff retry around a `LlmProvider::complete` call ([A6] / +//! hebb_app#13). +//! +//! The agent reasoning loop should not give up the first time a provider +//! returns 429 or 500 — those are usually transient. This module's +//! `retry_with_backoff` wraps an arbitrary async closure that returns +//! `Result` and retries on errors the classifier marks +//! retryable, applying: +//! +//! - exponential backoff with a configurable initial delay and multiplier; +//! - a cap on the per-attempt delay so the loop doesn't sleep for minutes +//! in a Cargo test by accident; +//! - random jitter so concurrent sessions don't synchronize a thundering +//! herd on the same provider; +//! - the server's `Retry-After` hint (via [`ProviderError::retry_after`]) +//! overrides the computed delay so a 429 with `Retry-After: 30` waits +//! exactly that long. +//! +//! Pure async; no tokio task spawns. The retry loop is just `tokio::time::sleep` +//! between calls. + +use std::future::Future; +use std::time::Duration; + +use rand::Rng; + +use super::ProviderError; + +/// Caller-facing tuning for [`retry_with_backoff`]. Defaults are picked +/// for an interactive agent turn: +/// +/// - 3 retries (4 total attempts). +/// - 250 ms initial delay, doubling each attempt, capped at 8 s. +/// - ±25% jitter on the computed delay. +/// +/// These are conservative — a 4-attempt turn budget keeps a stuck session +/// honest while still surviving a hiccup. The reasoning loop ([C2]) can +/// pass `RetryConfig::aggressive()` for background work where waiting +/// longer is cheaper than failing. +#[derive(Debug, Clone, Copy)] +pub struct RetryConfig { + /// Maximum number of *retries* (i.e. total attempts = `max_retries + 1`). + pub max_retries: u32, + /// Delay before the first retry. Each subsequent retry doubles it, + /// capped at `max_delay`. + pub initial_delay: Duration, + /// Upper bound on a single sleep. Without this, attempt 5 would sleep + /// for 16 × initial — surprising in a unit test. + pub max_delay: Duration, + /// Random jitter applied to the computed delay, in `[0.0, 1.0]`. A + /// jitter of `0.25` means the realised delay is uniformly sampled + /// from `[delay * 0.75, delay * 1.25]`. + pub jitter: f64, +} + +impl RetryConfig { + /// Conservative defaults — see [`RetryConfig`] doc. + pub fn interactive() -> Self { + Self { + max_retries: 3, + initial_delay: Duration::from_millis(250), + max_delay: Duration::from_secs(8), + jitter: 0.25, + } + } + + /// More patience: 6 retries, 500 ms initial, capped at 30 s. Useful + /// for background / batch turns where finishing is cheaper than + /// failing. + pub fn aggressive() -> Self { + Self { + max_retries: 6, + initial_delay: Duration::from_millis(500), + max_delay: Duration::from_secs(30), + jitter: 0.25, + } + } + + /// Tests/CI: no waiting between retries. Lets a test exercise the + /// retry loop without `tokio::time::pause()` boilerplate. + pub fn no_wait(max_retries: u32) -> Self { + Self { + max_retries, + initial_delay: Duration::ZERO, + max_delay: Duration::ZERO, + jitter: 0.0, + } + } +} + +impl Default for RetryConfig { + fn default() -> Self { + Self::interactive() + } +} + +/// Run `op` with retries on [`ProviderError::is_retryable`] failures. +/// +/// The closure is re-invoked from scratch each attempt — callers pass it +/// as `|| async { provider.complete(&req, key).await }` so the same +/// request goes out. Non-retryable errors are returned immediately. +/// +/// The realised delay between attempts is: +/// +/// 1. If the error carries a `retry_after` hint, use that (capped at +/// `cfg.max_delay`). +/// 2. Otherwise: `initial_delay * 2^(attempt - 1)`, capped at +/// `cfg.max_delay`, then jittered by `± cfg.jitter`. +/// +/// `tokio::time::sleep` is used between attempts. Tests can use +/// `RetryConfig::no_wait` to skip sleeping. +pub async fn retry_with_backoff(cfg: RetryConfig, mut op: F) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut attempt: u32 = 0; + loop { + match op().await { + Ok(value) => return Ok(value), + Err(err) => { + if !err.is_retryable() || attempt >= cfg.max_retries { + return Err(err); + } + let delay = next_delay(&cfg, attempt, err.retry_after()); + if !delay.is_zero() { + tokio::time::sleep(delay).await; + } + attempt += 1; + } + } + } +} + +fn next_delay(cfg: &RetryConfig, attempt: u32, retry_after: Option) -> Duration { + // Respect the server's hint when present. Cap it so a misbehaving + // provider can't pin the loop on a 10-minute wait. + if let Some(hint) = retry_after { + return hint.min(cfg.max_delay); + } + if cfg.initial_delay.is_zero() { + return Duration::ZERO; + } + let base = cfg + .initial_delay + .saturating_mul(1u32.checked_shl(attempt).unwrap_or(u32::MAX)); + let capped = base.min(cfg.max_delay); + apply_jitter(capped, cfg.jitter) +} + +fn apply_jitter(d: Duration, jitter: f64) -> Duration { + if jitter <= 0.0 || d.is_zero() { + return d; + } + let j = jitter.clamp(0.0, 1.0); + let factor = 1.0 + rand::thread_rng().gen_range(-j..=j); + let ms = (d.as_millis() as f64 * factor).max(0.0) as u64; + Duration::from_millis(ms) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::cell::Cell; + use std::time::Duration; + + /// A non-retryable error returns immediately; the closure is called + /// exactly once. + #[tokio::test] + async fn auth_error_is_not_retried() { + let calls = Cell::new(0u32); + let result: Result<(), _> = retry_with_backoff(RetryConfig::no_wait(5), || { + calls.set(calls.get() + 1); + async { Err(ProviderError::Auth("bad key".into())) } + }) + .await; + assert!(matches!(result, Err(ProviderError::Auth(_)))); + assert_eq!(calls.get(), 1); + } + + /// A retryable error retries up to `max_retries`, then surfaces the + /// last error. + #[tokio::test] + async fn transient_retries_until_budget_exhausted() { + let calls = Cell::new(0u32); + let result: Result<(), _> = retry_with_backoff(RetryConfig::no_wait(2), || { + calls.set(calls.get() + 1); + async { Err(ProviderError::Transient("5xx".into())) } + }) + .await; + assert!(matches!(result, Err(ProviderError::Transient(_)))); + // 1 initial + 2 retries = 3 total attempts. + assert_eq!(calls.get(), 3); + } + + /// Retries stop as soon as the closure succeeds. + #[tokio::test] + async fn retry_stops_on_first_success() { + let calls = Cell::new(0u32); + let result = retry_with_backoff(RetryConfig::no_wait(5), || { + let n = calls.get() + 1; + calls.set(n); + async move { + if n < 3 { + Err(ProviderError::Transient("flaky".into())) + } else { + Ok::(42) + } + } + }) + .await; + assert_eq!(result.unwrap(), 42); + assert_eq!(calls.get(), 3); + } + + /// RateLimit with no `retry_after` is retryable just like Transient. + #[tokio::test] + async fn rate_limit_without_hint_is_retried() { + let calls = Cell::new(0u32); + let _result: Result<(), _> = retry_with_backoff(RetryConfig::no_wait(1), || { + calls.set(calls.get() + 1); + async { + Err(ProviderError::RateLimit { + message: "slow down".into(), + retry_after: None, + }) + } + }) + .await; + assert_eq!(calls.get(), 2); + } + + #[test] + fn next_delay_uses_retry_after_when_present() { + let cfg = RetryConfig::interactive(); + let hint = Duration::from_secs(2); + // The Retry-After hint should win over the computed base. + assert_eq!(next_delay(&cfg, 0, Some(hint)), hint); + } + + #[test] + fn next_delay_caps_retry_after_at_max_delay() { + let cfg = RetryConfig { + max_delay: Duration::from_secs(5), + ..RetryConfig::interactive() + }; + let too_long = Duration::from_secs(600); + assert_eq!(next_delay(&cfg, 0, Some(too_long)), Duration::from_secs(5)); + } + + #[test] + fn next_delay_doubles_then_caps() { + // No jitter so the math is exact. + let cfg = RetryConfig { + max_retries: 5, + initial_delay: Duration::from_millis(100), + max_delay: Duration::from_millis(1_000), + jitter: 0.0, + }; + assert_eq!(next_delay(&cfg, 0, None), Duration::from_millis(100)); + assert_eq!(next_delay(&cfg, 1, None), Duration::from_millis(200)); + assert_eq!(next_delay(&cfg, 2, None), Duration::from_millis(400)); + assert_eq!(next_delay(&cfg, 3, None), Duration::from_millis(800)); + // Capped at max_delay. + assert_eq!(next_delay(&cfg, 4, None), Duration::from_millis(1_000)); + assert_eq!(next_delay(&cfg, 100, None), Duration::from_millis(1_000)); + } + + #[test] + fn jitter_zero_is_identity() { + let d = Duration::from_millis(500); + assert_eq!(apply_jitter(d, 0.0), d); + } + + #[test] + fn jitter_stays_within_window() { + let d = Duration::from_millis(1_000); + for _ in 0..1_000 { + let j = apply_jitter(d, 0.25); + assert!(j >= Duration::from_millis(750)); + assert!(j <= Duration::from_millis(1_250)); + } + } +}