Skip to main content

jules_api/streaming/
sse.rs

1//! Server-Sent Events (SSE) parsing.
2use serde::{Deserialize, Serialize};
3
4/// Represents a single Server-Sent Event (SSE).
5#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
6pub struct SseEvent {
7    /// The event type (e.g., "message").
8    pub event: Option<String>,
9    /// The event payload.
10    pub data: String,
11    /// An optional event ID.
12    pub id: Option<String>,
13    /// The retry time in milliseconds.
14    pub retry: Option<u64>,
15}
16
17/// A parser for buffering and yielding `SseEvent`s from a text stream.
18pub struct SseParser {
19    buffer: String,
20}
21
22impl Default for SseParser {
23    fn default() -> Self {
24        Self {
25            // Bolt optimization: Pre-allocate 8KB to avoid reallocations on initial streaming chunks.
26            buffer: String::with_capacity(8192),
27        }
28    }
29}
30
31impl SseParser {
32    /// Creates a new `SseParser`.
33    #[must_use]
34    pub fn new() -> Self {
35        Self::default()
36    }
37
38    /// Pushes a chunk of text into the parser, returning any complete `SseEvent`s parsed.
39    pub fn push(&mut self, chunk: &str) -> Vec<SseEvent> {
40        self.buffer.push_str(chunk);
41        let mut events = Vec::new();
42
43        let mut last_pos = 0;
44        while let Some(pos) = self.buffer[last_pos..].find("\n\n") {
45            let abs_pos = last_pos + pos;
46            let block = &self.buffer[last_pos..abs_pos];
47
48            if let Some(event) = Self::parse_block(block) {
49                events.push(event);
50            }
51
52            last_pos = abs_pos + 2;
53        }
54
55        if last_pos > 0 {
56            self.buffer.drain(..last_pos);
57        }
58
59        events
60    }
61
62    fn parse_block(block: &str) -> Option<SseEvent> {
63        if block.is_empty() {
64            return None;
65        }
66
67        let mut event = SseEvent {
68            // Bolt optimization: Pre-allocate `data` capacity to the block length (the absolute maximum possible size).
69            // This completely eliminates string reallocations when appending multiple `data:` lines in a hot streaming loop.
70            data: String::with_capacity(block.len()),
71            ..Default::default()
72        };
73        let mut has_data = false;
74
75        for line in block.lines() {
76            if line.starts_with(':') {
77                continue; // Comment
78            }
79
80            if let Some(colon_pos) = line.find(':') {
81                let field = &line[..colon_pos];
82                let mut value = &line[colon_pos + 1..];
83
84                if value.starts_with(' ') {
85                    value = &value[1..];
86                }
87
88                match field {
89                    "event" => event.event = Some(value.to_string()),
90                    "data" => {
91                        if has_data {
92                            event.data.push('\n');
93                        }
94                        event.data.push_str(value);
95                        has_data = true;
96                    }
97                    "id" => event.id = Some(value.to_string()),
98                    "retry" => {
99                        if let Ok(retry) = value.parse() {
100                            event.retry = Some(retry);
101                        }
102                    }
103                    _ => {} // Ignore unknown fields
104                }
105            }
106        }
107
108        if has_data || event.event.is_some() || event.id.is_some() || event.retry.is_some() {
109            Some(event)
110        } else {
111            None
112        }
113    }
114}
115
116#[cfg(test)]
117#[cfg(not(target_arch = "wasm32"))]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_sse_parser_basic() {
123        let mut parser = SseParser::new();
124        let events = parser.push("data: hello\n\n");
125        assert_eq!(events.len(), 1);
126        assert_eq!(events[0].data, "hello");
127    }
128
129    #[test]
130    fn test_sse_parser_fragmented() {
131        let mut parser = SseParser::new();
132        let mut events = parser.push("data: he");
133        assert!(events.is_empty());
134        events = parser.push("llo\n\n");
135        assert_eq!(events.len(), 1);
136        assert_eq!(events[0].data, "hello");
137    }
138
139    #[test]
140    fn test_sse_parser_multiline() {
141        let mut parser = SseParser::new();
142        let events = parser.push("data: line1\ndata: line2\n\n");
143        assert_eq!(events.len(), 1);
144        assert_eq!(events[0].data, "line1\nline2");
145    }
146
147    #[test]
148    fn test_sse_parser_all_fields() {
149        let mut parser = SseParser::new();
150        let events = parser.push("id: 123\nevent: message\ndata: payload\nretry: 5000\n\n");
151        assert_eq!(events.len(), 1);
152        assert_eq!(events[0].id, Some("123".to_string()));
153        assert_eq!(events[0].event, Some("message".to_string()));
154        assert_eq!(events[0].data, "payload");
155        assert_eq!(events[0].retry, Some(5000));
156    }
157}
158
159#[cfg(test)]
160#[cfg(not(target_arch = "wasm32"))]
161mod fuzz {
162    use super::*;
163    use proptest::prelude::*;
164
165    proptest! {
166        #[test]
167        fn test_sse_parser_doesnt_crash(s in ".*") {
168            let mut parser = SseParser::new();
169            let _ = parser.push(&s);
170        }
171
172        #[test]
173        fn test_sse_parser_multiple_pushes(chunks in proptest::collection::vec(".*", 1..10)) {
174            let mut parser = SseParser::new();
175            for chunk in chunks {
176                let _ = parser.push(&chunk);
177            }
178        }
179    }
180}