Skip to main content

gproxy_transform/transform/common/
sse.rs

1/// A provider-neutral SSE transport frame.
2///
3/// This is framing, not a model event IR. Pair modules still own event payload
4/// conversion after JSON decoding.
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct SseFrame {
7    pub event: Option<String>,
8    pub data: String,
9}
10
11impl SseFrame {
12    pub fn data(data: impl Into<String>) -> Self {
13        Self {
14            event: None,
15            data: data.into(),
16        }
17    }
18
19    pub fn event(event: impl Into<String>, data: impl Into<String>) -> Self {
20        Self {
21            event: Some(event.into()),
22            data: data.into(),
23        }
24    }
25
26    pub fn encode(&self) -> String {
27        let mut encoded = String::new();
28        if let Some(event) = &self.event {
29            encoded.push_str("event: ");
30            encoded.push_str(event);
31            encoded.push('\n');
32        }
33        for line in self.data.lines() {
34            encoded.push_str("data: ");
35            encoded.push_str(line);
36            encoded.push('\n');
37        }
38        encoded.push('\n');
39        encoded
40    }
41}
42
43/// Incremental SSE frame decoder: feed raw body chunks, drain complete frames.
44/// Tolerates CRLF, multi-line `data:`, and frames split across chunk
45/// boundaries — including chunks split inside a multi-byte UTF-8 character.
46/// Comments, `id:` and `retry:` lines are framing noise (dropped).
47#[derive(Debug, Default)]
48pub struct SseDecoder {
49    buf: String,
50    utf8: super::utf8::Utf8StreamDecoder,
51}
52
53impl SseDecoder {
54    pub fn new() -> Self {
55        Self::default()
56    }
57
58    /// Append a chunk and return all complete (blank-line-terminated) frames.
59    /// SSE is text by definition; genuinely invalid UTF-8 is replaced lossily,
60    /// while an incomplete trailing sequence waits for the next chunk.
61    pub fn push(&mut self, chunk: &[u8]) -> Vec<SseFrame> {
62        self.utf8.decode_into(chunk, &mut self.buf);
63        if self.buf.contains('\r') {
64            self.buf = self.buf.replace("\r\n", "\n");
65        }
66        let mut frames = Vec::new();
67        while let Some(pos) = self.buf.find("\n\n") {
68            let raw: String = self.buf.drain(..pos + 2).collect();
69            if let Some(frame) = parse_frame(&raw) {
70                frames.push(frame);
71            }
72        }
73        frames
74    }
75
76    /// Drain a trailing, unterminated frame at end of stream (some upstreams
77    /// omit the final blank line).
78    pub fn finish(&mut self) -> Option<SseFrame> {
79        self.utf8.flush(&mut self.buf);
80        let raw = std::mem::take(&mut self.buf);
81        parse_frame(&raw)
82    }
83}
84
85fn parse_frame(raw: &str) -> Option<SseFrame> {
86    let mut event = None;
87    let mut data_lines: Vec<&str> = Vec::new();
88    for line in raw.lines() {
89        if let Some(rest) = line.strip_prefix("event:") {
90            event = Some(rest.trim_start().to_owned());
91        } else if let Some(rest) = line.strip_prefix("data:") {
92            data_lines.push(rest.strip_prefix(' ').unwrap_or(rest));
93        }
94    }
95    if data_lines.is_empty() {
96        return None;
97    }
98    Some(SseFrame {
99        event,
100        data: data_lines.join("\n"),
101    })
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn frame_split_across_chunks() {
110        let mut d = SseDecoder::new();
111        assert!(d.push(b"event: ping\nda").is_empty());
112        let frames = d.push(b"ta: {\"a\":1}\n\n: comment\ndata: x");
113        assert_eq!(frames.len(), 1);
114        assert_eq!(frames[0].event.as_deref(), Some("ping"));
115        assert_eq!(frames[0].data, "{\"a\":1}");
116        // trailing unterminated frame surfaces on finish()
117        assert_eq!(d.finish().unwrap().data, "x");
118    }
119
120    #[test]
121    fn crlf_and_multiline_data() {
122        let mut d = SseDecoder::new();
123        let frames = d.push(b"data: l1\r\ndata: l2\r\n\r\n");
124        assert_eq!(frames.len(), 1);
125        assert_eq!(frames[0].data, "l1\nl2");
126    }
127
128    #[test]
129    fn multibyte_char_split_across_chunks() {
130        // "data: 汉字\n\n" split inside the 3-byte "汉" — must not yield U+FFFD.
131        let mut d = SseDecoder::new();
132        let bytes = "data: 汉字\n\n".as_bytes();
133        assert!(d.push(&bytes[..7]).is_empty()); // cuts "汉" after 1 byte
134        let frames = d.push(&bytes[7..]);
135        assert_eq!(frames.len(), 1);
136        assert_eq!(frames[0].data, "汉字");
137    }
138}