sie_sdk/retry/
mod.rs

1//! The retry state machine shared by every endpoint.
2//!
3//! The Python SDK repeats a near-identical loop in each method, with small per-endpoint
4//! differences buried inside it. Here those differences are a policy struct and the loop
5//! exists once.
6
7pub mod backoff;
8
9use std::time::{Duration, Instant};
10
11use crate::error::{Error, Result, TransportErrorKind, codes};
12use crate::http::HttpResponse;
13use crate::wire;
14
15/// Per-call knobs every request builder exposes.
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct RequestOptions {
18    /// Machine profile, optionally prefixed with a pool: `"pool/l4"` or `"l4"`.
19    pub gpu: Option<String>,
20    /// Whether to wait out provisioning and transport hiccups, or fail fast.
21    pub wait_for_capacity: bool,
22    /// Total wall-clock budget for the call, retries included.
23    pub provision_timeout: Duration,
24    /// Cap on `RESOURCE_EXHAUSTED` retries. Zero fails fast.
25    pub max_oom_retries: u32,
26}
27
28impl Default for RequestOptions {
29    fn default() -> Self {
30        Self {
31            gpu: None,
32            wait_for_capacity: true,
33            provision_timeout: backoff::DEFAULT_PROVISION_TIMEOUT,
34            max_oom_retries: backoff::RESOURCE_EXHAUSTED_MAX_RETRIES,
35        }
36    }
37}
38
39/// Which retry branches an endpoint participates in.
40///
41/// The asymmetries are deliberate, not accidents: generation is non-idempotent, so it
42/// never replays a request that may already have reached a worker.
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub(crate) struct RetryPolicy {
45    /// Poll through `LORA_LOADING`. Only `/v1/encode` can receive it.
46    pub lora_loading: bool,
47    /// Whether `MODEL_LOADING` polling requires `wait_for_capacity`.
48    pub model_loading_gated_on_wait: bool,
49    /// Whether a 504 may be replayed. Only true for idempotent endpoints.
50    pub retry_gateway_timeout: bool,
51    /// Whether a failure before the connection was established may be replayed.
52    pub retry_connect: bool,
53    /// Whether a failure after the request was written may be replayed.
54    pub retry_midflight_transport: bool,
55    /// Whether `RESOURCE_EXHAUSTED` retries require `wait_for_capacity`.
56    pub oom_gated_on_wait: bool,
57    /// Whether a 503 capacity code means the request could not be priced at all.
58    pub estimate_unroutable: bool,
59}
60
61impl RetryPolicy {
62    /// `/v1/score`, `/v1/extract`: idempotent, so everything is replayable.
63    pub(crate) const INFERENCE: Self = Self {
64        lora_loading: false,
65        model_loading_gated_on_wait: false,
66        retry_gateway_timeout: true,
67        retry_connect: true,
68        retry_midflight_transport: true,
69        oom_gated_on_wait: false,
70        estimate_unroutable: false,
71    };
72
73    /// `/v1/encode`: as [`Self::INFERENCE`], plus `LoRA` adapter polling.
74    pub(crate) const ENCODE: Self = Self {
75        lora_loading: true,
76        ..Self::INFERENCE
77    };
78
79    /// `/v1/generate` buffered: non-idempotent, and model loading is opt-in.
80    pub(crate) const GENERATE: Self = Self {
81        lora_loading: false,
82        model_loading_gated_on_wait: true,
83        retry_gateway_timeout: false,
84        retry_connect: true,
85        retry_midflight_transport: false,
86        oom_gated_on_wait: false,
87        estimate_unroutable: false,
88    };
89
90    /// `/v1/chat/completions`, `/v1/responses`, and every SSE stream.
91    pub(crate) const STREAM: Self = Self {
92        lora_loading: false,
93        model_loading_gated_on_wait: false,
94        retry_gateway_timeout: false,
95        retry_connect: true,
96        retry_midflight_transport: false,
97        oom_gated_on_wait: true,
98        estimate_unroutable: false,
99    };
100
101    /// Endpoints with no capacity semantics at all (models, health, files, jobs, pools).
102    ///
103    /// These are single-shot: a metadata call against a server that is down must fail in
104    /// seconds, not sit inside a provisioning budget.
105    pub(crate) const NONE: Self = Self {
106        lora_loading: false,
107        model_loading_gated_on_wait: false,
108        retry_gateway_timeout: false,
109        retry_connect: false,
110        retry_midflight_transport: false,
111        oom_gated_on_wait: false,
112        estimate_unroutable: false,
113    };
114
115    /// `/v1/estimate`: a capacity code means the request cannot be priced, not that it
116    /// should be retried.
117    pub(crate) const ESTIMATE: Self = Self {
118        estimate_unroutable: true,
119        ..Self::NONE
120    };
121}
122
123/// What the caller should do with a response the retry machine has inspected.
124#[derive(Debug, Clone, Copy, PartialEq, Eq)]
125pub(crate) enum Decision {
126    /// The response is terminal and successful; hand it to the parser.
127    Accept,
128    /// Sleep for this long, then replay the request.
129    Retry(Duration),
130}
131
132/// Mutable state for one logical call, spanning all its attempts.
133#[derive(Debug)]
134pub(crate) struct RetryState {
135    policy: RetryPolicy,
136    wait_for_capacity: bool,
137    max_oom_retries: u32,
138    budget: Duration,
139    start: Instant,
140    gpu: Option<String>,
141    model: Option<String>,
142    retries: u32,
143    oom_retries: u32,
144    lora_retries: u32,
145}
146
147impl RetryState {
148    pub(crate) fn new(policy: RetryPolicy, options: &RequestOptions, model: Option<&str>) -> Self {
149        Self {
150            policy,
151            wait_for_capacity: options.wait_for_capacity,
152            max_oom_retries: options.max_oom_retries,
153            budget: options.provision_timeout,
154            start: Instant::now(),
155            gpu: options.gpu.clone(),
156            model: model.map(str::to_string),
157            retries: 0,
158            oom_retries: 0,
159            lora_retries: 0,
160        }
161    }
162
163    /// Retries performed so far, reported back through [`crate::types::RequestMetadata`].
164    pub(crate) fn retries(&self) -> u32 {
165        self.retries
166    }
167
168    fn elapsed(&self) -> Duration {
169        self.start.elapsed()
170    }
171
172    fn remaining(&self) -> Duration {
173        self.budget.saturating_sub(self.elapsed())
174    }
175
176    /// The timeout for the next attempt, or the terminal error when the budget is spent.
177    pub(crate) fn attempt_timeout(&self, client_timeout: Duration) -> Result<Duration> {
178        let remaining = self.remaining();
179        if remaining.is_zero() {
180            return Err(Error::Provisioning {
181                message: format!(
182                    "Provision timeout ({:.1}s) exceeded before request could be sent",
183                    self.budget.as_secs_f64()
184                ),
185                gpu: self.gpu.clone(),
186                retry_after: None,
187            });
188        }
189        Ok(client_timeout.min(remaining))
190    }
191
192    /// Classify and act on a transport-level failure.
193    pub(crate) fn on_transport_error(
194        &mut self,
195        error: &reqwest::Error,
196        base_url: &str,
197    ) -> Result<Duration> {
198        let kind = classify_transport_error(error);
199        let retryable = self.wait_for_capacity
200            && match kind {
201                TransportErrorKind::Connect => {
202                    self.policy.retry_connect && is_transient_connect_error(error)
203                }
204                TransportErrorKind::Timeout | TransportErrorKind::MidFlight => {
205                    self.policy.retry_midflight_transport
206                }
207            };
208
209        if retryable && let Some(delay) = backoff::transport_delay(self.elapsed(), self.budget) {
210            tracing::debug!(
211                "{} retrying in {:.1}s (elapsed: {:.1}s, timeout: {:.1}s): {error}",
212                transport_label(kind),
213                delay.as_secs_f64(),
214                self.elapsed().as_secs_f64(),
215                self.budget.as_secs_f64(),
216            );
217            self.retries += 1;
218            return Ok(delay);
219        }
220
221        let message = match kind {
222            TransportErrorKind::Connect => format!("Failed to connect to {base_url}: {error}"),
223            TransportErrorKind::Timeout => format!("Request timed out: {error}"),
224            TransportErrorKind::MidFlight => format!(
225                "Connection lost mid-request; the peer closed the connection before sending a \
226                 complete response: {error}"
227            ),
228        };
229        Err(Error::Connection {
230            message,
231            kind,
232            source: None,
233        })
234    }
235
236    /// Inspect a response and decide whether to accept, retry, or fail.
237    ///
238    /// Only `>= 400` responses reach the error branches; a 2xx or 3xx is accepted and left
239    /// for the endpoint's parser.
240    pub(crate) fn on_response(&mut self, response: &HttpResponse) -> Result<Decision> {
241        // Terminal load failures short-circuit before any budget is engaged: retrying a
242        // gated repo or a missing dependency wastes the whole provision window.
243        wire::check_model_load_failed(response, self.model.as_deref(), self.retries)?;
244        if self.policy.estimate_unroutable {
245            wire::check_estimate_unroutable(response, self.retries)?;
246        }
247
248        if response.status == 503
249            && let Some(delay) = self.on_service_unavailable(response)?
250        {
251            self.retries += 1;
252            return Ok(Decision::Retry(delay));
253        }
254
255        if response.status == 504 {
256            if let Some(delay) = self.on_gateway_timeout(response) {
257                self.retries += 1;
258                return Ok(Decision::Retry(delay));
259            }
260            if !self.policy.retry_gateway_timeout {
261                return Err(Error::Server {
262                    message: "Gateway timed out (504) after the request was published to the queue; \
263                              a worker may already be generating. Not retried because generation is \
264                              non-idempotent (retrying could double-bill)."
265                        .to_string(),
266                    code: wire::error_code(response),
267                    status: 504,
268                    request: crate::http::metadata::parse(&response.headers, None, self.retries).map(Box::new),
269                });
270            }
271        }
272
273        if response.status >= 400 {
274            return Err(wire::handle_error(
275                response,
276                self.model.as_deref(),
277                self.retries,
278            ));
279        }
280        Ok(Decision::Accept)
281    }
282
283    /// The 503 branches. `Ok(None)` means "not a capacity code, fall through to the
284    /// terminal handler".
285    fn on_service_unavailable(&mut self, response: &HttpResponse) -> Result<Option<Duration>> {
286        let Some(code) = wire::error_code(response) else {
287            return Ok(None);
288        };
289        let hint = backoff::retry_after(&response.headers);
290
291        match code.as_str() {
292            codes::PROVISIONING => self.provisioning_delay(hint).map(Some),
293            codes::LORA_LOADING if self.policy.lora_loading => self.lora_delay(hint).map(Some),
294            codes::MODEL_LOADING => {
295                if self.policy.model_loading_gated_on_wait && !self.wait_for_capacity {
296                    return Ok(None);
297                }
298                self.model_loading_delay(hint).map(Some)
299            }
300            codes::RESOURCE_EXHAUSTED => {
301                if self.policy.oom_gated_on_wait && !self.wait_for_capacity {
302                    return Err(self.resource_exhausted(response));
303                }
304                self.oom_delay(response, hint).map(Some)
305            }
306            _ => Ok(None),
307        }
308    }
309
310    fn provisioning_delay(&self, hint: Option<Duration>) -> Result<Duration> {
311        let gpu_label = self.gpu.as_deref().unwrap_or("default");
312        if !self.wait_for_capacity {
313            return Err(Error::Provisioning {
314                message: format!(
315                    "No capacity available for GPU '{gpu_label}'. Server is provisioning."
316                ),
317                gpu: self.gpu.clone(),
318                retry_after: hint,
319            });
320        }
321        let remaining = self.remaining();
322        if remaining.is_zero() {
323            return Err(Error::Provisioning {
324                message: format!(
325                    "Provisioning timeout after {:.1}s waiting for GPU '{gpu_label}'",
326                    self.elapsed().as_secs_f64()
327                ),
328                gpu: self.gpu.clone(),
329                retry_after: hint,
330            });
331        }
332        // A server hint is honoured verbatim; only the SDK's own default is jittered.
333        Ok(match hint {
334            Some(hint) => hint.min(remaining),
335            None => backoff::apply_jitter(backoff::DEFAULT_RETRY_DELAY.min(remaining)),
336        })
337    }
338
339    fn model_loading_delay(&self, hint: Option<Duration>) -> Result<Duration> {
340        let remaining = self.remaining();
341        if remaining.is_zero() {
342            return Err(Error::ModelLoading {
343                message: format!(
344                    "Model loading timeout after {:.1}s for '{}'",
345                    self.elapsed().as_secs_f64(),
346                    self.model.as_deref().unwrap_or("model")
347                ),
348                model: self.model.clone(),
349            });
350        }
351        Ok(backoff::retry_after_or(hint, backoff::MODEL_LOADING_DELAY).min(remaining))
352    }
353
354    fn lora_delay(&mut self, hint: Option<Duration>) -> Result<Duration> {
355        self.lora_retries += 1;
356        let remaining = self.remaining();
357        if self.lora_retries > backoff::LORA_LOADING_MAX_RETRIES || remaining.is_zero() {
358            return Err(Error::LoraLoading {
359                message: format!(
360                    "LoRA loading timeout after {} retries",
361                    self.lora_retries - 1
362                ),
363                lora: None,
364                model: self.model.clone(),
365            });
366        }
367        // Clamped to the remaining budget, unlike the Python SDK, which can overshoot
368        // `provision_timeout_s` by up to ten seconds on this branch.
369        Ok(backoff::retry_after_or(hint, backoff::LORA_LOADING_DELAY).min(remaining))
370    }
371
372    fn oom_delay(&mut self, response: &HttpResponse, hint: Option<Duration>) -> Result<Duration> {
373        let remaining = self.remaining();
374        if self.oom_retries >= self.max_oom_retries || remaining.is_zero() {
375            return Err(self.resource_exhausted(response));
376        }
377        let delay = backoff::oom_backoff(hint, self.oom_retries);
378        // Sleeping the whole remaining budget would surface a timeout instead of the root
379        // cause, so report the exhaustion now.
380        if delay >= remaining {
381            return Err(self.resource_exhausted(response));
382        }
383        if self.oom_retries == 0 {
384            tracing::warn!(
385                "Server resource exhausted, retrying in {:.1}s (attempt 1/{}, elapsed: {:.1}s, timeout: {:.1}s)",
386                delay.as_secs_f64(),
387                self.max_oom_retries,
388                self.elapsed().as_secs_f64(),
389                self.budget.as_secs_f64(),
390            );
391        } else {
392            tracing::info!(
393                "Server resource exhausted, retrying in {:.1}s (attempt {}/{})",
394                delay.as_secs_f64(),
395                self.oom_retries + 1,
396                self.max_oom_retries,
397            );
398        }
399        self.oom_retries += 1;
400        Ok(delay)
401    }
402
403    fn resource_exhausted(&self, response: &HttpResponse) -> Error {
404        Error::ResourceExhausted {
405            message: format!(
406                "Server resource exhausted after {} retry attempt(s) for model '{}'",
407                self.oom_retries,
408                self.model.as_deref().unwrap_or("unknown")
409            ),
410            model: self.model.clone(),
411            retries: self.oom_retries,
412            request: crate::http::metadata::parse(&response.headers, None, self.retries)
413                .map(Box::new),
414        }
415    }
416
417    fn on_gateway_timeout(&self, response: &HttpResponse) -> Option<Duration> {
418        if !self.policy.retry_gateway_timeout || !self.wait_for_capacity {
419            return None;
420        }
421        let remaining = self.remaining();
422        if remaining.is_zero() {
423            return None;
424        }
425        let hint = backoff::retry_after(&response.headers);
426        let delay = backoff::retry_after_or(hint, backoff::MODEL_LOADING_DELAY).min(remaining);
427        tracing::info!(
428            "Gateway timeout (504), retrying in {:.1}s (elapsed: {:.1}s, timeout: {:.1}s)",
429            delay.as_secs_f64(),
430            self.elapsed().as_secs_f64(),
431            self.budget.as_secs_f64(),
432        );
433        Some(delay)
434    }
435}
436
437fn transport_label(kind: TransportErrorKind) -> &'static str {
438    match kind {
439        TransportErrorKind::Connect => "Connect error,",
440        TransportErrorKind::Timeout => "Request timeout,",
441        TransportErrorKind::MidFlight => "Transient transport error,",
442    }
443}
444
445pub(crate) fn classify_transport_error(error: &reqwest::Error) -> TransportErrorKind {
446    if error.is_timeout() {
447        TransportErrorKind::Timeout
448    } else if error.is_connect() {
449        TransportErrorKind::Connect
450    } else {
451        TransportErrorKind::MidFlight
452    }
453}
454
455/// Errnos that describe a peer or network that may recover on its own.
456const TRANSIENT_ERRNOS: &[i32] = &[
457    libc_errno::ECONNREFUSED,
458    libc_errno::ECONNRESET,
459    libc_errno::ETIMEDOUT,
460    libc_errno::EHOSTUNREACH,
461    libc_errno::ENETUNREACH,
462    libc_errno::ENETDOWN,
463    libc_errno::EHOSTDOWN,
464];
465
466/// Whether a connect-time failure is worth replaying.
467///
468/// Walks the source chain looking for an OS errno. A TLS handshake failure surfaces as
469/// `InvalidData` with no errno and is never transient: the peer's certificate or protocol
470/// will not fix itself within a provision window. When nothing conclusive is found the
471/// answer is "transient", matching the Python SDK, which defaults to retrying on platforms
472/// that do not surface an errno.
473pub(crate) fn is_transient_connect_error(error: &reqwest::Error) -> bool {
474    let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error);
475    while let Some(current) = source {
476        if let Some(io_error) = current.downcast_ref::<std::io::Error>() {
477            if let Some(errno) = io_error.raw_os_error() {
478                return TRANSIENT_ERRNOS.contains(&errno);
479            }
480            if io_error.kind() == std::io::ErrorKind::InvalidData {
481                return false;
482            }
483        }
484        source = current.source();
485    }
486    true
487}
488
489/// The handful of errno values the SDK cares about, without pulling in a libc dependency.
490#[allow(non_snake_case)]
491mod libc_errno {
492    pub const ECONNREFUSED: i32 = 111;
493    pub const ECONNRESET: i32 = 104;
494    pub const ETIMEDOUT: i32 = 110;
495    pub const EHOSTUNREACH: i32 = 113;
496    pub const ENETUNREACH: i32 = 101;
497    pub const ENETDOWN: i32 = 100;
498    pub const EHOSTDOWN: i32 = 112;
499}
500
501#[cfg(test)]
502mod tests {
503    use super::*;
504    use bytes::Bytes;
505    use reqwest::header::HeaderMap;
506
507    fn response(status: u16, code: Option<&str>, retry_after: Option<&str>) -> HttpResponse {
508        let mut headers = HeaderMap::new();
509        headers.insert(
510            reqwest::header::CONTENT_TYPE,
511            "application/json".parse().unwrap(),
512        );
513        if let Some(value) = retry_after {
514            headers.insert(reqwest::header::RETRY_AFTER, value.parse().unwrap());
515        }
516        let body = match code {
517            Some(code) => format!(r#"{{"error": {{"code": "{code}", "message": "m"}}}}"#),
518            None => "{}".to_string(),
519        };
520        HttpResponse {
521            status,
522            headers,
523            body: Bytes::from(body),
524        }
525    }
526
527    fn state(policy: RetryPolicy, options: RequestOptions) -> RetryState {
528        RetryState::new(policy, &options, Some("bge-m3"))
529    }
530
531    #[test]
532    fn success_is_accepted() {
533        let mut state = state(RetryPolicy::ENCODE, RequestOptions::default());
534        assert_eq!(
535            state.on_response(&response(200, None, None)).unwrap(),
536            Decision::Accept
537        );
538        assert_eq!(state.retries(), 0);
539    }
540
541    #[test]
542    fn provisioning_retries_and_honours_the_hint() {
543        let mut state = state(RetryPolicy::ENCODE, RequestOptions::default());
544        let decision = state
545            .on_response(&response(503, Some(codes::PROVISIONING), Some("7")))
546            .unwrap();
547        assert_eq!(decision, Decision::Retry(Duration::from_secs(7)));
548        assert_eq!(state.retries(), 1);
549    }
550
551    #[test]
552    fn provisioning_fails_fast_without_wait_for_capacity() {
553        let mut state = state(
554            RetryPolicy::ENCODE,
555            RequestOptions {
556                wait_for_capacity: false,
557                ..Default::default()
558            },
559        );
560        let err = state
561            .on_response(&response(503, Some(codes::PROVISIONING), Some("5")))
562            .unwrap_err();
563        assert!(matches!(err, Error::Provisioning { .. }));
564        assert_eq!(err.retry_after(), Some(Duration::from_secs(5)));
565    }
566
567    #[test]
568    fn lora_loading_only_applies_to_encode() {
569        let mut encode = state(RetryPolicy::ENCODE, RequestOptions::default());
570        assert_eq!(
571            encode
572                .on_response(&response(503, Some(codes::LORA_LOADING), None))
573                .unwrap(),
574            Decision::Retry(backoff::LORA_LOADING_DELAY)
575        );
576
577        let mut score = state(RetryPolicy::INFERENCE, RequestOptions::default());
578        let err = score
579            .on_response(&response(503, Some(codes::LORA_LOADING), None))
580            .unwrap_err();
581        assert!(matches!(err, Error::Server { status: 503, .. }));
582    }
583
584    #[test]
585    fn lora_loading_gives_up_after_ten_retries() {
586        let mut state = state(RetryPolicy::ENCODE, RequestOptions::default());
587        for _ in 0..backoff::LORA_LOADING_MAX_RETRIES {
588            state
589                .on_response(&response(503, Some(codes::LORA_LOADING), None))
590                .unwrap();
591        }
592        let err = state
593            .on_response(&response(503, Some(codes::LORA_LOADING), None))
594            .unwrap_err();
595        match err {
596            Error::LoraLoading { message, .. } => {
597                assert!(message.contains("after 10 retries"), "{message}");
598            }
599            other => panic!("unexpected: {other:?}"),
600        }
601    }
602
603    #[test]
604    fn model_loading_gating_differs_between_generate_and_encode() {
605        let fail_fast = RequestOptions {
606            wait_for_capacity: false,
607            ..Default::default()
608        };
609
610        let mut encode = state(RetryPolicy::ENCODE, fail_fast.clone());
611        assert_eq!(
612            encode
613                .on_response(&response(503, Some(codes::MODEL_LOADING), None))
614                .unwrap(),
615            Decision::Retry(backoff::MODEL_LOADING_DELAY),
616        );
617
618        let mut generate = state(RetryPolicy::GENERATE, fail_fast);
619        let err = generate
620            .on_response(&response(503, Some(codes::MODEL_LOADING), None))
621            .unwrap_err();
622        assert!(matches!(err, Error::Server { status: 503, .. }));
623    }
624
625    #[test]
626    fn oom_retries_are_capped_and_then_surface_the_root_cause() {
627        let mut state = state(RetryPolicy::ENCODE, RequestOptions::default());
628        for attempt in 0..backoff::RESOURCE_EXHAUSTED_MAX_RETRIES {
629            let decision = state
630                .on_response(&response(503, Some(codes::RESOURCE_EXHAUSTED), None))
631                .unwrap();
632            let Decision::Retry(delay) = decision else {
633                panic!("expected a retry on attempt {attempt}")
634            };
635            assert!(delay <= backoff::RESOURCE_EXHAUSTED_MAX_DELAY);
636        }
637        let err = state
638            .on_response(&response(503, Some(codes::RESOURCE_EXHAUSTED), None))
639            .unwrap_err();
640        match err {
641            Error::ResourceExhausted {
642                retries, message, ..
643            } => {
644                assert_eq!(retries, 3);
645                assert!(message.contains("bge-m3"), "{message}");
646            }
647            other => panic!("unexpected: {other:?}"),
648        }
649    }
650
651    #[test]
652    fn oom_is_disabled_by_zero_max_retries() {
653        let mut state = state(
654            RetryPolicy::ENCODE,
655            RequestOptions {
656                max_oom_retries: 0,
657                ..Default::default()
658            },
659        );
660        let err = state
661            .on_response(&response(503, Some(codes::RESOURCE_EXHAUSTED), None))
662            .unwrap_err();
663        assert!(matches!(err, Error::ResourceExhausted { retries: 0, .. }));
664    }
665
666    #[test]
667    fn oom_is_gated_on_wait_for_capacity_only_on_stream_paths() {
668        let fail_fast = RequestOptions {
669            wait_for_capacity: false,
670            ..Default::default()
671        };
672
673        let mut buffered = state(RetryPolicy::ENCODE, fail_fast.clone());
674        assert!(matches!(
675            buffered.on_response(&response(503, Some(codes::RESOURCE_EXHAUSTED), None)),
676            Ok(Decision::Retry(_))
677        ));
678
679        let mut streaming = state(RetryPolicy::STREAM, fail_fast);
680        assert!(matches!(
681            streaming.on_response(&response(503, Some(codes::RESOURCE_EXHAUSTED), None)),
682            Err(Error::ResourceExhausted { .. })
683        ));
684    }
685
686    #[test]
687    fn gateway_timeout_is_replayed_only_for_idempotent_endpoints() {
688        let mut encode = state(RetryPolicy::ENCODE, RequestOptions::default());
689        assert_eq!(
690            encode.on_response(&response(504, None, Some("3"))).unwrap(),
691            Decision::Retry(Duration::from_secs(3))
692        );
693
694        let mut generate = state(RetryPolicy::GENERATE, RequestOptions::default());
695        let err = generate
696            .on_response(&response(504, None, None))
697            .unwrap_err();
698        match err {
699            Error::Server {
700                message, status, ..
701            } => {
702                assert_eq!(status, 504);
703                assert!(message.contains("double-bill"), "{message}");
704            }
705            other => panic!("unexpected: {other:?}"),
706        }
707    }
708
709    #[test]
710    fn gateway_timeout_is_terminal_without_wait_for_capacity() {
711        let mut encode = state(
712            RetryPolicy::ENCODE,
713            RequestOptions {
714                wait_for_capacity: false,
715                ..Default::default()
716            },
717        );
718        let err = encode.on_response(&response(504, None, None)).unwrap_err();
719        assert!(matches!(err, Error::Server { status: 504, .. }));
720    }
721
722    #[test]
723    fn model_load_failure_short_circuits_before_the_retry_budget() {
724        let mut headers = HeaderMap::new();
725        headers.insert(
726            reqwest::header::CONTENT_TYPE,
727            "application/json".parse().unwrap(),
728        );
729        let failed = HttpResponse {
730            status: 502,
731            headers,
732            body: Bytes::from(
733                r#"{"error": {"code": "MODEL_LOAD_FAILED", "message": "gated", "error_class": "GATED"}}"#,
734            ),
735        };
736        let mut state = state(RetryPolicy::ENCODE, RequestOptions::default());
737        assert!(matches!(
738            state.on_response(&failed),
739            Err(Error::ModelLoadFailed { .. })
740        ));
741        assert_eq!(state.retries(), 0);
742    }
743
744    #[test]
745    fn exhausted_budget_refuses_to_send() {
746        let state = state(
747            RetryPolicy::ENCODE,
748            RequestOptions {
749                provision_timeout: Duration::ZERO,
750                ..Default::default()
751            },
752        );
753        let err = state.attempt_timeout(Duration::from_secs(30)).unwrap_err();
754        match err {
755            Error::Provisioning { message, .. } => {
756                assert!(message.contains("before request could be sent"));
757            }
758            other => panic!("unexpected: {other:?}"),
759        }
760    }
761
762    #[test]
763    fn attempt_timeout_is_clamped_to_the_remaining_budget() {
764        let state = state(
765            RetryPolicy::ENCODE,
766            RequestOptions {
767                provision_timeout: Duration::from_secs(2),
768                ..Default::default()
769            },
770        );
771        assert!(state.attempt_timeout(Duration::from_secs(30)).unwrap() <= Duration::from_secs(2));
772        assert_eq!(
773            state.attempt_timeout(Duration::from_millis(500)).unwrap(),
774            Duration::from_millis(500)
775        );
776    }
777}
sie_sdk/retry/mod.rs

sie_sdk/retry/
mod.rs
