allama_http/sse/
chat.rs

1//! The chat stream driver over a response's SSE events: the SSE form of
2//! [`FramedChatStream`](crate::framing::FramedChatStream).
3
4use super::SseStream;
5#[cfg(doc)]
6use crate::framing::ChatFrameDecoder;
7use crate::framing::FramedChatStream;
8
9/// A chat stream decoded from a response's SSE events by an adapter's
10/// [`ChatFrameDecoder`]`<`[`SseEvent`](super::SseEvent)`>`: the
11/// [`FramedChatStream`] over an [`SseStream`], with the driver guarantees
12/// the decoder trait lists.
13///
14/// [`SseChatStream::new`](FramedChatStream::new) classifies a failure with
15/// [`SseError::into_error`](super::SseError::into_error), for a body whose
16/// error converts into a [`TransportError`](crate::TransportError), as an
17/// [`HttpResponse`](crate::HttpResponse) body's does;
18/// [`with_error_map`](FramedChatStream::with_error_map) takes any other
19/// body, with a map that keeps the rules of `into_error`: a timeout stays a
20/// timeout, invalid UTF-8 is a stream decode error, an event past the size
21/// cap a stream protocol error, and the messages are the adapter's own.
22///
23/// It is unboxed and `Send` (for a `Send` body); box it into a
24/// [`ChatStream`](allama::ChatStream) with
25/// [`Prepared::finish_stream`](allama::adapter::chat::Prepared::finish_stream):
26///
27/// ```
28/// use allama::{Error, Result, StreamEvent, StreamErrorKind};
29/// use allama_http::TransportError;
30/// use allama_http::framing::ChatFrameDecoder;
31/// use allama_http::sse::{SseChatStream, SseEvent, SseStream};
32/// use futures_util::{FutureExt, StreamExt, stream};
33///
34/// /// `data:` text deltas, then `event: done`.
35/// struct Lines {
36///     started: bool,
37///     done: bool,
38/// }
39///
40/// impl ChatFrameDecoder<SseEvent> for Lines {
41///     fn frame(&mut self, event: SseEvent) -> Vec<Result<StreamEvent>> {
42///         let mut out = Vec::new();
43///         if !self.started {
44///             self.started = true;
45///             out.push(Ok(StreamEvent::response_start(None, None)));
46///             out.push(Ok(StreamEvent::block_start(0, allama::StreamBlockType::Text)));
47///         }
48///         if event.event == "done" {
49///             self.done = true;
50///             out.push(Ok(StreamEvent::block_stop(0)));
51///             out.push(Ok(StreamEvent::ResponseStop));
52///         } else {
53///             out.push(Ok(StreamEvent::text_delta(0, event.data)));
54///         }
55///         out
56///     }
57///
58///     fn end(&mut self) -> Vec<Result<StreamEvent>> {
59///         let error = Error::stream(StreamErrorKind::Incomplete, "the stream ended early");
60///         vec![Err(error.authored().with_provider("lines"))]
61///     }
62///
63///     fn fail(&mut self, error: Error) -> Vec<Result<StreamEvent>> {
64///         vec![Err(error.with_provider("lines"))]
65///     }
66///
67///     fn is_done(&self) -> bool {
68///         self.done
69///     }
70/// }
71///
72/// let body = stream::iter([Ok::<_, TransportError>(
73///     &b"data: Hel\n\ndata: lo\n\nevent: done\ndata: {}\n\n"[..],
74/// )]);
75/// let stream = SseChatStream::new(SseStream::new(body), Lines { started: false, done: false });
76/// let events: Vec<_> = stream.collect().now_or_never().unwrap();
77/// assert_eq!(events.len(), 6);
78/// assert!(matches!(events[5], Ok(StreamEvent::ResponseStop)));
79/// ```
80pub type SseChatStream<S, D> = FramedChatStream<SseStream<S>, D>;
81
82#[cfg(all(test, not(target_arch = "wasm32")))]
83mod tests {
84    use std::sync::{Arc, Mutex};
85
86    use allama::{ChatStream, Error, StreamErrorKind};
87    use futures_util::StreamExt;
88
89    use super::*;
90    use crate::framing::ChatFrameDecoder;
91    use crate::sse::{SseError, SseEvent};
92    use crate::{TransportError, TransportErrorKind};
93    use allama::Result;
94    use allama::StreamEvent;
95    use futures_core::Stream;
96
97    /// What the decoder was called with, shared with the test.
98    #[derive(Debug, Default)]
99    struct Calls {
100        events: Vec<String>,
101        ends: usize,
102        fails: Vec<Error>,
103    }
104
105    /// A decoder that yields the scripted items for each event (by data),
106    /// and records its calls.
107    struct Scripted {
108        calls: Arc<Mutex<Calls>>,
109        script: fn(&str) -> Vec<Result<StreamEvent>>,
110        end: Vec<Result<StreamEvent>>,
111        done_after: Option<&'static str>,
112        done: bool,
113    }
114
115    impl Scripted {
116        fn new(script: fn(&str) -> Vec<Result<StreamEvent>>) -> (Self, Arc<Mutex<Calls>>) {
117            let calls = Arc::default();
118            let decoder = Self {
119                calls: Arc::clone(&calls),
120                script,
121                end: Vec::new(),
122                done_after: None,
123                done: false,
124            };
125            (decoder, calls)
126        }
127    }
128
129    impl ChatFrameDecoder<SseEvent> for Scripted {
130        fn frame(&mut self, event: SseEvent) -> Vec<Result<StreamEvent>> {
131            self.calls.lock().unwrap().events.push(event.data.clone());
132            self.done |= self.done_after == Some(event.data.as_str());
133            (self.script)(&event.data)
134        }
135
136        fn end(&mut self) -> Vec<Result<StreamEvent>> {
137            self.calls.lock().unwrap().ends += 1;
138            std::mem::take(&mut self.end)
139        }
140
141        fn fail(&mut self, error: Error) -> Vec<Result<StreamEvent>> {
142            self.calls.lock().unwrap().fails.push(error.clone());
143            vec![Err(error)]
144        }
145
146        fn is_done(&self) -> bool {
147            self.done
148        }
149    }
150
151    fn delta(data: &str) -> Vec<Result<StreamEvent>> {
152        vec![Ok(StreamEvent::text_delta(0, data))]
153    }
154
155    /// A body of `chunks`, which counts its polls.
156    fn body(
157        chunks: Vec<std::result::Result<&'static str, TransportError>>,
158        polls: &Arc<Mutex<usize>>,
159    ) -> impl Stream<Item = std::result::Result<&'static str, TransportError>> + Send + Unpin + 'static
160    {
161        let polls = Arc::clone(polls);
162        futures_util::stream::iter(chunks).inspect(move |_| *polls.lock().unwrap() += 1)
163    }
164
165    async fn run<D: ChatFrameDecoder<SseEvent>>(
166        chunks: Vec<std::result::Result<&'static str, TransportError>>,
167        decoder: D,
168    ) -> (Vec<Result<StreamEvent>>, usize) {
169        let polls = Arc::default();
170        let stream = SseChatStream::new(SseStream::new(body(chunks, &polls)), decoder);
171        let items: Vec<_> = stream.collect().await;
172        let polls = *polls.lock().unwrap();
173        (items, polls)
174    }
175
176    #[tokio::test]
177    async fn end_runs_once_when_the_body_ends() {
178        let (mut decoder, calls) = Scripted::new(delta);
179        decoder.end = vec![Ok(StreamEvent::ResponseStop)];
180        let (items, _) = run(vec![Ok("data: a\n\ndata: b\n\n")], decoder).await;
181        assert_eq!(items.len(), 3);
182        assert!(matches!(items[2], Ok(StreamEvent::ResponseStop)));
183        let calls = calls.lock().unwrap();
184        assert_eq!(calls.events, ["a", "b"]);
185        assert_eq!(calls.ends, 1);
186        assert!(calls.fails.is_empty());
187    }
188
189    #[tokio::test]
190    async fn an_empty_end_terminates() {
191        let (decoder, calls) = Scripted::new(delta);
192        let (items, _) = run(vec![Ok("data: a\n\n")], decoder).await;
193        assert_eq!(items.len(), 1);
194        assert_eq!(calls.lock().unwrap().ends, 1);
195    }
196
197    #[tokio::test]
198    async fn fail_runs_once_when_the_body_fails_and_nothing_follows() {
199        let (decoder, calls) = Scripted::new(delta);
200        let (items, polls) = run(
201            vec![
202                Ok("data: a\n\n"),
203                Err(TransportError::new(TransportErrorKind::Body, "reset")),
204                Ok("data: b\n\n"),
205            ],
206            decoder,
207        )
208        .await;
209        assert_eq!(items.len(), 2);
210        let error = items[1].as_ref().unwrap_err();
211        assert!(matches!(
212            error,
213            Error::Stream {
214                kind: StreamErrorKind::Transport,
215                ..
216            }
217        ));
218        assert_eq!(error.upstream_message(), Some("reset"));
219        assert_eq!(polls, 2, "the body is not polled after it failed");
220        let calls = calls.lock().unwrap();
221        assert_eq!(calls.fails.len(), 1);
222        assert_eq!(calls.ends, 0);
223    }
224
225    #[tokio::test]
226    async fn a_timeout_stays_a_timeout() {
227        let (decoder, _) = Scripted::new(delta);
228        let (items, _) = run(
229            vec![Err(TransportError::new(
230                TransportErrorKind::Timeout,
231                "idle",
232            ))],
233            decoder,
234        )
235        .await;
236        assert!(matches!(items[..], [Err(Error::Timeout { .. })]));
237    }
238
239    #[tokio::test]
240    async fn the_stream_ends_after_a_decoder_error() {
241        let (decoder, calls) = Scripted::new(|data| match data {
242            "bad" => vec![
243                Ok(StreamEvent::text_delta(0, "before")),
244                Err(Error::stream_decode("bad")),
245                Ok(StreamEvent::text_delta(0, "dropped")),
246            ],
247            data => delta(data),
248        });
249        let (items, polls) = run(
250            vec![
251                Ok("data: a\n\n"),
252                Ok("data: bad\n\n"),
253                Ok("data: c\n\n"),
254                Ok("data: d\n\n"),
255            ],
256            decoder,
257        )
258        .await;
259        assert_eq!(items.len(), 3, "{items:?}");
260        assert!(items[2].is_err());
261        assert_eq!(polls, 2, "the body is not polled after the error");
262        let calls = calls.lock().unwrap();
263        assert_eq!(calls.events, ["a", "bad"]);
264        assert_eq!(calls.ends, 0);
265        assert!(calls.fails.is_empty());
266    }
267
268    #[tokio::test]
269    async fn the_stream_ends_once_the_decoder_is_done() {
270        let (mut decoder, calls) = Scripted::new(delta);
271        decoder.done_after = Some("last");
272        let (items, polls) = run(
273            vec![
274                Ok("data: a\n\ndata: last\n\n"),
275                Ok("data: ignored\n\n"),
276                Err(TransportError::new(TransportErrorKind::Body, "never read")),
277            ],
278            decoder,
279        )
280        .await;
281        assert_eq!(items.len(), 2);
282        assert_eq!(polls, 1);
283        let calls = calls.lock().unwrap();
284        assert_eq!(calls.events, ["a", "last"]);
285        assert_eq!(calls.ends, 0);
286        assert!(calls.fails.is_empty());
287    }
288
289    #[tokio::test]
290    async fn sse_failures_are_classified() {
291        let polls = Arc::default();
292        let invalid = futures_util::stream::iter([Ok::<_, TransportError>(&b"data: \xFF\n\n"[..])]);
293        let (decoder, _) = Scripted::new(delta);
294        let items: Vec<_> = SseChatStream::new(SseStream::new(invalid), decoder)
295            .collect()
296            .await;
297        let error = items[0].as_ref().unwrap_err();
298        assert!(matches!(
299            error,
300            Error::Stream {
301                kind: StreamErrorKind::Decode,
302                ..
303            }
304        ));
305        assert!(!error.is_retryable());
306        assert_eq!(error.message_origin(), allama::MessageOrigin::Authored);
307
308        let (decoder, _) = Scripted::new(delta);
309        let events =
310            SseStream::new(body(vec![Ok("data: 0123456789abcdef!\n\n")], &polls)).max_event_size(8);
311        let items: Vec<_> = SseChatStream::new(events, decoder).collect().await;
312        assert!(matches!(
313            items[..],
314            [Err(Error::Stream {
315                kind: StreamErrorKind::Protocol,
316                ..
317            })]
318        ));
319    }
320
321    #[tokio::test]
322    async fn an_error_map_classifies_other_bodies() {
323        let body = futures_util::stream::iter([
324            Ok::<_, String>("data: a\n\n"),
325            Err("platform failure".to_owned()),
326        ]);
327        let (decoder, _) = Scripted::new(delta);
328        let stream =
329            SseChatStream::with_error_map(SseStream::new(body), decoder, |error| match error {
330                SseError::Transport(text) => Error::stream(StreamErrorKind::Transport, "failed")
331                    .authored()
332                    .with_upstream_message(text),
333                other => Error::stream_decode(other.to_string()).authored(),
334            });
335        let items: Vec<_> = stream.collect().await;
336        assert_eq!(
337            items[1].as_ref().unwrap_err().upstream_message(),
338            Some("platform failure")
339        );
340    }
341
342    #[test]
343    fn the_stream_is_send() {
344        fn send<T: Send + 'static>() {}
345        send::<SseChatStream<crate::Body, Scripted>>();
346        let (decoder, _) = Scripted::new(delta);
347        let events = SseStream::new(futures_util::stream::empty::<
348            std::result::Result<bytes::Bytes, TransportError>,
349        >());
350        let _: ChatStream = Box::pin(SseChatStream::new(events, decoder));
351    }
352}
allama_http/sse/chat.rs

allama_http/sse/
chat.rs
