jules_api/streaming/
sse.rs1use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
6pub struct SseEvent {
7 pub event: Option<String>,
9 pub data: String,
11 pub id: Option<String>,
13 pub retry: Option<u64>,
15}
16
17pub struct SseParser {
19 buffer: String,
20}
21
22impl Default for SseParser {
23 fn default() -> Self {
24 Self {
25 buffer: String::with_capacity(8192),
27 }
28 }
29}
30
31impl SseParser {
32 #[must_use]
34 pub fn new() -> Self {
35 Self::default()
36 }
37
38 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 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; }
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 _ => {} }
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}