audacity-sdk 0.1.0

Rust SDK for the Audacity Investments AI gateway — Amazon Bedrock Converse-compatible API surface
Documentation
//! 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,
}

/// Incremental SSE parser.  Feed raw bytes with [`push`]; drain parsed tokens
/// with [`drain`].
pub(crate) struct SseParser {
    buf: BytesMut,
    tokens: Vec<SseToken>,
}

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

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

    /// 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(lf_pos) = self.buf.iter().position(|&b| b == b'\n') {
            // 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 truncated multi-byte
            // sequence means the network read split inside a codepoint — leave
            // the bytes (including the LF we found) in the buffer and wait for
            // more data.
            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);
        }
    }
}

/// 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 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())]);
    }
}