audacity-sdk 0.2.0

//! Incremental SSE parser — handles chunk-boundary splits, multi-byte UTF-8,
//! comments, missing final newline, and `data: [DONE]`.

use bytes::{Buf, BytesMut};

/// Token emitted by the SSE parser.
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum SseToken {
    /// A `data:` payload (already stripped of the leading optional space).
    Data(String),
    /// The stream-terminating `data: [DONE]` sentinel.
    Done,
}

/// Spec §1: SDKs must tolerate SSE lines up to 32 MiB; beyond that the stream
/// is aborted with a `ModelStreamErrorException`.
const MAX_LINE_BYTES: usize = 32 * 1024 * 1024;

/// Incremental SSE parser.  Feed raw bytes with [`push`]; drain parsed tokens
/// with [`drain`].
pub(crate) struct SseParser {
    buf: BytesMut,
    tokens: Vec<SseToken>,
    /// Prefix of `buf` already scanned and known to contain no LF — avoids
    /// rescanning from index 0 when a long line trickles in across many chunks.
    scanned: usize,
    overflowed: bool,
}

impl SseParser {
    pub fn new() -> Self {
        Self {
            buf: BytesMut::new(),
            tokens: Vec::new(),
            scanned: 0,
            overflowed: false,
        }
    }

    /// Append raw bytes to the internal buffer and parse as many complete lines
    /// as possible.
    pub fn push(&mut self, chunk: &[u8]) {
        if self.overflowed {
            return;
        }
        self.buf.extend_from_slice(chunk);
        self.drain_lines();
        if self.buf.len() > MAX_LINE_BYTES {
            self.overflowed = true;
        }
    }

    /// True once a single line exceeded [`MAX_LINE_BYTES`]; the stream must be
    /// aborted by the caller.
    pub fn overflowed(&self) -> bool {
        self.overflowed
    }

    /// Flush any buffered data that may not end with a newline (e.g. the very
    /// last chunk before the connection closes).  Call once when the byte
    /// stream ends.
    pub fn flush(&mut self) {
        if !self.buf.is_empty() {
            // pretend there's a trailing newline so the last line is processed
            self.buf.extend_from_slice(b"\n");
            self.drain_lines();
        }
    }

    /// Take all tokens parsed so far.
    pub fn drain(&mut self) -> impl Iterator<Item = SseToken> + '_ {
        self.tokens.drain(..)
    }

    // ── internal ─────────────────────────────────────────────────────────────

    fn drain_lines(&mut self) {
        while let Some(rel) = self.buf[self.scanned..].iter().position(|&b| b == b'\n') {
            let lf_pos = self.scanned + rel;
            // Determine the line end, stripping an optional CR (CRLF support).
            let line_end = if lf_pos > 0 && self.buf[lf_pos - 1] == b'\r' {
                lf_pos - 1
            } else {
                lf_pos
            };

            // Decode the line borrowed straight from the buffer — only a final
            // `SseToken::Data` payload allocates.  A codepoint can never span
            // the LF we just found, so a complete line that fails UTF-8
            // validation is bad wire data (not a network split) — skip it.
            if let Ok(line) = std::str::from_utf8(&self.buf[..line_end]) {
                if let Some(token) = parse_line(line) {
                    self.tokens.push(token);
                }
            }
            self.buf.advance(lf_pos + 1);
            self.scanned = 0;
        }
        // No LF in the remainder — remember how far we scanned.
        self.scanned = self.buf.len();
    }
}

/// Interpret one complete SSE line, returning a token for `data:` lines.
fn parse_line(line: &str) -> Option<SseToken> {
    // Empty line → SSE event boundary (we emit per `data:` line); comment
    // lines start with ':'.  Other SSE fields (`event:`, `id:`, `retry:`) —
    // ignore via the `data:` prefix check.
    if line.is_empty() || line.starts_with(':') {
        return None;
    }
    let rest = line.strip_prefix("data:")?;
    // Strip one optional leading space per SSE spec.
    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 invalid_utf8_complete_line_is_skipped() {
        // A complete line (LF present) with invalid UTF-8 is bad wire data —
        // it must be dropped without wedging the parser.
        let mut p = SseParser::new();
        p.push(b"data: \xff\xfe bad\ndata: ok\n");
        assert_eq!(collect(&mut p), vec![SseToken::Data("ok".into())]);
        // Subsequent lines keep flowing.
        p.push(b"data: [DONE]\n");
        assert_eq!(collect(&mut p), vec![SseToken::Done]);
        assert!(!p.overflowed());
    }

    #[test]
    fn oversized_line_sets_overflow() {
        let mut p = SseParser::new();
        p.push(b"data: ");
        // Feed > 32 MiB without a newline in large chunks.
        let chunk = vec![b'a'; 8 * 1024 * 1024];
        for _ in 0..5 {
            p.push(&chunk);
        }
        assert!(p.overflowed());
        assert!(collect(&mut p).is_empty());
    }

    #[test]
    fn long_partial_line_then_completion() {
        // Exercises the incremental scan offset: many pushes without LF, then
        // the newline arrives and the whole line parses.
        let mut p = SseParser::new();
        p.push(b"data: ");
        for _ in 0..1000 {
            p.push(b"xyzzy");
        }
        assert!(collect(&mut p).is_empty());
        p.push(b"\n");
        let tokens = collect(&mut p);
        assert_eq!(tokens.len(), 1);
        match &tokens[0] {
            SseToken::Data(s) => assert_eq!(s.len(), 5000),
            other => panic!("expected Data, got {other:?}"),
        }
    }

    #[test]
    fn split_mid_utf8() {
        // U+00E9 (é) = 0xC3 0xA9 (2 bytes in UTF-8)
        let full = b"data: caf\xc3\xa9\n";
        let mut p = SseParser::new();
        // split after the first byte of the two-byte sequence
        p.push(&full[..full.len() - 2]); // everything up to 0xC3
                                         // buffer should hold the partial line; no token yet
        assert!(collect(&mut p).is_empty());
        p.push(&full[full.len() - 2..]); // 0xA9 + newline
        let tokens = collect(&mut p);
        assert_eq!(tokens, vec![SseToken::Data("caf\u{00e9}".into())]);
    }
}