use bytes::{Buf, BytesMut};
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum SseToken {
Data(String),
Done,
}
pub(crate) struct SseParser {
buf: BytesMut,
tokens: Vec<SseToken>,
}
impl SseParser {
pub fn new() -> Self {
Self {
buf: BytesMut::new(),
tokens: Vec::new(),
}
}
pub fn push(&mut self, chunk: &[u8]) {
self.buf.extend_from_slice(chunk);
self.drain_lines();
}
pub fn flush(&mut self) {
if !self.buf.is_empty() {
self.buf.extend_from_slice(b"\n");
self.drain_lines();
}
}
pub fn drain(&mut self) -> impl Iterator<Item = SseToken> + '_ {
self.tokens.drain(..)
}
fn drain_lines(&mut self) {
while let Some(lf_pos) = self.buf.iter().position(|&b| b == b'\n') {
let line_end = if lf_pos > 0 && self.buf[lf_pos - 1] == b'\r' {
lf_pos - 1
} else {
lf_pos
};
let Ok(line) = std::str::from_utf8(&self.buf[..line_end]) else {
break;
};
if let Some(token) = parse_line(line) {
self.tokens.push(token);
}
self.buf.advance(lf_pos + 1);
}
}
}
fn parse_line(line: &str) -> Option<SseToken> {
if line.is_empty() || line.starts_with(':') {
return None;
}
let rest = line.strip_prefix("data:")?;
let payload = rest.strip_prefix(' ').unwrap_or(rest);
if payload == "[DONE]" {
Some(SseToken::Done)
} else {
Some(SseToken::Data(payload.to_owned()))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn collect(parser: &mut SseParser) -> Vec<SseToken> {
parser.drain().collect()
}
#[test]
fn basic_data_line() {
let mut p = SseParser::new();
p.push(b"data: {\"a\":1}\n\n");
let tokens = collect(&mut p);
assert_eq!(tokens, vec![SseToken::Data("{\"a\":1}".into())]);
}
#[test]
fn done_sentinel() {
let mut p = SseParser::new();
p.push(b"data: [DONE]\n");
assert_eq!(collect(&mut p), vec![SseToken::Done]);
}
#[test]
fn comment_ignored() {
let mut p = SseParser::new();
p.push(b": keep-alive\ndata: hi\n");
let tokens = collect(&mut p);
assert_eq!(tokens, vec![SseToken::Data("hi".into())]);
}
#[test]
fn data_without_space() {
let mut p = SseParser::new();
p.push(b"data:{\"x\":2}\n");
assert_eq!(collect(&mut p), vec![SseToken::Data("{\"x\":2}".into())]);
}
#[test]
fn split_across_chunks() {
let mut p = SseParser::new();
p.push(b"da");
assert!(collect(&mut p).is_empty());
p.push(b"ta: hello\n");
assert_eq!(collect(&mut p), vec![SseToken::Data("hello".into())]);
}
#[test]
fn split_mid_line() {
let mut p = SseParser::new();
p.push(b"data: {\"tex");
assert!(collect(&mut p).is_empty());
p.push(b"t\":\"hi\"}\n");
assert_eq!(
collect(&mut p),
vec![SseToken::Data("{\"text\":\"hi\"}".into())]
);
}
#[test]
fn multiple_events_one_chunk() {
let mut p = SseParser::new();
p.push(b"data: a\ndata: b\ndata: [DONE]\n");
let tokens = collect(&mut p);
assert_eq!(
tokens,
vec![
SseToken::Data("a".into()),
SseToken::Data("b".into()),
SseToken::Done,
]
);
}
#[test]
fn crlf_line_endings() {
let mut p = SseParser::new();
p.push(b"data: hello\r\n");
assert_eq!(collect(&mut p), vec![SseToken::Data("hello".into())]);
}
#[test]
fn flush_without_trailing_newline() {
let mut p = SseParser::new();
p.push(b"data: last");
assert!(collect(&mut p).is_empty());
p.flush();
assert_eq!(collect(&mut p), vec![SseToken::Data("last".into())]);
}
#[test]
fn split_mid_utf8() {
let full = b"data: caf\xc3\xa9\n";
let mut p = SseParser::new();
p.push(&full[..full.len() - 2]); assert!(collect(&mut p).is_empty());
p.push(&full[full.len() - 2..]); let tokens = collect(&mut p);
assert_eq!(tokens, vec![SseToken::Data("caf\u{00e9}".into())]);
}
}