autumn_plugin_markstream/
sse.rs

1//! Server-Sent Events that stream Markdown into a [`Markstream`](crate::Markstream).
2//!
3//! # Protocol
4//!
5//! | Event | Data | Effect in the browser |
6//! |---|---|---|
7//! | `chunk` | JSON string | Add the text to the end. |
8//! | `replace` | JSON string | Replace all text. |
9//! | `done` | `null` | Mark the text final. Close the stream. |
10//! | `error` | `null` | Mark the text final and failed. Close the stream. |
11//!
12//! The data is JSON, so each string round-trips: leading spaces, `\r`,
13//! `\n` and Unicode stay the same. `done` and `error` carry `null`,
14//! because a browser does not dispatch an event with no data.
15//!
16//! The init module closes the stream and does not reconnect. A reconnect
17//! would start the stream again and show the text two times. A stream that stops before `done`
18//! shows the error state.
19//!
20//! ```rust,no_run
21//! use autumn_plugin_markstream::sse::{markstream_sse, stream};
22//!
23//! # fn handler() -> impl autumn_web::reexports::axum::response::IntoResponse {
24//! markstream_sse(stream::iter(["# Title\n", "Some ", "*text*."]))
25//! # }
26//! ```
27
28use std::convert::Infallible;
29
30use autumn_web::sse::{Event, Sse};
31use futures_util::StreamExt as _;
32/// Re-exported from `futures-util`, so you can build chunk streams.
33pub use futures_util::{Stream, stream};
34
35/// Event name: add text.
36pub const EVENT_CHUNK: &str = "chunk";
37/// Event name: replace all text.
38pub const EVENT_REPLACE: &str = "replace";
39/// Event name: the stream is complete.
40pub const EVENT_DONE: &str = "done";
41/// Event name: the stream failed.
42pub const EVENT_ERROR: &str = "error";
43
44/// One markstream event.
45#[derive(Debug, Clone, PartialEq, Eq)]
46#[must_use]
47#[non_exhaustive]
48pub enum StreamEvent {
49    /// Add this text to the end.
50    Chunk(String),
51    /// Replace all text with this text.
52    Replace(String),
53    /// The stream is complete.
54    Done,
55    /// The stream failed. The event has no details.
56    Error,
57}
58
59impl StreamEvent {
60    /// Makes a [`StreamEvent::Chunk`].
61    pub fn chunk(text: impl Into<String>) -> Self {
62        Self::Chunk(text.into())
63    }
64
65    /// Makes a [`StreamEvent::Replace`].
66    pub fn replace(text: impl Into<String>) -> Self {
67        Self::Replace(text.into())
68    }
69
70    /// The SSE event name.
71    #[must_use]
72    pub const fn name(&self) -> &'static str {
73        match self {
74            Self::Chunk(_) => EVENT_CHUNK,
75            Self::Replace(_) => EVENT_REPLACE,
76            Self::Done => EVENT_DONE,
77            Self::Error => EVENT_ERROR,
78        }
79    }
80
81    /// The SSE data: a JSON string, or `null`.
82    #[must_use]
83    pub fn data(&self) -> String {
84        match self {
85            Self::Chunk(text) | Self::Replace(text) => {
86                serde_json::Value::String(text.clone()).to_string()
87            }
88            Self::Done | Self::Error => "null".to_owned(),
89        }
90    }
91}
92
93impl From<StreamEvent> for Event {
94    fn from(event: StreamEvent) -> Self {
95        Self::default().event(event.name()).data(event.data())
96    }
97}
98
99/// Turns text chunks into markstream events.
100///
101/// It skips empty chunks and sends [`StreamEvent::Done`] at the end.
102pub fn stream_events<S, T>(chunks: S) -> impl Stream<Item = StreamEvent>
103where
104    S: Stream<Item = T>,
105    T: Into<String>,
106{
107    try_stream_events(chunks.map(Ok::<T, Infallible>))
108}
109
110/// Turns fallible text chunks into markstream events.
111///
112/// It skips empty chunks. At the first `Err`, it sends
113/// [`StreamEvent::Error`] and stops. If no error occurs, it sends
114/// [`StreamEvent::Done`] at the end. The last event is always one `Done` or
115/// one `Error`. The error value does not go to the client. Log it before
116/// this function if you need it.
117pub fn try_stream_events<S, T, E>(chunks: S) -> impl Stream<Item = StreamEvent>
118where
119    S: Stream<Item = Result<T, E>>,
120    T: Into<String>,
121{
122    stream::unfold(Some(Box::pin(chunks)), |state| async move {
123        let mut chunks = state?;
124        loop {
125            match chunks.next().await {
126                Some(Ok(text)) => {
127                    let text = text.into();
128                    if !text.is_empty() {
129                        return Some((StreamEvent::Chunk(text), Some(chunks)));
130                    }
131                }
132                Some(Err(_)) => return Some((StreamEvent::Error, None)),
133                None => return Some((StreamEvent::Done, None)),
134            }
135        }
136    })
137}
138
139/// Makes an SSE response from text chunks.
140///
141/// Use it as the handler response for a [`Markstream::stream`] URL. It
142/// sends keep-alive comments every 15 seconds.
143///
144/// [`Markstream::stream`]: crate::Markstream::stream
145pub fn markstream_sse<S, T>(chunks: S) -> Sse<impl Stream<Item = Result<Event, Infallible>>>
146where
147    S: Stream<Item = T> + Send + 'static,
148    T: Into<String> + 'static,
149{
150    try_markstream_sse(chunks.map(Ok::<T, Infallible>))
151}
152
153/// Makes an SSE response from fallible text chunks.
154///
155/// See [`try_stream_events`] for the error rules.
156pub fn try_markstream_sse<S, T, E>(chunks: S) -> Sse<impl Stream<Item = Result<Event, Infallible>>>
157where
158    S: Stream<Item = Result<T, E>> + Send + 'static,
159    T: Into<String> + 'static,
160    E: 'static,
161{
162    let events = try_stream_events(chunks).map(|event| Ok(Event::from(event)));
163    Sse::new(events).keep_alive(autumn_web::sse::keep_alive())
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use crate::test_support::parse_sse;
170    use autumn_web::reexports::axum::{self, response::IntoResponse as _};
171    use autumn_web::sse::Event;
172    use futures_util::stream;
173    use proptest::prelude::*;
174
175    /// Reference model: the events for `items`, in order.
176    fn model(items: &[Result<String, ()>]) -> Vec<StreamEvent> {
177        let mut out = Vec::new();
178        for item in items {
179            match item {
180                Ok(s) if s.is_empty() => {}
181                Ok(s) => out.push(StreamEvent::Chunk(s.clone())),
182                Err(()) => {
183                    out.push(StreamEvent::Error);
184                    return out;
185                }
186            }
187        }
188        out.push(StreamEvent::Done);
189        out
190    }
191
192    fn block_on<F: std::future::Future>(f: F) -> F::Output {
193        tokio::runtime::Builder::new_current_thread()
194            .enable_all()
195            .build()
196            .expect("runtime")
197            .block_on(f)
198    }
199
200    /// The events a browser dispatches for the response that `make`
201    /// builds. `make` runs in the runtime: keep-alive needs a reactor.
202    fn wire<R: axum::response::IntoResponse>(make: impl FnOnce() -> R) -> Vec<(String, String)> {
203        let body = block_on(async move {
204            axum::body::to_bytes(make().into_response().into_body(), usize::MAX).await
205        })
206        .expect("body");
207        parse_sse(std::str::from_utf8(&body).expect("utf-8"))
208    }
209
210    #[test]
211    fn rust_and_js_agree_on_event_names() {
212        let core = include_str!("../assets/core.js");
213        assert_eq!(
214            crate::test_support::js_array(core, "EVENTS"),
215            [EVENT_CHUNK, EVENT_REPLACE, EVENT_DONE, EVENT_ERROR]
216        );
217    }
218
219    #[test]
220    fn event_names_and_payloads() {
221        let cases = [
222            (StreamEvent::chunk("a"), "chunk", r#""a""#),
223            (StreamEvent::replace("b\n"), "replace", r#""b\n""#),
224            (StreamEvent::Done, "done", "null"),
225            (StreamEvent::Error, "error", "null"),
226        ];
227        for (event, name, data) in cases {
228            assert_eq!(event.name(), name);
229            assert_eq!(event.data(), data);
230        }
231        assert_eq!(
232            [EVENT_CHUNK, EVENT_REPLACE, EVENT_DONE, EVENT_ERROR],
233            ["chunk", "replace", "done", "error"]
234        );
235    }
236
237    #[test]
238    fn each_event_dispatches_in_a_browser() {
239        // An empty data buffer does not dispatch. `done` must carry data.
240        let events = [
241            StreamEvent::chunk(""),
242            StreamEvent::Done,
243            StreamEvent::Error,
244        ];
245        let body = stream::iter(events.map(|e| Ok::<_, std::convert::Infallible>(e.into())));
246        let got = wire(|| axum::response::sse::Sse::new(body));
247        let names: Vec<_> = got.iter().map(|(n, _)| n.as_str()).collect();
248        assert_eq!(names, ["chunk", "done", "error"]);
249    }
250
251    #[test]
252    fn errors_stop_the_stream_and_hide_details() {
253        let items = stream::iter([Ok("a"), Err("db password is hunter2"), Ok("b")]);
254        let got = block_on(try_stream_events(items).collect::<Vec<_>>());
255        assert_eq!(got, [StreamEvent::chunk("a"), StreamEvent::Error]);
256        let body = wire(|| {
257            try_markstream_sse(stream::iter([
258                Ok::<_, String>("a"),
259                Err("hunter2".to_owned()),
260            ]))
261        });
262        assert_eq!(body.last().map(|(n, _)| n.as_str()), Some("error"));
263        assert!(body.iter().all(|(_, d)| !d.contains("hunter2")), "{body:?}");
264    }
265
266    #[tokio::test]
267    async fn sse_response_has_event_stream_type() {
268        let response = markstream_sse(stream::iter(["x"])).into_response();
269        assert_eq!(response.headers()["content-type"], "text/event-stream");
270    }
271
272    proptest! {
273        #[test]
274        fn events_match_the_model(items in prop::collection::vec(
275            prop_oneof![4 => any::<String>().prop_map(Ok), 1 => Just(Err(()))], 0..12)
276        ) {
277            let got = block_on(try_stream_events(stream::iter(items.clone())).collect::<Vec<_>>());
278            prop_assert_eq!(&got, &model(&items));
279            let terminal = got.iter().filter(|e| matches!(e, StreamEvent::Done | StreamEvent::Error)).count();
280            prop_assert_eq!(terminal, 1);
281            prop_assert!(matches!(got.last(), Some(StreamEvent::Done | StreamEvent::Error)));
282        }
283
284        #[test]
285        fn chunks_round_trip_through_the_wire(chunks in prop::collection::vec(any::<String>(), 0..8)) {
286            let input = chunks.clone();
287            let got = wire(move || markstream_sse(stream::iter(input)));
288            let mut text = String::new();
289            for (name, data) in &got[..got.len() - 1] {
290                prop_assert_eq!(name, "chunk");
291                let s: String = serde_json::from_str(data).expect("JSON string");
292                text.push_str(&s);
293            }
294            prop_assert_eq!(got.last().map(|(n, _)| n.as_str()), Some("done"));
295            prop_assert_eq!(text, chunks.concat());
296        }
297
298        #[test]
299        fn replace_round_trips(s in any::<String>()) {
300            let event = StreamEvent::replace(s.clone());
301            let got = wire(move || axum::response::sse::Sse::new(stream::iter([
302                Ok::<_, std::convert::Infallible>(Event::from(event)),
303            ])));
304            prop_assert_eq!(got.len(), 1);
305            let back: String = serde_json::from_str(&got[0].1).expect("JSON");
306            prop_assert_eq!(back, s);
307        }
308    }
309}
autumn_plugin_markstream/sse.rs

autumn_plugin_markstream/
sse.rs
