newt_core/
retry.rs

1//! Shared HTTP retry/backoff for inference backends.
2//!
3//! Lived in `newt-inference` until Step 9.7 moved it here so the relocated
4//! agentic loop ([`crate::agentic`]) can share it without creating a
5//! `newt-inference` ⇄ `newt-core` dependency cycle; `newt_inference::retry`
6//! re-exports this module, so all existing import paths still work.
7//!
8//! Both `LocalOllamaBackend` and `LocalVllmBackend` (in `newt-inference`)
9//! drive their single `try_complete` attempt through [`with_backoff`], so the
10//! retry policy and the error-classification rules live in exactly one place.
11//!
12//! ## Why this exists
13//!
14//! Hosted OpenAI-compatible endpoints (e.g. NVIDIA's inference API) fail
15//! *intermittently* — transient connection resets and `429 Too Many Requests`
16//! under load. The previous per-backend loops only retried connection errors
17//! and `5xx`, so a `429` surfaced as a hard error, and the fixed
18//! `[250, 500, 1000]ms` schedule gave up after ~1.75s. [`RetryPolicy`] widens
19//! both: `408`/`429` are retryable, and the backoff is true exponential with a
20//! configurable ceiling and jitter.
21
22use std::future::Future;
23use std::time::Duration;
24
25/// Whether a failed attempt is worth retrying.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum Retryability {
28    /// Transient — retrying may succeed (connection failure, timeout, `408`,
29    /// `429`, or any `5xx`).
30    Retry,
31    /// Permanent for this request — retrying will not help (other `4xx`, or a
32    /// malformed success body).
33    Fatal,
34}
35
36/// Classify a backend error by its message.
37///
38/// Backends produce two error shapes that this inspects:
39/// - transport/timeout: `"<backend> request failed: <source>"`
40/// - HTTP status:        `"<backend> returned <code>: <body>"`
41///
42/// Anything else (e.g. a JSON decode error on a `200`) is [`Retryability::Fatal`].
43pub fn classify(err: &anyhow::Error) -> Retryability {
44    let msg = err.to_string();
45
46    // Transport-level failure from reqwest (connection refused, reset, DNS,
47    // or a client timeout). Always worth another attempt.
48    if msg.contains("request failed") {
49        return Retryability::Retry;
50    }
51
52    // HTTP status surfaced as "<backend> returned <code>: <body>".
53    if let Some(code) = status_code_in(&msg) {
54        if is_retryable_status(code) {
55            return Retryability::Retry;
56        }
57        return Retryability::Fatal;
58    }
59
60    Retryability::Fatal
61}
62
63/// `408 Request Timeout`, `429 Too Many Requests`, and every `5xx` are
64/// retryable; all other statuses are not.
65fn is_retryable_status(code: u16) -> bool {
66    code == 408 || code == 429 || (500..600).contains(&code)
67}
68
69/// Pull the first status code out of an error message.
70///
71/// Recognises two formats produced by the backends in this workspace:
72/// - `"<backend> returned <code> …"` — local backends (Ollama, vLLM)
73/// - `"inference endpoint <code> …"` — hosted OpenAI-compatible endpoints
74///   (NVIDIA inference API, LiteLLM proxies, etc.)
75///
76/// The code is the leading run of ASCII digits after the matched prefix
77/// (e.g. `StatusCode` Display is `"503 Service Unavailable"`, digits first).
78fn status_code_in(msg: &str) -> Option<u16> {
79    const PREFIXES: &[&str] = &["returned ", "inference endpoint "];
80    for prefix in PREFIXES {
81        if let Some(after) = msg.split_once(prefix).map(|(_, r)| r) {
82            if let Some(code) = after
83                .split(|c: char| !c.is_ascii_digit())
84                .find(|s: &&str| !s.is_empty())
85                .and_then(|s| s.parse().ok())
86            {
87                return Some(code);
88            }
89        }
90    }
91    None
92}
93
94/// Exponential-backoff retry policy.
95///
96/// `delay(attempt) = min(max, base * 2^(attempt-1))`, optionally spread with
97/// equal jitter (half fixed, half random) to avoid synchronized retries when
98/// several agents hammer the same hosted endpoint.
99#[derive(Debug, Clone)]
100pub struct RetryPolicy {
101    /// Number of retries *after* the first attempt. `0` disables retrying.
102    pub max_retries: u32,
103    /// Base delay before the first retry.
104    pub base: Duration,
105    /// Ceiling for any single backoff delay.
106    pub max: Duration,
107    /// Whether to apply equal jitter to each delay.
108    pub jitter: bool,
109}
110
111impl Default for RetryPolicy {
112    /// Production default: 4 retries (5 attempts total), 500 ms base doubling
113    /// to an 8 s ceiling, with jitter — roughly 15 s of resilience against a
114    /// flaky hosted endpoint.
115    fn default() -> Self {
116        Self {
117            max_retries: 4,
118            base: Duration::from_millis(500),
119            max: Duration::from_secs(8),
120            jitter: true,
121        }
122    }
123}
124
125impl RetryPolicy {
126    /// Build from environment variables, falling back to [`RetryPolicy::default`].
127    ///
128    /// Variables (all optional, unset/invalid values are silently ignored):
129    /// - `NEWT_HTTP_MAX_RETRIES` — retry count after the first attempt
130    /// - `NEWT_HTTP_BACKOFF_BASE_MS` — base delay in milliseconds
131    /// - `NEWT_HTTP_BACKOFF_MAX_MS` — ceiling delay in milliseconds
132    /// - `NEWT_HTTP_JITTER` — `0`/`false`/`off` disables jitter
133    pub fn from_env() -> Self {
134        Self::from_env_or(Self::default())
135    }
136
137    /// Like [`from_env`](Self::from_env) but starts from `base` instead of
138    /// [`Default`]. Use this when a caller has different default thresholds
139    /// from the crate default but still wants the env vars to override them.
140    pub fn from_env_or(mut base: Self) -> Self {
141        if let Some(n) = env_parse::<u32>("NEWT_HTTP_MAX_RETRIES") {
142            base.max_retries = n;
143        }
144        if let Some(ms) = env_parse::<u64>("NEWT_HTTP_BACKOFF_BASE_MS") {
145            base.base = Duration::from_millis(ms);
146        }
147        if let Some(ms) = env_parse::<u64>("NEWT_HTTP_BACKOFF_MAX_MS") {
148            base.max = Duration::from_millis(ms);
149        }
150        if let Ok(v) = std::env::var("NEWT_HTTP_JITTER") {
151            base.jitter = !matches!(
152                v.trim().to_ascii_lowercase().as_str(),
153                "0" | "false" | "off"
154            );
155        }
156        base
157    }
158
159    /// A policy tuned for slow, home-lab local inference endpoints (e.g. a DGX
160    /// that can drop for 30–60 s under load). More patient than the default
161    /// hosted-API policy:
162    /// - 6 retries (7 attempts total)
163    /// - 2 s base doubling to a 30 s ceiling, with jitter
164    /// - Total resilience window: ~90 s
165    ///
166    /// Env-var overrides (`NEWT_HTTP_MAX_RETRIES` etc.) still apply.
167    pub fn for_local_inference() -> Self {
168        Self::from_env_or(Self {
169            max_retries: 6,
170            base: Duration::from_secs(2),
171            max: Duration::from_secs(30),
172            jitter: true,
173        })
174    }
175
176    /// A zero-delay policy with `max_retries` retries — for tests that want to
177    /// exercise the retry loop without sleeping.
178    pub fn immediate(max_retries: u32) -> Self {
179        Self {
180            max_retries,
181            base: Duration::ZERO,
182            max: Duration::ZERO,
183            jitter: false,
184        }
185    }
186
187    /// Delay before the `attempt`-th retry (`attempt` is 1-based: `1` is the
188    /// first retry). Deterministic when `jitter` is false.
189    pub fn delay_for(&self, attempt: u32) -> Duration {
190        let capped_ms = self.base_delay_ms(attempt);
191        if !self.jitter || capped_ms == 0 {
192            return Duration::from_millis(capped_ms);
193        }
194        // Equal jitter: keep half fixed, randomize the other half. This never
195        // collapses to ~0 (unlike full jitter), so backoff still grows.
196        let half = capped_ms / 2;
197        let span = capped_ms - half; // == ceil(capped/2)
198        Duration::from_millis(half + jitter_u64() % (span + 1))
199    }
200
201    /// The deterministic (pre-jitter) capped exponential delay in ms.
202    fn base_delay_ms(&self, attempt: u32) -> u64 {
203        let shift = attempt.saturating_sub(1).min(31);
204        let factor = 1u64 << shift;
205        let raw = (self.base.as_millis() as u64).saturating_mul(factor);
206        raw.min(self.max.as_millis() as u64)
207    }
208}
209
210/// Drive a fallible async operation under `policy`, calling `on_retry` before
211/// each sleep so the caller can log or display a retry indicator.
212///
213/// `on_retry(attempt, delay)` is called synchronously before sleeping:
214/// `attempt` is 1-based (1 = first retry), `delay` is the sleep duration.
215///
216/// Calls `op` until it succeeds, the error is [`Retryability::Fatal`], or
217/// `policy.max_retries` is exhausted. On exhaustion the *last* error is
218/// returned.
219pub async fn with_backoff_notify<T, F, Fut, N>(
220    policy: &RetryPolicy,
221    mut op: F,
222    mut on_retry: N,
223) -> anyhow::Result<T>
224where
225    F: FnMut() -> Fut,
226    Fut: Future<Output = anyhow::Result<T>>,
227    N: FnMut(u32, Duration),
228{
229    let mut retries = 0u32;
230    loop {
231        match op().await {
232            Ok(value) => return Ok(value),
233            Err(err) => {
234                if classify(&err) == Retryability::Fatal || retries >= policy.max_retries {
235                    return Err(err);
236                }
237                retries += 1;
238                let delay = policy.delay_for(retries);
239                on_retry(retries, delay);
240                tracing::warn!(
241                    attempt = retries,
242                    delay_ms = delay.as_millis() as u64,
243                    error = %err,
244                    "retrying inference request"
245                );
246                tokio::time::sleep(delay).await;
247            }
248        }
249    }
250}
251
252/// Drive a fallible async operation under `policy`.
253///
254/// Convenience wrapper around [`with_backoff_notify`] with a no-op callback.
255/// Calls `op` until it succeeds, the error is [`Retryability::Fatal`], or
256/// `policy.max_retries` is exhausted — sleeping `policy.delay_for(attempt)`
257/// between attempts. On exhaustion the *last* error is returned (so the caller
258/// still sees e.g. the final `503`).
259pub async fn with_backoff<T, F, Fut>(policy: &RetryPolicy, op: F) -> anyhow::Result<T>
260where
261    F: FnMut() -> Fut,
262    Fut: Future<Output = anyhow::Result<T>>,
263{
264    with_backoff_notify(policy, op, |_, _| {}).await
265}
266
267/// Parse an environment variable, returning `None` if unset or unparseable.
268fn env_parse<T: std::str::FromStr>(key: &str) -> Option<T> {
269    std::env::var(key).ok()?.trim().parse().ok()
270}
271
272/// Cheap, dependency-free jitter source. Not cryptographic — it only needs to
273/// de-correlate retry timing across attempts and processes. Seeded once from
274/// the wall clock (entropy only, never used as a coordination primitive),
275/// then advanced by a SplitMix64 step + xorshift per call.
276fn jitter_u64() -> u64 {
277    use std::sync::atomic::{AtomicU64, Ordering};
278    static STATE: AtomicU64 = AtomicU64::new(0);
279    // One-time seed from the clock's sub-second noise. Concurrent first-callers
280    // may both seed; that's harmless — any nonzero start works.
281    if STATE.load(Ordering::Relaxed) == 0 {
282        let seed = std::time::SystemTime::now()
283            .duration_since(std::time::UNIX_EPOCH)
284            .map(|d| d.subsec_nanos() as u64)
285            .unwrap_or(0)
286            ^ 0x9E37_79B9_7F4A_7C15;
287        STATE.store(seed | 1, Ordering::Relaxed);
288    }
289    let mut x = STATE.fetch_add(0x9E37_79B9_7F4A_7C15, Ordering::Relaxed);
290    x ^= x >> 12;
291    x ^= x << 25;
292    x ^= x >> 27;
293    x.wrapping_mul(0x2545_F491_4F6C_DD1D)
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299    use std::cell::Cell;
300
301    fn err(msg: &str) -> anyhow::Error {
302        anyhow::anyhow!("{msg}")
303    }
304
305    #[test]
306    fn classify_transport_failure_is_retry() {
307        assert_eq!(
308            classify(&err("vLLM request failed: error sending request for url")),
309            Retryability::Retry
310        );
311        assert_eq!(
312            classify(&err("Ollama request failed: connection refused")),
313            Retryability::Retry
314        );
315    }
316
317    #[test]
318    fn classify_429_and_408_and_5xx_are_retry() {
319        // The core regression: 429 must be retryable.
320        assert_eq!(
321            classify(&err("vLLM returned 429 Too Many Requests: slow down")),
322            Retryability::Retry
323        );
324        assert_eq!(
325            classify(&err("vLLM returned 408 Request Timeout:")),
326            Retryability::Retry
327        );
328        assert_eq!(
329            classify(&err("Ollama returned 503 Service Unavailable: down")),
330            Retryability::Retry
331        );
332        assert_eq!(
333            classify(&err("vLLM returned 500 Internal Server Error: boom")),
334            Retryability::Retry
335        );
336    }
337
338    #[test]
339    fn classify_other_4xx_is_fatal() {
340        assert_eq!(
341            classify(&err("vLLM returned 400 Bad Request: nope")),
342            Retryability::Fatal
343        );
344        assert_eq!(
345            classify(&err("vLLM returned 404 Not Found:")),
346            Retryability::Fatal
347        );
348    }
349
350    #[test]
351    fn classify_unknown_shape_is_fatal() {
352        assert_eq!(
353            classify(&err("error decoding response body: expected value")),
354            Retryability::Fatal
355        );
356    }
357
358    #[test]
359    fn status_code_parsing() {
360        assert_eq!(
361            status_code_in("vLLM returned 503 Service Unavailable: x"),
362            Some(503)
363        );
364        assert_eq!(status_code_in("Ollama returned 429: y"), Some(429));
365        assert_eq!(status_code_in("no status here"), None);
366    }
367
368    #[test]
369    fn classify_inference_endpoint_5xx_is_retry() {
370        // Regression: hosted endpoints format errors as "inference endpoint <code> …"
371        // which previously fell through status_code_in and was classified Fatal.
372        assert_eq!(
373            classify(&err(
374                r#"inference endpoint 500 Internal Server Error: {"error":{"message":"instance_id not found"}}"#
375            )),
376            Retryability::Retry
377        );
378        assert_eq!(
379            classify(&err("inference endpoint 503 Service Unavailable: down")),
380            Retryability::Retry
381        );
382        assert_eq!(
383            classify(&err("inference endpoint 429 Too Many Requests: slow")),
384            Retryability::Retry
385        );
386    }
387
388    #[test]
389    fn status_code_parsing_inference_endpoint_format() {
390        assert_eq!(
391            status_code_in("inference endpoint 500 Internal Server Error: body"),
392            Some(500)
393        );
394        assert_eq!(
395            status_code_in("inference endpoint 429 Too Many Requests: body"),
396            Some(429)
397        );
398        assert_eq!(
399            status_code_in("inference endpoint 400 Bad Request: body"),
400            Some(400)
401        );
402    }
403
404    #[test]
405    fn base_delay_is_exponential_and_capped() {
406        let p = RetryPolicy {
407            max_retries: 10,
408            base: Duration::from_millis(500),
409            max: Duration::from_secs(8),
410            jitter: false,
411        };
412        assert_eq!(p.base_delay_ms(1), 500);
413        assert_eq!(p.base_delay_ms(2), 1000);
414        assert_eq!(p.base_delay_ms(3), 2000);
415        assert_eq!(p.base_delay_ms(4), 4000);
416        assert_eq!(p.base_delay_ms(5), 8000);
417        // Capped at max thereafter.
418        assert_eq!(p.base_delay_ms(6), 8000);
419        assert_eq!(p.base_delay_ms(30), 8000);
420    }
421
422    #[test]
423    fn delay_without_jitter_is_deterministic() {
424        let p = RetryPolicy {
425            jitter: false,
426            ..RetryPolicy::default()
427        };
428        assert_eq!(p.delay_for(1), Duration::from_millis(500));
429        assert_eq!(p.delay_for(2), Duration::from_millis(1000));
430    }
431
432    #[test]
433    fn jittered_delay_stays_within_equal_jitter_band() {
434        let p = RetryPolicy {
435            max_retries: 4,
436            base: Duration::from_millis(1000),
437            max: Duration::from_secs(8),
438            jitter: true,
439        };
440        // attempt 1 → capped 1000ms → band [500, 1000].
441        for _ in 0..200 {
442            let ms = p.delay_for(1).as_millis() as u64;
443            assert!((500..=1000).contains(&ms), "delay {ms} outside [500,1000]");
444        }
445    }
446
447    #[test]
448    fn immediate_policy_never_sleeps() {
449        let p = RetryPolicy::immediate(3);
450        assert_eq!(p.max_retries, 3);
451        assert_eq!(p.delay_for(1), Duration::ZERO);
452        assert_eq!(p.delay_for(3), Duration::ZERO);
453    }
454
455    #[tokio::test]
456    async fn with_backoff_succeeds_after_transient_failures() {
457        let calls = Cell::new(0u32);
458        let result: anyhow::Result<&str> = with_backoff(&RetryPolicy::immediate(5), || {
459            let n = calls.get() + 1;
460            calls.set(n);
461            async move {
462                if n < 3 {
463                    Err(err("vLLM returned 429 Too Many Requests: slow"))
464                } else {
465                    Ok("ok")
466                }
467            }
468        })
469        .await;
470        assert_eq!(result.unwrap(), "ok");
471        assert_eq!(calls.get(), 3, "should have retried twice then succeeded");
472    }
473
474    #[tokio::test]
475    async fn with_backoff_stops_on_fatal() {
476        let calls = Cell::new(0u32);
477        let result: anyhow::Result<&str> = with_backoff(&RetryPolicy::immediate(5), || {
478            calls.set(calls.get() + 1);
479            async move { Err(err("vLLM returned 400 Bad Request: nope")) }
480        })
481        .await;
482        assert!(result.is_err());
483        assert_eq!(calls.get(), 1, "fatal error must not retry");
484    }
485
486    #[tokio::test]
487    async fn with_backoff_gives_up_after_max_retries() {
488        let calls = Cell::new(0u32);
489        let result: anyhow::Result<&str> = with_backoff(&RetryPolicy::immediate(3), || {
490            calls.set(calls.get() + 1);
491            async move { Err(err("vLLM returned 503 Service Unavailable: down")) }
492        })
493        .await;
494        let e = result.unwrap_err();
495        assert!(e.to_string().contains("503"), "last error preserved: {e}");
496        assert_eq!(calls.get(), 4, "1 initial + 3 retries");
497    }
498
499    #[tokio::test]
500    async fn with_backoff_notify_fires_callback_before_each_sleep() {
501        let calls = Cell::new(0u32);
502        let notified = Cell::new(0u32);
503        let result: anyhow::Result<&str> = with_backoff_notify(
504            &RetryPolicy::immediate(3),
505            || {
506                let n = calls.get() + 1;
507                calls.set(n);
508                async move {
509                    if n <= 2 {
510                        Err(err("vLLM returned 503 Service Unavailable: down"))
511                    } else {
512                        Ok("ok")
513                    }
514                }
515            },
516            |_, _| {
517                notified.set(notified.get() + 1);
518            },
519        )
520        .await;
521        assert_eq!(result.unwrap(), "ok");
522        assert_eq!(calls.get(), 3, "two retries then success");
523        assert_eq!(
524            notified.get(),
525            2,
526            "callback fired once before each retry sleep"
527        );
528    }
529
530    #[test]
531    fn for_local_inference_is_more_patient_than_default() {
532        let local = RetryPolicy::for_local_inference();
533        let default = RetryPolicy::default();
534        assert!(
535            local.max_retries > default.max_retries,
536            "local policy must allow more retries"
537        );
538        assert!(
539            local.max > default.max,
540            "local policy must have a longer backoff ceiling"
541        );
542    }
543
544    #[test]
545    fn from_env_or_starts_from_provided_base() {
546        // No env vars set — result must equal the base.
547        let base = RetryPolicy {
548            max_retries: 9,
549            base: Duration::from_secs(3),
550            max: Duration::from_secs(60),
551            jitter: false,
552        };
553        let result = RetryPolicy::from_env_or(base.clone());
554        assert_eq!(result.max_retries, 9);
555        assert_eq!(result.base, Duration::from_secs(3));
556        assert_eq!(result.max, Duration::from_secs(60));
557        assert!(!result.jitter);
558    }
559}
newt_core/retry.rs

newt_core/
retry.rs
