rama_http/service/web/endpoint/response/
sse.rs

1//! Server-Sent Events (SSE) response.
2
3use rama_core::error::BoxError;
4use rama_core::futures::Stream;
5use rama_core::futures::TryStream;
6use rama_http_headers::{CacheControl, Connection, ContentType};
7use rama_http_types::{
8    Body, Response,
9    sse::{
10        Event, EventDataWrite,
11        server::{KeepAlive, KeepAliveStream, SseResponseBody},
12    },
13};
14
15use super::{Headers, IntoResponse};
16
17/// An SSE response
18#[must_use]
19#[derive(Debug, Clone)]
20pub struct Sse<S> {
21    stream: S,
22}
23
24impl<S> Sse<S> {
25    /// Create a new [`Sse`] response that will respond with the given stream of
26    /// [`Event`]s.
27    pub fn new<T>(stream: S) -> Self
28    where
29        S: TryStream<Ok = Event<T>> + Send + 'static,
30        S::Error: Into<BoxError>,
31        T: EventDataWrite,
32    {
33        Self { stream }
34    }
35
36    /// Configure the interval between keep-alive messages.
37    ///
38    /// Defaults to no keep-alive messages.
39    pub fn with_keep_alive<T, E>(self, keep_alive: KeepAlive<T>) -> Sse<KeepAliveStream<S, T>>
40    where
41        S: Stream<Item = Result<Event<T>, E>>,
42        E: Into<BoxError>,
43        T: EventDataWrite,
44    {
45        Sse {
46            stream: KeepAliveStream::new(keep_alive, self.stream),
47        }
48    }
49}
50
51impl<S, E, T> IntoResponse for Sse<S>
52where
53    S: Stream<Item = Result<Event<T>, E>> + Send + 'static,
54    E: Into<BoxError>,
55    T: EventDataWrite,
56{
57    fn into_response(self) -> Response {
58        (
59            Headers((
60                CacheControl::default().with_no_cache(),
61                ContentType::text_event_stream(),
62                // will be automatically filtered out for h2+
63                Connection::keep_alive(),
64            )),
65            Body::new(SseResponseBody::new(self.stream)),
66        )
67            .into_response()
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use crate::service::{client::HttpClientExt as _, web::Router};
75    use ahash::{HashMap, HashMapExt as _};
76    use rama_core::futures::stream;
77    use rama_core::stream::StreamExt as _;
78    use rama_core::{Service as _, combinators::Either};
79    use rama_http_types::sse::JsonEventData;
80    use rama_utils::str::smol_str::SmolStr;
81    use std::{convert::Infallible, time::Duration};
82
83    #[tokio::test]
84    async fn basic() {
85        let client = Router::new()
86            .with_get("/", async || {
87                let stream = stream::iter(vec![
88                    Event::default()
89                        .with_data(Either::A("one"))
90                        .try_with_comment(SmolStr::new_static("this is a comment"))
91                        .unwrap(),
92                    Event::default().with_data(Either::B(JsonEventData(
93                        serde_json::json!({ "foo": "bar" }),
94                    ))),
95                    Event::default()
96                        .try_with_event(SmolStr::new_static("three"))
97                        .unwrap()
98                        .with_retry(30_000)
99                        .try_with_id(SmolStr::new_static("unique-id"))
100                        .unwrap(),
101                ])
102                .map(Ok::<_, Infallible>);
103                Sse::new(stream)
104            })
105            .boxed();
106
107        let response = client.get("http://example.com").send().await.unwrap();
108
109        assert_eq!(response.headers()["content-type"], "text/event-stream");
110        assert_eq!(response.headers()["cache-control"], "no-cache");
111
112        let mut stream = response.into_body();
113
114        let event_fields =
115            parse_event(std::str::from_utf8(&stream.chunk().await.unwrap().unwrap()).unwrap());
116        assert_eq!(&event_fields["data"], "one");
117        assert_eq!(&event_fields["comment"], "this is a comment");
118
119        let event_fields =
120            parse_event(std::str::from_utf8(&stream.chunk().await.unwrap().unwrap()).unwrap());
121        assert_eq!(&event_fields["data"], "{\"foo\":\"bar\"}");
122        assert!(!event_fields.contains_key("comment"));
123
124        let event_fields =
125            parse_event(std::str::from_utf8(&stream.chunk().await.unwrap().unwrap()).unwrap());
126        assert_eq!(&event_fields["event"], "three");
127        assert_eq!(&event_fields["retry"], "30000");
128        assert_eq!(&event_fields["id"], "unique-id");
129        assert!(!event_fields.contains_key("comment"));
130
131        assert!(stream.chunk().await.unwrap().is_none());
132    }
133
134    #[tokio::test(start_paused = true)]
135    async fn keep_alive() {
136        const DELAY: Duration = Duration::from_secs(5);
137
138        let client = Router::new()
139            .with_get("/", async || {
140                let stream = stream::repeat_with(|| Event::default().with_data("msg"))
141                    .map(Ok::<_, Infallible>)
142                    .throttle(DELAY);
143
144                Sse::new(stream).with_keep_alive(
145                    KeepAlive::<&'static str>::new()
146                        .with_interval(Duration::from_secs(1))
147                        .try_with_text("keep-alive-text")
148                        .unwrap(),
149                )
150            })
151            .boxed();
152
153        let mut stream = client
154            .get("http://example.com")
155            .send()
156            .await
157            .unwrap()
158            .into_body();
159
160        for _ in 0..5 {
161            // first message should be an event
162            let event_fields =
163                parse_event(std::str::from_utf8(&stream.chunk().await.unwrap().unwrap()).unwrap());
164            assert_eq!(&event_fields["data"], "msg");
165
166            // then 4 seconds of keep-alive messages
167            for _ in 0..4 {
168                tokio::time::sleep(Duration::from_secs(1)).await;
169                let event_fields = parse_event(
170                    std::str::from_utf8(&stream.chunk().await.unwrap().unwrap()).unwrap(),
171                );
172                assert_eq!(&event_fields["comment"], "keep-alive-text");
173            }
174        }
175    }
176
177    #[tokio::test(start_paused = true)]
178    async fn keep_alive_ends_when_the_stream_ends() {
179        const DELAY: Duration = Duration::from_secs(5);
180
181        let client = Router::new()
182            .with_get("/", async || {
183                let stream = stream::repeat_with(|| Event::default().with_data("msg"))
184                    .map(Ok::<_, Infallible>)
185                    .throttle(DELAY)
186                    .take(2);
187
188                Sse::new(stream).with_keep_alive(
189                    KeepAlive::<&'static str>::new()
190                        .with_interval(Duration::from_secs(1))
191                        .try_with_text("keep-alive-text")
192                        .unwrap(),
193                )
194            })
195            .boxed();
196
197        let mut stream = client
198            .get("http://example.com")
199            .send()
200            .await
201            .unwrap()
202            .into_body();
203
204        // first message should be an event
205        let event_fields =
206            parse_event(std::str::from_utf8(&stream.chunk().await.unwrap().unwrap()).unwrap());
207        assert_eq!(&event_fields["data"], "msg");
208
209        // then 4 seconds of keep-alive messages
210        for _ in 0..4 {
211            tokio::time::sleep(Duration::from_secs(1)).await;
212            let event_fields =
213                parse_event(std::str::from_utf8(&stream.chunk().await.unwrap().unwrap()).unwrap());
214            assert_eq!(&event_fields["comment"], "keep-alive-text");
215        }
216
217        // then the last event
218        let event_fields =
219            parse_event(std::str::from_utf8(&stream.chunk().await.unwrap().unwrap()).unwrap());
220        assert_eq!(&event_fields["data"], "msg");
221
222        // then no more events or keep-alive messages
223        assert!(stream.chunk().await.unwrap().is_none());
224    }
225
226    fn parse_event(payload: &str) -> HashMap<String, String> {
227        let mut fields = HashMap::new();
228
229        let mut lines = payload.lines().peekable();
230        while let Some(line) = lines.next() {
231            if line.is_empty() {
232                assert_eq!(None, lines.next());
233                break;
234            }
235
236            let (mut key, value) = line.split_once(':').unwrap();
237            let value = value.trim();
238            if key.is_empty() {
239                key = "comment";
240            }
241            fields.insert(key.to_owned(), value.to_owned());
242        }
243
244        fields
245    }
246}
rama_http/service/web/endpoint/response/sse.rs

rama_http/service/web/endpoint/response/
sse.rs
