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    ///
62    /// Frames are parsed by BORROWING slices of the buffer behind an advancing
63    /// cursor; consumed bytes are drained once at the end. (The former
64    /// per-frame `drain(..).collect::<String>()` walked the frame char by char
65    /// AND memmoved the buffer tail per frame — the top CPU hotspot of every
66    /// streaming profile.)
67    pub fn push(&mut self, chunk: &[u8]) -> Vec<SseFrame> {
68        self.utf8.decode_into(chunk, &mut self.buf);
69        if self.buf.contains('\r') {
70            self.buf = self.buf.replace("\r\n", "\n");
71        }
72        let mut frames = Vec::new();
73        let mut cursor = 0;
74        while let Some(pos) = self.buf[cursor..].find("\n\n") {
75            let end = cursor + pos + 2;
76            if let Some(frame) = parse_frame(&self.buf[cursor..end]) {
77                frames.push(frame);
78            }
79            cursor = end;
80        }
81        if cursor > 0 {
82            self.buf.drain(..cursor);
83        }
84        frames
85    }
86
87    /// Drain a trailing, unterminated frame at end of stream (some upstreams
88    /// omit the final blank line).
89    pub fn finish(&mut self) -> Option<SseFrame> {
90        self.utf8.flush(&mut self.buf);
91        let raw = std::mem::take(&mut self.buf);
92        parse_frame(&raw)
93    }
94}
95
96fn parse_frame(raw: &str) -> Option<SseFrame> {
97    let mut event = None;
98    let mut data_lines: Vec<&str> = Vec::new();
99    for line in raw.lines() {
100        if let Some(rest) = line.strip_prefix("event:") {
101            event = Some(rest.trim_start().to_owned());
102        } else if let Some(rest) = line.strip_prefix("data:") {
103            data_lines.push(rest.strip_prefix(' ').unwrap_or(rest));
104        }
105    }
106    if data_lines.is_empty() {
107        return None;
108    }
109    Some(SseFrame {
110        event,
111        data: data_lines.join("\n"),
112    })
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[test]
120    fn frame_split_across_chunks() {
121        let mut d = SseDecoder::new();
122        assert!(d.push(b"event: ping\nda").is_empty());
123        let frames = d.push(b"ta: {\"a\":1}\n\n: comment\ndata: x");
124        assert_eq!(frames.len(), 1);
125        assert_eq!(frames[0].event.as_deref(), Some("ping"));
126        assert_eq!(frames[0].data, "{\"a\":1}");
127        // trailing unterminated frame surfaces on finish()
128        assert_eq!(d.finish().unwrap().data, "x");
129    }
130
131    #[test]
132    fn crlf_and_multiline_data() {
133        let mut d = SseDecoder::new();
134        let frames = d.push(b"data: l1\r\ndata: l2\r\n\r\n");
135        assert_eq!(frames.len(), 1);
136        assert_eq!(frames[0].data, "l1\nl2");
137    }
138
139    #[test]
140    fn multibyte_char_split_across_chunks() {
141        // "data: 汉字\n\n" split inside the 3-byte "汉" — must not yield U+FFFD.
142        let mut d = SseDecoder::new();
143        let bytes = "data: 汉字\n\n".as_bytes();
144        assert!(d.push(&bytes[..7]).is_empty()); // cuts "汉" after 1 byte
145        let frames = d.push(&bytes[7..]);
146        assert_eq!(frames.len(), 1);
147        assert_eq!(frames[0].data, "汉字");
148    }
149}
gproxy_transform/transform/common/sse.rs

gproxy_transform/transform/common/
sse.rs
