allama_http/
timeout.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4use std::time::Duration;
5
6use bytes::Bytes;
7use futures_core::Stream;
8use futures_timer::Delay;
9use futures_util::future::{Either, select};
10
11use crate::{Body, HttpRequest, HttpResponse, HttpTransport, TransportError, TransportErrorKind};
12
13/// Time bounds for a streaming request, enforced above the transport.
14///
15/// A streaming response has no total deadline (a long answer is not a
16/// failure), but a dead endpoint or a stalled connection must not hang a
17/// caller forever. Two bounds catch that:
18///
19/// - `setup`: from sending the request until the response status and
20///   headers arrive.
21/// - `idle`: the longest wait for a body chunk, the first one included.
22///   Only waiting on the transport counts: for each chunk, the clock starts
23///   when the reader polls and the chunk is not ready yet, so the time
24///   before the reader's first poll or between its polls does not count.
25///
26/// [`send`](Self::send) enforces both on any [`HttpTransport`]; provider
27/// crates call it for their streaming requests (non-streaming requests carry
28/// a whole-exchange [`HttpRequest::timeout`] instead). Expiry fails with a
29/// [`TransportErrorKind::Timeout`] error, which providers surface as
30/// [`allama::Error::Timeout`], retryable and fallback-eligible.
31///
32/// The defaults ([`DEFAULT_SETUP`](Self::DEFAULT_SETUP),
33/// [`DEFAULT_IDLE`](Self::DEFAULT_IDLE)) are deliberately generous: they
34/// detect dead connections, not slow models. Reasoning models can stay
35/// silent for minutes before their first token, and an idle bound that cuts
36/// such a stream off bills the reasoning and returns nothing. Tighter,
37/// workload-specific bounds belong in `allama::wrappers::TimeoutProvider`,
38/// which wraps the decoded stream and fires first when its limits are
39/// shorter. It can also bound a stream's total duration
40/// (`TimeoutPolicy::stream_total`), which this transport layer leaves to
41/// the application.
42///
43/// Timers come from `futures-timer` (the same timers `TimeoutProvider`
44/// uses), so the bounds work on any async runtime and, through the browser
45/// or worker `setTimeout`, on `wasm32-unknown-unknown`.
46///
47/// ```
48/// use std::time::Duration;
49/// use allama_http::StreamTimeouts;
50///
51/// let tight = StreamTimeouts::default().idle(Some(Duration::from_secs(30)));
52/// assert_eq!(tight.setup, Some(StreamTimeouts::DEFAULT_SETUP));
53///
54/// // A local server that may load a model first: no bounds at all.
55/// let local = StreamTimeouts::unbounded();
56/// assert_eq!((local.setup, local.idle), (None, None));
57/// ```
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59#[non_exhaustive]
60pub struct StreamTimeouts {
61    /// Bound on the wait for the response status and headers.
62    pub setup: Option<Duration>,
63    /// Bound on the wait for each body chunk, the first one included,
64    /// counted from the reader's poll that finds the chunk not ready.
65    pub idle: Option<Duration>,
66}
67
68impl Default for StreamTimeouts {
69    /// [`DEFAULT_SETUP`](Self::DEFAULT_SETUP) and
70    /// [`DEFAULT_IDLE`](Self::DEFAULT_IDLE).
71    fn default() -> Self {
72        Self {
73            setup: Some(Self::DEFAULT_SETUP),
74            idle: Some(Self::DEFAULT_IDLE),
75        }
76    }
77}
78
79impl StreamTimeouts {
80    /// Default [`setup`](Self::setup) bound: 120 s, the same budget the
81    /// provider crates give a whole non-streaming call. A streaming endpoint
82    /// normally answers with headers right away and keeps the model's
83    /// latency for the body.
84    pub const DEFAULT_SETUP: Duration = Duration::from_secs(120);
85
86    /// Default [`idle`](Self::idle) bound: 300 s. Longer than a non-streaming
87    /// call's whole default budget on purpose: streaming is how callers reach
88    /// reasoning models that think for minutes before their first token, and
89    /// not every endpoint sends keep-alive events meanwhile.
90    pub const DEFAULT_IDLE: Duration = Duration::from_secs(300);
91
92    /// No bounds: the stream may wait forever (for example for a local
93    /// server that loads the model before answering).
94    #[must_use]
95    pub const fn unbounded() -> Self {
96        Self {
97            setup: None,
98            idle: None,
99        }
100    }
101
102    /// Set (or clear) the setup bound.
103    #[must_use]
104    pub const fn setup(mut self, timeout: Option<Duration>) -> Self {
105        self.setup = timeout;
106        self
107    }
108
109    /// Set (or clear) the idle bound.
110    #[must_use]
111    pub const fn idle(mut self, timeout: Option<Duration>) -> Self {
112        self.idle = timeout;
113        self
114    }
115
116    /// Send `request` through `transport` with these bounds: the response
117    /// headers must arrive within `setup`, and its body then fails when no
118    /// chunk arrives within `idle`.
119    ///
120    /// Expiry drops the transport's future or body, cancelling the request.
121    ///
122    /// # Errors
123    ///
124    /// The transport's errors, or a [`TransportErrorKind::Timeout`] error
125    /// when the setup bound expires. An idle expiry is the body's last item.
126    pub async fn send<T>(
127        self,
128        transport: &T,
129        request: HttpRequest,
130    ) -> Result<HttpResponse, TransportError>
131    where
132        T: HttpTransport + ?Sized,
133    {
134        let sent = transport.send(request);
135        let response = match self.setup {
136            Some(limit) => match select(sent, Delay::new(limit)).await {
137                Either::Left((result, _)) => result?,
138                Either::Right(((), _)) => return Err(expired("stream setup", limit)),
139            },
140            None => sent.await?,
141        };
142        let Some(idle) = self.idle else {
143            return Ok(response);
144        };
145        let status = response.status;
146        let headers = response.headers.clone();
147        let body = IdleTimeoutBody::new(response.into_body(), idle);
148        Ok(HttpResponse::new(status, headers, Box::pin(body)))
149    }
150}
151
152/// A body that fails when no chunk arrives within the idle bound, then ends.
153///
154/// For each chunk, the first one included, the clock starts when the reader
155/// polls and the chunk is not ready yet: time the reader spends before its
156/// first poll or between polls does not count. Bodies such as hyper's
157/// HTTP/1 bodies read from the socket only while polled, so a timer left
158/// running across a slow reader would expire on the first poll after it.
159struct IdleTimeoutBody {
160    inner: Option<Body>,
161    idle: Duration,
162    timer: Delay,
163    /// The timer is running: set on the first `Pending` of each wait (the
164    /// first one included), cleared when a chunk is handed out.
165    armed: bool,
166}
167
168impl IdleTimeoutBody {
169    fn new(inner: Body, idle: Duration) -> Self {
170        Self {
171            inner: Some(inner),
172            idle,
173            timer: Delay::new(idle),
174            // Disarmed until the reader first waits: an adapter may do
175            // other work between the headers and its first read.
176            armed: false,
177        }
178    }
179}
180
181impl Stream for IdleTimeoutBody {
182    type Item = Result<Bytes, TransportError>;
183
184    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
185        let this = &mut *self;
186        let Some(inner) = this.inner.as_mut() else {
187            return Poll::Ready(None);
188        };
189        match inner.as_mut().poll_next(cx) {
190            Poll::Ready(Some(item)) => {
191                this.armed = false;
192                Poll::Ready(Some(item))
193            }
194            Poll::Ready(None) => {
195                this.inner = None;
196                Poll::Ready(None)
197            }
198            Poll::Pending => {
199                if !this.armed {
200                    this.timer.reset(this.idle);
201                    this.armed = true;
202                }
203                if Pin::new(&mut this.timer).poll(cx).is_pending() {
204                    return Poll::Pending;
205                }
206                this.inner = None;
207                Poll::Ready(Some(Err(expired(
208                    "wait for the next body chunk",
209                    this.idle,
210                ))))
211            }
212        }
213    }
214}
215
216/// The error for an expired bound, phrased like `TimeoutProvider`'s:
217/// `"<what> exceeded <limit>"`.
218fn expired(what: &str, limit: Duration) -> TransportError {
219    TransportError::new(
220        TransportErrorKind::Timeout,
221        format!("{what} exceeded {limit:?}"),
222    )
223    .authored()
224}
225
226#[cfg(all(test, not(target_arch = "wasm32")))]
227mod tests {
228    use futures_util::StreamExt;
229
230    use super::*;
231    use crate::{Headers, ReplayResponse, ReplayTransport, SendFuture};
232
233    const SHORT: Duration = Duration::from_millis(20);
234
235    /// Headers after `headers_after`; then each chunk after `gap`, or a body
236    /// that never produces anything when `gap` is `None`.
237    struct Slow {
238        headers_after: Duration,
239        gap: Option<Duration>,
240    }
241
242    impl HttpTransport for Slow {
243        fn send(&self, _request: HttpRequest) -> SendFuture<'_> {
244            let gap = self.gap;
245            let headers_after = self.headers_after;
246            Box::pin(async move {
247                Delay::new(headers_after).await;
248                let body: Body = match gap {
249                    Some(gap) => Box::pin(futures_util::stream::iter(["a", "b", "c"]).then(
250                        move |chunk| async move {
251                            Delay::new(gap).await;
252                            Ok(Bytes::from_static(chunk.as_bytes()))
253                        },
254                    )),
255                    None => Box::pin(futures_util::stream::pending()),
256                };
257                Ok(HttpResponse::new(200, Headers::new(), body))
258            })
259        }
260    }
261
262    fn request() -> HttpRequest {
263        HttpRequest::post("https://x", "{}")
264    }
265
266    #[tokio::test]
267    async fn slow_headers_fail_the_setup_bound() {
268        let transport = Slow {
269            headers_after: Duration::from_secs(5),
270            gap: Some(Duration::ZERO),
271        };
272        let err = StreamTimeouts::unbounded()
273            .setup(Some(SHORT))
274            .send(&transport, request())
275            .await
276            .unwrap_err();
277        assert_eq!(err.kind(), TransportErrorKind::Timeout);
278        let err = err.into_error();
279        assert!(matches!(err, allama::Error::Timeout { .. }));
280        assert_eq!(err.message(), "stream setup exceeded 20ms");
281        // Written by the bound itself: the redacted log keeps it.
282        assert_eq!(
283            err.as_log().to_string(),
284            r#"kind=timeout message="stream setup exceeded 20ms""#
285        );
286    }
287
288    #[tokio::test]
289    async fn a_stalled_body_fails_the_idle_bound_then_ends() {
290        let transport = Slow {
291            headers_after: Duration::ZERO,
292            gap: None,
293        };
294        let response = StreamTimeouts::default()
295            .idle(Some(SHORT))
296            .send(&transport, request())
297            .await
298            .unwrap();
299        assert_eq!(response.status, 200);
300        let mut body = response.into_body();
301        let err = body.next().await.unwrap().unwrap_err();
302        assert_eq!(err.kind(), TransportErrorKind::Timeout);
303        assert_eq!(
304            err.to_string(),
305            "request timed out: wait for the next body chunk exceeded 20ms"
306        );
307        assert!(body.next().await.is_none());
308    }
309
310    #[tokio::test]
311    async fn steady_chunks_reset_the_idle_bound() {
312        let transport = Slow {
313            headers_after: Duration::ZERO,
314            gap: Some(Duration::from_millis(30)),
315        };
316        // Three 30 ms gaps: 90 ms in total, but never 60 ms of silence.
317        let response = StreamTimeouts::unbounded()
318            .idle(Some(Duration::from_millis(60)))
319            .send(&transport, request())
320            .await
321            .unwrap();
322        assert_eq!(response.text().await.unwrap(), "abc");
323    }
324
325    /// The idle bound measures the transport, not the consumer: time spent
326    /// between polls does not count, even though the body is read only
327    /// while polled.
328    #[tokio::test]
329    async fn slow_consumers_do_not_trip_the_idle_bound() {
330        let transport = Slow {
331            headers_after: Duration::ZERO,
332            gap: Some(SHORT),
333        };
334        let response = StreamTimeouts::unbounded()
335            .idle(Some(Duration::from_millis(100)))
336            .send(&transport, request())
337            .await
338            .unwrap();
339        let mut body = response.into_body();
340        let mut chunks = Vec::new();
341        while let Some(chunk) = body.next().await {
342            chunks.push(chunk.unwrap());
343            Delay::new(Duration::from_millis(150)).await;
344        }
345        assert_eq!(chunks.concat(), b"abc");
346    }
347
348    /// The idle bound starts at the first wait, not at the headers: an
349    /// adapter that reads the body late is not timed out for it.
350    #[tokio::test]
351    async fn a_late_first_poll_does_not_trip_the_idle_bound() {
352        let transport = Slow {
353            headers_after: Duration::ZERO,
354            gap: Some(Duration::from_millis(10)),
355        };
356        let response = StreamTimeouts::unbounded()
357            .idle(Some(Duration::from_millis(50)))
358            .send(&transport, request())
359            .await
360            .unwrap();
361        // The first chunk is not ready when first polled: a timer started at
362        // the headers would have expired by then.
363        Delay::new(Duration::from_millis(100)).await;
364        assert_eq!(response.text().await.unwrap(), "abc");
365    }
366
367    /// A first chunk that never comes still trips the bound, measured from
368    /// the first poll.
369    #[tokio::test]
370    async fn a_stalled_first_chunk_fails_the_idle_bound() {
371        let body: Body = Box::pin(futures_util::stream::pending());
372        let mut body = IdleTimeoutBody::new(body, Duration::from_millis(50));
373        Delay::new(Duration::from_millis(100)).await;
374        let started = std::time::Instant::now();
375        let err = body.next().await.unwrap().unwrap_err();
376        assert_eq!(err.kind(), TransportErrorKind::Timeout);
377        assert!(
378            started.elapsed() >= Duration::from_millis(50),
379            "{:?}",
380            started.elapsed()
381        );
382        assert!(body.next().await.is_none());
383    }
384
385    /// A body that stalls after a chunk still trips the bound, measured from
386    /// the poll that found nothing ready.
387    #[tokio::test]
388    async fn a_body_stalled_after_a_chunk_fails_the_idle_bound() {
389        let body: Body = Box::pin(
390            futures_util::stream::iter([Ok(Bytes::from_static(b"a"))])
391                .chain(futures_util::stream::pending()),
392        );
393        let mut body = IdleTimeoutBody::new(body, SHORT);
394        assert_eq!(body.next().await.unwrap().unwrap(), "a");
395        Delay::new(SHORT * 3).await;
396        let started = std::time::Instant::now();
397        let err = body.next().await.unwrap().unwrap_err();
398        assert_eq!(err.kind(), TransportErrorKind::Timeout);
399        assert!(started.elapsed() >= SHORT, "{:?}", started.elapsed());
400        assert!(body.next().await.is_none());
401    }
402
403    #[tokio::test]
404    async fn unbounded_passes_responses_and_errors_through() {
405        let transport = ReplayTransport::new([ReplayResponse::new(201, "ok")]);
406        let response = StreamTimeouts::unbounded()
407            .send(&transport, request())
408            .await
409            .unwrap();
410        assert_eq!(response.status, 201);
411        assert_eq!(response.text().await.unwrap(), "ok");
412
413        // The queue is empty now: the transport's own error comes back.
414        let err = StreamTimeouts::default()
415            .send(&transport, request())
416            .await
417            .unwrap_err();
418        assert_eq!(err.kind(), TransportErrorKind::Connect);
419    }
420}
421
422/// Run on `wasm32` with `wasm-bindgen-test`: the timers go through the
423/// host's `setTimeout`.
424#[cfg(all(test, target_arch = "wasm32"))]
425mod wasm_tests {
426    use futures_util::StreamExt;
427    use wasm_bindgen_test::wasm_bindgen_test;
428
429    use super::*;
430    use crate::{Headers, SendFuture};
431
432    struct Stalled;
433
434    impl HttpTransport for Stalled {
435        fn send(&self, _request: HttpRequest) -> SendFuture<'_> {
436            Box::pin(async {
437                Ok(HttpResponse::new(
438                    200,
439                    Headers::new(),
440                    Box::pin(futures_util::stream::pending()),
441                ))
442            })
443        }
444    }
445
446    #[wasm_bindgen_test]
447    async fn a_stalled_body_fails_the_idle_bound() {
448        let response = StreamTimeouts::default()
449            .idle(Some(Duration::from_millis(20)))
450            .send(&Stalled, HttpRequest::get("https://x"))
451            .await
452            .unwrap();
453        let err = response.into_body().next().await.unwrap().unwrap_err();
454        assert_eq!(err.kind(), TransportErrorKind::Timeout);
455    }
456}
allama_http/timeout.rs

allama_http/
timeout.rs
