yapi_ai/
sse.rs

1//! Server-sent events decoding, as the provider SDKs pi uses decode them.
2
3/// One server-sent event.
4#[derive(Clone, Debug, Default, PartialEq, Eq)]
5pub struct Event {
6    /// The `event:` field, if any.
7    pub event: Option<String>,
8    /// `data:` lines joined by `\n`.
9    pub data: String,
10}
11
12/// Incremental decoder: feed bytes as they arrive, collect complete events.
13///
14/// Lines end at `\n`, `\r\n` or `\r`, and a blank line dispatches the event. Comments
15/// (`:`) and unknown fields are ignored. Bytes of a UTF-8 character split across
16/// chunks are kept until the rest arrives.
17#[derive(Debug, Default)]
18pub struct Decoder {
19    pending: Vec<u8>,
20    buffer: String,
21    event: Option<String>,
22    data: Vec<String>,
23}
24
25impl Decoder {
26    /// Decodes a chunk and returns the events it completes.
27    pub fn push(&mut self, chunk: &[u8]) -> Vec<Event> {
28        let text = yapi_types::js::decode_utf8_stream(&mut self.pending, chunk);
29        self.buffer.push_str(&text);
30        self.lines()
31    }
32
33    /// Ends the stream: decodes what is left, including an unterminated last line and
34    /// an undispatched event.
35    pub fn finish(&mut self) -> Vec<Event> {
36        if !self.pending.is_empty() {
37            let text = String::from_utf8_lossy(&self.pending).into_owned();
38            self.pending.clear();
39            self.buffer.push_str(&text);
40        }
41        let mut events = self.lines();
42        if !self.buffer.is_empty() {
43            let line = std::mem::take(&mut self.buffer);
44            events.extend(self.line(&line));
45        }
46        events.extend(self.flush());
47        events
48    }
49
50    fn lines(&mut self) -> Vec<Event> {
51        let mut events = Vec::new();
52        while let Some(index) = self.buffer.find(['\r', '\n']) {
53            let mut next = index + 1;
54            // A `\r` at the end of the buffer may be the first half of `\r\n`.
55            if self.buffer.as_bytes()[index] == b'\r' {
56                match self.buffer.as_bytes().get(next) {
57                    Some(b'\n') => next += 1,
58                    None => break,
59                    Some(_) => {}
60                }
61            }
62            let line = self.buffer[..index].to_owned();
63            self.buffer.drain(..next);
64            events.extend(self.line(&line));
65        }
66        events
67    }
68
69    fn line(&mut self, line: &str) -> Option<Event> {
70        if line.is_empty() {
71            return self.flush();
72        }
73        if line.starts_with(':') {
74            return None;
75        }
76        let (field, value) = match line.split_once(':') {
77            Some((field, value)) => (field, value.strip_prefix(' ').unwrap_or(value)),
78            None => (line, ""),
79        };
80        match field {
81            "event" => self.event = Some(value.to_owned()),
82            "data" => self.data.push(value.to_owned()),
83            _ => {}
84        }
85        None
86    }
87
88    fn flush(&mut self) -> Option<Event> {
89        if self.event.is_none() && self.data.is_empty() {
90            return None;
91        }
92        Some(Event {
93            event: self.event.take(),
94            data: std::mem::take(&mut self.data).join("\n"),
95        })
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    fn decode(chunks: &[&[u8]]) -> Vec<Event> {
104        let mut decoder = Decoder::default();
105        let mut events = Vec::new();
106        for chunk in chunks {
107            events.extend(decoder.push(chunk));
108        }
109        events.extend(decoder.finish());
110        events
111    }
112
113    fn event(name: Option<&str>, data: &str) -> Event {
114        Event {
115            event: name.map(str::to_owned),
116            data: data.to_owned(),
117        }
118    }
119
120    #[test]
121    fn splits_events_across_chunks() {
122        let events = decode(&[
123            b"event: a\nda",
124            b"ta: {\"x\":1}\n\nevent: b\r\ndata: 2\r",
125            b"\n\r\n",
126        ]);
127        assert_eq!(
128            events,
129            [event(Some("a"), "{\"x\":1}"), event(Some("b"), "2")]
130        );
131    }
132
133    #[test]
134    fn joins_data_skips_comments() {
135        let events = decode(&[b": ping\ndata: one\ndata:two\nid: 3\n\n"]);
136        assert_eq!(events, [event(None, "one\ntwo")]);
137    }
138
139    #[test]
140    fn split_utf8_and_flush() {
141        let crab = "data: 🦀".as_bytes();
142        let events = decode(&[&crab[..8], &crab[8..]]);
143        assert_eq!(events, [event(None, "🦀")]);
144    }
145}
yapi_ai/sse.rs

yapi_ai/
sse.rs
