fix(sampler): retry Cloudflare 52x and other 5xx during outages

Port of upstream grok-build 1.0.0/1.0.1-window retry policy.

- is_retryable_api_status: 429 plus any 5xx except Cloudflare 525/526
  (origin TLS handshake / invalid certificate — never clears on its
  own), replacing the 500|502|503|504|520 allowlist so the 521-524,
  529, and 530 edge pages an outage serves actually retry.
- The generic retry path clamps a server Retry-After to
  MAX_RETRY_BACKOFF (30s) and jitters it: Cloudflare 52x pages ask for
  60-120s, which across the budget stalled a turn ~28 minutes. The 429
  path keeps the full (parse-capped) Retry-After under its low attempt
  cap, and jitter is shared via jittered().

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KSy9S31pZNWXtRZNVJ1sDf

crates/codegen/kigi-sampler/src/retry.rs

@@ -5,8 +5,11 @@
 //!
 //! # Retry behavior summary
 //!
-//! **Retried** (up to [`DEFAULT_MAX_RETRIES`] = 15, ~6 min with 30s backoff cap):
-//! - 500, 502, 503, 504, 520 (server errors)
+//! **Retried** (up to [`DEFAULT_MAX_RETRIES`] = 15, ~6 min with the
+//! [`MAX_RETRY_BACKOFF`] cap — a server `Retry-After` is clamped to it too):
+//! - Any 5xx except Cloudflare 525/526 — covers the Cloudflare edge pages
+//!   (520-524 origin unreachable/timed out, 530 edge 1xxx) and upstream
+//!   overload (529); see `is_retryable_api_status`
 //! - Connection errors (timeout, refused, reset)
 //! - `EventStreamError` / `StreamError` (mid-stream failures)
 //! - `EmptyResponse` (model returned no content/tool calls)
@@ -77,16 +80,35 @@ pub fn doom_loop_backoff(retry_count: u32) -> Duration {
     Duration::from_millis(hasher.finish() % 251)
 }
 
-/// Exponential backoff (2s, 4s, 8s, ..., capped 30s) with +/-20% jitter
-/// to prevent thundering-herd retry storms.
+/// Longest single wait on the generic retry path — the exponential-backoff
+/// ceiling, and the clamp for a server `Retry-After`. Cloudflare answers 52x
+/// with `Retry-After: 60`-`120`; honoring that verbatim across the retry
+/// budget would stall a turn ~28 min instead of the ~6 min budget. The 429
+/// path deliberately waits the full `Retry-After` instead, bounded by
+/// [`RATE_LIMIT_RETRY_THRESHOLD`] attempts (and by the parse-level 120s cap
+/// on the header).
+pub const MAX_RETRY_BACKOFF: Duration = Duration::from_secs(30);
+
+/// Exponential backoff (2s, 4s, 8s, ..., capped at [`MAX_RETRY_BACKOFF`])
+/// with +/-20% jitter to prevent thundering-herd retry storms.
 pub fn retry_backoff_with_jitter(retry_count: u32) -> Duration {
+    let shift = retry_count.saturating_sub(1);
+    let base_ms = 2000u64
+        .checked_shl(shift)
+        .unwrap_or(u64::MAX)
+        .min(MAX_RETRY_BACKOFF.as_millis() as u64);
+    jittered(Duration::from_millis(base_ms))
+}
+
+/// +/-20% jitter around `base`, de-syncing clients that failed at the
+/// same instant (e.g. a mass Cloudflare 52x event during an origin outage).
+fn jittered(base: Duration) -> Duration {
     use std::hash::{Hash, Hasher};
     use std::sync::atomic::{AtomicU64, Ordering};
 
     static JITTER_SEQ: AtomicU64 = AtomicU64::new(0);
 
-    let shift = retry_count.saturating_sub(1);
-    let base_ms = 2000u64.checked_shl(shift).unwrap_or(u64::MAX).min(30_000);
+    let base_ms = base.as_millis() as u64;
     let jitter_range = base_ms / 5;
     let mut hasher = std::hash::DefaultHasher::new();
     JITTER_SEQ.fetch_add(1, Ordering::Relaxed).hash(&mut hasher);
@@ -210,9 +232,11 @@ pub fn classify_error(
         if next_attempt >= max_retries {
             return RetryDecision::Fatal(clone_error(err));
         }
+        // A server `Retry-After` is clamped: Cloudflare 52x pages ask for
+        // 60-120s, which across the budget would stall a turn ~28 min.
         let backoff = err
             .retry_after()
-            .map(Duration::from_secs)
+            .map(|secs| jittered(Duration::from_secs(secs).min(MAX_RETRY_BACKOFF)))
             .unwrap_or_else(|| retry_backoff_with_jitter(next_attempt));
         if next_attempt == 1 {
             return RetryDecision::RetryWithClientRebuild { backoff };
@@ -595,6 +619,22 @@ mod tests {
         }
     }
 
+    #[test]
+    fn classify_5xx_retry_after_is_clamped_to_backoff_cap() {
+        // Cloudflare answers 52x with Retry-After: 60-120; honoring it
+        // verbatim across the retry budget would stall a turn ~28 min.
+        let err = api_err_with_retry_after(StatusCode::from_u16(522).unwrap(), 120);
+        match classify_error(&err, 1, 15, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::Retry { backoff } => {
+                assert!(
+                    backoff <= MAX_RETRY_BACKOFF.mul_f64(1.2),
+                    "52x Retry-After must clamp to ~MAX_RETRY_BACKOFF, got {backoff:?}"
+                );
+            }
+            other => panic!("expected Retry, got {other:?}"),
+        }
+    }
+
     #[test]
     fn classify_5xx_first_retry_rebuilds_client() {
         let err = api_err(StatusCode::INTERNAL_SERVER_ERROR, "boom");

crates/codegen/kigi-sampling-types/src/error.rs

@@ -239,9 +239,7 @@ impl SamplingError {
             SamplingError::InvalidConfiguration(_) => false,
             SamplingError::Http(err) => is_retryable_reqwest(err),
             SamplingError::Serialization(_) => false,
-            SamplingError::Api { status, .. } => {
-                matches!(status.as_u16(), 429 | 500 | 502 | 503 | 504 | 520)
-            }
+            SamplingError::Api { status, .. } => is_retryable_api_status(status.as_u16()),
             SamplingError::EventStreamError(_) => true,
             SamplingError::StreamError { .. } => true,
             SamplingError::IdleTimeout { .. } => false,
@@ -342,6 +340,15 @@ pub fn try_parse_stream_error(data: &str) -> Option<SamplingError> {
     })
 }
 
+/// 429 and any 5xx except Cloudflare 525/526 (origin TLS handshake /
+/// invalid certificate — a broken origin certificate never clears on its
+/// own). The 5xx sweep covers the Cloudflare edge pages (520-524 origin
+/// unreachable/timed out, 530 edge 1xxx) and upstream overload (529),
+/// which an outage serves in place of the API's own statuses.
+pub fn is_retryable_api_status(status: u16) -> bool {
+    status == 429 || ((500..=599).contains(&status) && !matches!(status, 525 | 526))
+}
+
 /// True when an error message indicates a context-window overflow. Backends report
 /// this inconsistently with no stable error code, so we match the message text; it's
 /// deterministic (re-sending the same payload always fails), so callers must not retry.
@@ -379,6 +386,25 @@ pub fn is_retryable_reqwest(err: &reqwest::Error) -> bool {
 mod tests {
     use super::*;
 
+    #[test]
+    fn retryable_api_status_covers_edge_5xx_and_excludes_broken_tls() {
+        for status in [429u16, 500, 502, 503, 504, 520, 521, 522, 523, 524, 529, 530, 599] {
+            assert!(is_retryable_api_status(status), "should retry {status}");
+        }
+        // Cloudflare 525/526: origin TLS handshake / invalid certificate —
+        // a broken origin certificate never clears on its own.
+        for status in [525u16, 526, 400, 401, 403, 404, 408, 422, 200] {
+            assert!(!is_retryable_api_status(status), "must not retry {status}");
+        }
+        let api = SamplingError::Api {
+            status: StatusCode::from_u16(522).unwrap(),
+            message: "cloudflare: connection timed out".into(),
+            model_metadata: None,
+            retry_after_secs: None,
+        };
+        assert!(api.is_retryable(), "522 edge page must be retryable");
+    }
+
     #[test]
     fn context_length_error_matches_backend_variants() {
         for msg in [