sse.rs

crates/cortex-core/src/provider/sse.rs

1. //! A byte-level Server-Sent Events parser.
2. //!
3. //! Network chunks can split anywhere: in the middle of a line, a `\r\n` pair,
4. //! or a multi-byte UTF-8 character. The parser buffers raw bytes and only
5. //! decodes complete lines, so every way of chunking the same stream yields the
6. //! same events (the property test below checks this).
8. /// One dispatched event: its `event:` name, if any, and its joined `data:` lines.
9. #[derive(Debug, Clone, PartialEq, Eq)]
10. pub struct SseEvent {
11. pub event: Option<String>,
12. pub data: String,
13. }
15. #[derive(Debug, Default)]
16. pub struct SseParser {
17. buf: Vec<u8>,
18. data: Vec<String>,
19. event: Option<String>,
20. /// A `\r` ended the last chunk; a `\n` at the start of the next one belongs to it.
21. pending_cr: bool,
22. }
24. impl SseParser {
25. pub fn new() -> Self {
26. Self::default()
27. }
29. /// Feed bytes; returns every event completed by them.
30. pub fn push(&mut self, chunk: &[u8]) -> Vec<SseEvent> {
31. let mut out = Vec::new();
32. let mut bytes = chunk;
33. if self.pending_cr {
34. self.pending_cr = false;
35. if bytes.first() == Some(&b'\n') {
36. bytes = &bytes[1..];
37. }
38. }
39. self.buf.extend_from_slice(bytes);
40. let mut start = 0;
41. let mut i = 0;
42. while i < self.buf.len() {
43. match self.buf[i] {
44. b'\n' => {
45. let line = self.buf[start..i].to_vec();
46. self.line(&line, &mut out);
47. i += 1;
48. start = i;
49. }
50. b'\r' => {
51. let line = self.buf[start..i].to_vec();
52. self.line(&line, &mut out);
53. i += 1;
54. if i < self.buf.len() {
55. if self.buf[i] == b'\n' {
56. i += 1;
57. }
58. } else {
59. self.pending_cr = true;
60. }
61. start = i;
62. }
63. _ => i += 1,
64. }
65. }
66. self.buf.drain(..start);
67. out
68. }
70. /// End of stream: an event left without its blank line is still delivered.
71. pub fn finish(&mut self) -> Vec<SseEvent> {
72. let mut out = Vec::new();
73. if !self.buf.is_empty() {
74. let line = std::mem::take(&mut self.buf);
75. self.line(&line, &mut out);
76. }
77. self.dispatch(&mut out);
78. out
79. }
81. fn line(&mut self, raw: &[u8], out: &mut Vec<SseEvent>) {
82. if raw.is_empty() {
83. self.dispatch(out);
84. return;
85. }
86. let line = String::from_utf8_lossy(raw);
87. if line.starts_with(':') {
88. return; // comment / keepalive
89. }
90. let (field, value) = match line.find(':') {
91. Some(at) => {
92. let v = &line[at + 1..];
93. (&line[..at], v.strip_prefix(' ').unwrap_or(v))
94. }
95. None => (&line[..], ""),
96. };
97. match field {
98. "data" => self.data.push(value.to_string()),
99. "event" => self.event = Some(value.to_string()),
100. _ => {} // id, retry and unknown fields are not used
101. }
102. }
104. fn dispatch(&mut self, out: &mut Vec<SseEvent>) {
105. if self.data.is_empty() {
106. self.event = None;
107. return;
108. }
109. out.push(SseEvent { event: self.event.take(), data: std::mem::take(&mut self.data).join("\n") });
110. }
111. }
113. #[cfg(test)]
114. mod tests {
115. use super::*;
116. use proptest::prelude::*;
118. fn all(input: &[u8]) -> Vec<SseEvent> {
119. let mut p = SseParser::new();
120. let mut out = p.push(input);
121. out.extend(p.finish());
122. out
123. }
125. fn data(events: &[SseEvent]) -> Vec<&str> {
126. events.iter().map(|e| e.data.as_str()).collect()
127. }
129. #[test]
130. fn basic_lf_and_crlf() {
131. assert_eq!(data(&all(b"data: a\n\ndata: b\n\n")), ["a", "b"]);
132. assert_eq!(data(&all(b"data: a\r\n\r\ndata: b\r\n\r\n")), ["a", "b"]);
133. assert_eq!(data(&all(b"data: a\r\rdata: b\r\r")), ["a", "b"]);
134. }
136. #[test]
137. fn comments_multiline_and_event_names() {
138. let ev = all(b": keepalive\n\nevent: error\ndata: one\ndata: two\n\ndata:x\n\n");
139. assert_eq!(ev.len(), 2);
140. assert_eq!(ev[0].event.as_deref(), Some("error"));
141. assert_eq!(ev[0].data, "one\ntwo");
142. assert_eq!(ev[1], SseEvent { event: None, data: "x".into() });
143. }
145. #[test]
146. fn unterminated_final_event_is_delivered() {
147. assert_eq!(data(&all(b"data: [DONE]")), ["[DONE]"]);
148. }
150. #[test]
151. fn utf8_split_across_chunks() {
152. let text = "data: caf\u{e9} \u{1F600}\n\n".as_bytes();
153. let mut p = SseParser::new();
154. let mut out = Vec::new();
155. for b in text {
156. out.extend(p.push(std::slice::from_ref(b)));
157. }
158. assert_eq!(out[0].data, "caf\u{e9} \u{1F600}");
159. }
161. #[test]
162. fn crlf_split_between_chunks() {
163. let mut p = SseParser::new();
164. let mut out = p.push(b"data: a\r");
165. out.extend(p.push(b"\n\r"));
166. out.extend(p.push(b"\ndata: b\r\n\r\n"));
167. assert_eq!(data(&out), ["a", "b"]);
168. }
170. proptest! {
171. #[test]
172. fn any_chunking_gives_the_same_events(cuts in proptest::collection::vec(0usize..400, 0..30)) {
173. let stream = "data: {\"a\":1}\r\n\r\n: ping\n\nevent: x\ndata: l1\ndata: l2 \u{e9}\u{4e2d}\n\ndata: [DONE]\n\n".as_bytes();
174. let expected = all(stream);
175. let mut points: Vec<usize> = cuts.into_iter().map(|c| c % (stream.len() + 1)).collect();
176. points.sort();
177. let mut p = SseParser::new();
178. let mut got = Vec::new();
179. let mut last = 0;
180. for at in points {
181. got.extend(p.push(&stream[last..at]));
182. last = at;
183. }
184. got.extend(p.push(&stream[last..]));
185. got.extend(p.finish());
186. prop_assert_eq!(got, expected);
187. }
188. }
189. }