microclaw_engine/llm/
sse.rs

1#[derive(Default)]
2pub(crate) struct SseEventParser {
3    pending: Vec<u8>,
4    data_lines: Vec<String>,
5}
6
7impl SseEventParser {
8    pub(crate) fn decode_line(line: Vec<u8>) -> String {
9        match String::from_utf8(line) {
10            Ok(line) => line,
11            Err(err) => String::from_utf8_lossy(&err.into_bytes()).into_owned(),
12        }
13    }
14
15    pub(crate) fn push_chunk(&mut self, chunk: impl AsRef<[u8]>) -> Vec<String> {
16        self.pending.extend_from_slice(chunk.as_ref());
17        let mut events = Vec::new();
18
19        while let Some(pos) = self.pending.iter().position(|b| *b == b'\n') {
20            let mut line = self.pending.drain(..=pos).collect::<Vec<_>>();
21            if line.last() == Some(&b'\n') {
22                line.pop();
23            }
24            if line.last() == Some(&b'\r') {
25                line.pop();
26            }
27            let line = Self::decode_line(line);
28            if let Some(event_data) = self.handle_line(&line) {
29                events.push(event_data);
30            }
31        }
32
33        events
34    }
35
36    pub(crate) fn finish(&mut self) -> Vec<String> {
37        let mut events = Vec::new();
38        if !self.pending.is_empty() {
39            let mut line = std::mem::take(&mut self.pending);
40            if line.last() == Some(&b'\r') {
41                line.pop();
42            }
43            let line = Self::decode_line(line);
44            if let Some(event_data) = self.handle_line(&line) {
45                events.push(event_data);
46            }
47        }
48        if let Some(event_data) = self.flush_event() {
49            events.push(event_data);
50        }
51        events
52    }
53
54    pub(crate) fn handle_line(&mut self, line: &str) -> Option<String> {
55        if line.is_empty() {
56            return self.flush_event();
57        }
58        if line.starts_with(':') {
59            return None;
60        }
61
62        let (field, value) = match line.split_once(':') {
63            Some((f, v)) => {
64                let v = v.strip_prefix(' ').unwrap_or(v);
65                (f, v)
66            }
67            None => (line, ""),
68        };
69
70        if field == "data" {
71            self.data_lines.push(value.to_string());
72        }
73        None
74    }
75
76    pub(crate) fn flush_event(&mut self) -> Option<String> {
77        if self.data_lines.is_empty() {
78            return None;
79        }
80        let data = self.data_lines.join("\n");
81        self.data_lines.clear();
82        Some(data)
83    }
84}
85
86// ---------------------------------------------------------------------------
87// Provider trait
88// ---------------------------------------------------------------------------
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    #[allow(unused_imports)]
94    use crate::llm::test_prelude::*;
95
96    #[test]
97    pub(crate) fn test_sse_event_parser_multiline_data() {
98        let mut parser = SseEventParser::default();
99        let events = parser
100            .push_chunk("event: message\n: keep-alive\ndata: {\"type\":\"x\",\ndata: \"v\":1}\n\n");
101        assert_eq!(events.len(), 1);
102        assert_eq!(events[0], "{\"type\":\"x\",\n\"v\":1}");
103    }
104
105    #[test]
106    pub(crate) fn test_sse_event_parser_finish_flushes_unterminated_event() {
107        let mut parser = SseEventParser::default();
108        let events = parser.push_chunk("data: hello");
109        assert!(events.is_empty());
110        let tail = parser.finish();
111        assert_eq!(tail, vec!["hello".to_string()]);
112    }
113
114    #[test]
115    pub(crate) fn test_sse_event_parser_preserves_split_utf8_bytes() {
116        let mut parser = SseEventParser::default();
117        let raw = "data: 主要\n\n".as_bytes();
118        let split = raw.iter().position(|b| *b >= 0x80).unwrap() + 1;
119        let head = parser.push_chunk(&raw[..split]);
120        assert!(head.is_empty());
121        let tail = parser.push_chunk(&raw[split..]);
122        assert_eq!(tail, vec!["主要".to_string()]);
123    }
124}
microclaw_engine/llm/sse.rs

microclaw_engine/llm/
sse.rs
