sova_sse/
lib.rs

1//! SSE channel helpers: fan-out, `Last-Event-ID`, keep-alive.
2
3use futures_util::stream::{self, Stream, StreamExt};
4use sova_core::{Request, Response};
5use std::collections::VecDeque;
6use std::convert::Infallible;
7use std::pin::Pin;
8use std::sync::{Arc, Mutex};
9use std::time::Duration;
10use tokio::sync::broadcast;
11use tokio_stream::wrappers::BroadcastStream;
12
13/// In-process pub/sub for SSE clients.
14#[derive(Clone)]
15pub struct SseChannel {
16    tx: broadcast::Sender<SseEvent>,
17    history: Arc<Mutex<VecDeque<SseEvent>>>,
18    history_cap: usize,
19}
20
21#[derive(Clone, Debug)]
22pub struct SseEvent {
23    pub id: Option<String>,
24    pub event: Option<String>,
25    pub data: String,
26}
27
28impl SseEvent {
29    pub fn data(data: impl Into<String>) -> Self {
30        Self {
31            id: None,
32            event: None,
33            data: data.into(),
34        }
35    }
36
37    pub fn id(mut self, id: impl Into<String>) -> Self {
38        self.id = Some(id.into());
39        self
40    }
41
42    pub fn event(mut self, name: impl Into<String>) -> Self {
43        self.event = Some(name.into());
44        self
45    }
46
47    pub fn format(&self) -> String {
48        let mut out = String::new();
49        if let Some(id) = &self.id {
50            out.push_str(&format!("id: {id}\n"));
51        }
52        if let Some(ev) = &self.event {
53            out.push_str(&format!("event: {ev}\n"));
54        }
55        for line in self.data.lines() {
56            out.push_str("data: ");
57            out.push_str(line);
58            out.push('\n');
59        }
60        out.push('\n');
61        out
62    }
63}
64
65impl SseChannel {
66    pub fn new(capacity: usize) -> Self {
67        let (tx, _) = broadcast::channel(capacity.max(16));
68        Self {
69            tx,
70            history: Arc::new(Mutex::new(VecDeque::new())),
71            history_cap: 100,
72        }
73    }
74
75    pub fn history_cap(mut self, n: usize) -> Self {
76        self.history_cap = n.max(1);
77        self
78    }
79
80    pub fn publish(&self, event: SseEvent) {
81        if event.id.is_some() {
82            let mut h = self.history.lock().unwrap();
83            h.push_back(event.clone());
84            while h.len() > self.history_cap {
85                h.pop_front();
86            }
87        }
88        let _ = self.tx.send(event);
89    }
90
91    pub fn subscribe(&self) -> broadcast::Receiver<SseEvent> {
92        self.tx.subscribe()
93    }
94
95    /// Events after `last_id` from history (exclusive), then live stream.
96    pub fn replay_after(&self, last_id: Option<&str>) -> Vec<SseEvent> {
97        let h = self.history.lock().unwrap();
98        let Some(last) = last_id else {
99            return Vec::new();
100        };
101        let mut out = Vec::new();
102        let mut past = false;
103        for e in h.iter() {
104            if past {
105                out.push(e.clone());
106            } else if e.id.as_deref() == Some(last) {
107                past = true;
108            }
109        }
110        out
111    }
112}
113
114/// Build an SSE [`Response`] with optional keep-alive comments.
115pub fn sse_response(
116    req: &Request,
117    channel: &SseChannel,
118    keep_alive: Duration,
119) -> Response {
120    let last_id = req
121        .header("last-event-id")
122        .map(|s| s.to_string());
123    let replay = channel.replay_after(last_id.as_deref());
124    let rx = channel.subscribe();
125
126    let live = BroadcastStream::new(rx).filter_map(|r| async move {
127        r.ok()
128    });
129
130    let replay_stream = stream::iter(replay);
131    let events = replay_stream.chain(live);
132
133    let keepalive = stream::unfold((), move |_| async move {
134        tokio::time::sleep(keep_alive).await;
135        Some((SseEvent::data(""), ()))
136    })
137    .map(|_e| ": keepalive\n\n".to_string());
138
139    let data_stream = events.map(|e| {
140        if e.data.is_empty() && e.id.is_none() && e.event.is_none() {
141            ": keepalive\n\n".to_string()
142        } else {
143            e.format()
144        }
145    });
146
147    let merged = stream::select(
148        data_stream.map(Ok::<_, Infallible>),
149        keepalive.map(Ok::<_, Infallible>),
150    );
151
152    Response::sse(merged)
153}
154
155/// Type alias for boxed SSE byte streams if needed by callers.
156pub type SseStream = Pin<Box<dyn Stream<Item = Result<String, Infallible>> + Send>>;
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn format_and_replay() {
164        let ch = SseChannel::new(8);
165        ch.publish(SseEvent::data("a").id("1"));
166        ch.publish(SseEvent::data("b").id("2"));
167        let replay = ch.replay_after(Some("1"));
168        assert_eq!(replay.len(), 1);
169        assert_eq!(replay[0].data, "b");
170        assert!(SseEvent::data("x").id("9").format().contains("id: 9"));
171    }
172
173    #[test]
174    fn replay_skips_unknown_last_event_id() {
175        let ch = SseChannel::new(8);
176        ch.publish(SseEvent::data("a").id("1"));
177        ch.publish(SseEvent::data("b").id("2"));
178        assert!(ch.replay_after(Some("missing")).is_empty());
179        assert!(ch.replay_after(None).is_empty());
180    }
181
182    #[test]
183    fn replay_respects_history_cap() {
184        let ch = SseChannel::new(8).history_cap(2);
185        ch.publish(SseEvent::data("a").id("1"));
186        ch.publish(SseEvent::data("b").id("2"));
187        ch.publish(SseEvent::data("c").id("3"));
188        assert!(ch.replay_after(Some("1")).is_empty(), "oldest id evicted");
189        let replay = ch.replay_after(Some("2"));
190        assert_eq!(replay.len(), 1);
191        assert_eq!(replay[0].id.as_deref(), Some("3"));
192    }
193
194    #[tokio::test]
195    async fn keepalive_comment_in_stream() {
196        use http::Method;
197        use http_body_util::BodyExt;
198        use sova_core::extend::Body;
199
200        let ch = SseChannel::new(8);
201        let req = Request::new(Method::GET, "/events");
202        let mut res = sse_response(&req, &ch, Duration::from_millis(30));
203        let body = res.take_body();
204        let Body::Stream(mut stream) = body else {
205            panic!("expected sse stream body");
206        };
207        let frame = tokio::time::timeout(Duration::from_millis(200), stream.frame())
208            .await
209            .expect("timeout waiting for keepalive")
210            .expect("stream ended")
211            .expect("frame error");
212        let chunk = frame.into_data().expect("data frame");
213        let text = String::from_utf8_lossy(&chunk);
214        assert!(text.contains("keepalive"), "expected comment keepalive, got {text:?}");
215    }
216}
sova_sse/lib.rs

sova_sse/
lib.rs
