rust-mcp-transport 2.0.0

use bytes::{Bytes, BytesMut};
use std::collections::HashMap;

use super::SseEvent;
const BUFFER_CAPACITY: usize = 1024;

/// A parser for Server-Sent Events (SSE) that processes incoming byte chunks into `SseEvent`s.
/// This Parser is specifically designed for MCP messages and with no multi-line data support
///
/// This struct maintains a buffer to accumulate incoming data and parses it into SSE events
/// based on the SSE protocol. It handles fields like `event`, `data`, and `id` as defined
/// in the SSE specification.
#[derive(Debug)]
pub struct SseParser {
    pub buffer: BytesMut,
}

impl SseParser {
    /// Creates a new `SseParser` with an empty buffer pre-allocated to a default capacity.
    ///
    /// The buffer is initialized with a capacity of `BUFFER_CAPACITY` to
    /// optimize for typical SSE message sizes.
    ///
    /// # Returns
    /// A new `SseParser` instance with an empty buffer.
    pub fn new() -> Self {
        Self {
            buffer: BytesMut::with_capacity(BUFFER_CAPACITY),
        }
    }

    /// Processes a new chunk of bytes and parses it into a vector of `SseEvent`s.
    ///
    /// This method appends the incoming `bytes` to the internal buffer, splits it into
    /// complete lines (delimited by `\n`), and parses each line according to the SSE
    /// protocol. It supports `event`, `id`, and `data` fields, as well as comments
    /// (lines starting with `:`). Empty lines are skipped, and incomplete lines remain
    /// in the buffer for future processing.
    ///
    /// # Parameters
    /// - `bytes`: The incoming chunk of bytes to parse.
    ///
    /// # Returns
    /// A vector of `SseEvent`s parsed from the complete lines in the buffer. If no
    /// complete events are found, an empty vector is returned.
    pub fn process_new_chunk(&mut self, bytes: Bytes) -> Vec<SseEvent> {
        self.buffer.extend_from_slice(&bytes);

        // Collect complete lines (ending in \n)-keep ALL lines, including empty ones for \n\n detection
        let mut lines = Vec::new();
        while let Some(pos) = self.buffer.iter().position(|&b| b == b'\n') {
            let line = self.buffer.split_to(pos + 1).freeze();
            lines.push(line);
        }

        let mut events = Vec::new();
        let mut current_message_lines: Vec<Bytes> = Vec::new();

        for line in lines {
            current_message_lines.push(line);

            // Check if we've hit a double newline (end of message)
            if current_message_lines.len() >= 2
                && current_message_lines
                    .last()
                    .is_some_and(|b| b.as_ref() == b"\n")
            {
                // Process the complete message (exclude the last empty lines for parsing)
                let message_lines: Vec<_> = current_message_lines
                    .drain(..current_message_lines.len() - 1)
                    .filter(|l| l.as_ref() != b"\n") // Filter internal empties
                    .collect();

                if let Some(event) = self.parse_sse_message(&message_lines) {
                    events.push(event);
                }
            }
        }

        // Put back any incomplete message, preserving partial bytes
        // that were after the last \n (self.buffer after the while loop).
        // this prevents data loss when data is received in multiple chunks (like from Cloudflare workers (Issue: #199))
        if !current_message_lines.is_empty() {
            let mut new_buf = BytesMut::new();
            for line in current_message_lines.into_iter() {
                new_buf.extend_from_slice(&line);
            }
            new_buf.extend_from_slice(&self.buffer);
            self.buffer = new_buf;
        }

        events
    }

    fn parse_sse_message(&self, lines: &[Bytes]) -> Option<SseEvent> {
        let mut fields: HashMap<String, String> = HashMap::new();
        let mut data_parts: Vec<String> = Vec::new();
        // Track whether we saw a `data:` line, even with an empty value, so we
        // can distinguish "no data field at all" from "data: " (empty data).
        let mut saw_data_line = false;

        for line_bytes in lines {
            let line_str = String::from_utf8_lossy(line_bytes);

            // Skip comments and empty lines
            if line_str.is_empty() || line_str.starts_with(':') {
                continue;
            }

            // Per the SSE spec, "data: " and "data:" are both valid; the value
            // may be empty. We accept either form. The trailing newline is
            // retained for backwards compatibility with existing consumers.
            let (key, value) = if let Some(value) = line_str.strip_prefix("data: ") {
                ("data", value.to_string())
            } else if line_str.trim_end_matches('\n') == "data:"
                || line_str.trim_end_matches('\n').is_empty() && line_str.starts_with("data:")
            {
                ("data", String::new())
            } else if let Some(value) = line_str.strip_prefix("event: ") {
                ("event", value.trim().to_string())
            } else if let Some(value) = line_str.strip_prefix("id: ") {
                ("id", value.trim().to_string())
            } else if let Some(value) = line_str.strip_prefix("retry: ") {
                ("retry", value.trim().to_string())
            } else {
                // Invalid line; skip
                continue;
            };

            if key == "data" {
                saw_data_line = true;
                // Skip empty data values. The legacy behavior preserves a
                // trailing newline in the value (e.g. "hello\n" for the
                // line "data: hello\n"), so a bare "\n" represents an
                // empty data line and must not be pushed.
                if !value.is_empty() && value != "\n" {
                    data_parts.push(value);
                }
            } else {
                fields.insert(key.to_string(), value);
            }
        }

        // Build data (concat multi-line data with \n) , should not occur in MCP tho
        let data = if data_parts.is_empty() {
            None
        } else {
            let full_data = data_parts.join("\n");
            Some(Bytes::copy_from_slice(full_data.as_bytes())) // Use copy_from_slice for efficiency
        };

        // Get event (default to None)
        let event = fields.get("event").cloned();
        let id = fields.get("id").cloned();
        let retry = fields
            .get("retry")
            .and_then(|r| r.trim().parse::<u64>().ok());

        // Emit a priming event (no data) if it carries an `id` or `retry`
        // (SEP-1699 resumability checkpoint / reconnection timing). Skip pure
        // junk messages with no fields at all.
        if data.is_none() && id.is_none() && retry.is_none() && event.is_none() && !saw_data_line {
            return None;
        }

        Some(SseEvent {
            event,
            data,
            id,
            retry,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;

    #[test]
    fn test_single_data_event() {
        let mut parser = SseParser::new();
        let input = Bytes::from("data: hello\n\n");
        let events = parser.process_new_chunk(input);

        assert_eq!(events.len(), 1);
        assert_eq!(
            events[0].data.as_deref(),
            Some(Bytes::from("hello\n").as_ref())
        );
        assert!(events[0].event.is_none());
        assert!(events[0].id.is_none());
    }

    #[test]
    fn test_event_with_id_and_data() {
        let mut parser = SseParser::new();
        let input = Bytes::from("event: message\nid: 123\ndata: hello\n\n");
        let events = parser.process_new_chunk(input);

        assert_eq!(events.len(), 1);
        assert_eq!(events[0].event.as_deref(), Some("message"));
        assert_eq!(events[0].id.as_deref(), Some("123"));
        assert_eq!(
            events[0].data.as_deref(),
            Some(Bytes::from("hello\n").as_ref())
        );
    }

    #[test]
    fn test_event_chunks_in_different_orders() {
        let mut parser = SseParser::new();
        let input = Bytes::from("data: hello\nevent: message\nid: 123\n\n");
        let events = parser.process_new_chunk(input);

        assert_eq!(events.len(), 1);
        assert_eq!(events[0].event.as_deref(), Some("message"));
        assert_eq!(events[0].id.as_deref(), Some("123"));
        assert_eq!(
            events[0].data.as_deref(),
            Some(Bytes::from("hello\n").as_ref())
        );
    }

    #[test]
    fn test_comment_line_ignored() {
        let mut parser = SseParser::new();
        let input = Bytes::from(": this is a comment\n\n");
        let events = parser.process_new_chunk(input);
        assert_eq!(events.len(), 0);
    }

    #[test]
    fn test_event_with_empty_data() {
        // SEP-1699: A `data:` line with no value is a valid SSE field. When
        // accompanied by `id` and/or `retry` it functions as a priming event
        // for resumability and reconnection. Standalone `data:` with no other
        // fields is preserved as an empty-data event.
        let mut parser = SseParser::new();
        let input = Bytes::from("data:\n\n");
        let events = parser.process_new_chunk(input);
        assert_eq!(events.len(), 1);
        assert!(events[0].data.is_none());
        assert!(events[0].id.is_none());
        assert!(events[0].retry.is_none());
    }

    #[test]
    fn test_priming_event_with_id_and_retry() {
        // SEP-1699 priming event: `id` + `retry` + empty `data` line, no result.
        let mut parser = SseParser::new();
        let input = Bytes::from("id: event-1\nretry: 500\ndata: \n\n");
        let events = parser.process_new_chunk(input);
        assert_eq!(events.len(), 1);
        assert_eq!(events[0].id.as_deref(), Some("event-1"));
        assert_eq!(events[0].retry, Some(500));
        assert!(events[0].data.is_none());
    }

    #[test]
    fn test_partial_chunks() {
        let mut parser = SseParser::new();

        let part1 = Bytes::from("data: hello");
        let part2 = Bytes::from(" world\n\n");

        let events1 = parser.process_new_chunk(part1);
        assert_eq!(events1.len(), 0); // incomplete

        let events2 = parser.process_new_chunk(part2);
        assert_eq!(events2.len(), 1);
        assert_eq!(
            events2[0].data.as_deref(),
            Some(Bytes::from("hello world\n").as_ref())
        );
    }

    #[test]
    fn test_data_line_across_chunks_with_preceding_event_line() {
        let mut parser = SseParser::new();

        let part1 = Bytes::from("event: message\ndata: ");
        let part2 = Bytes::from("hello\n\n");

        let events1 = parser.process_new_chunk(part1);
        assert_eq!(events1.len(), 0);

        let events2 = parser.process_new_chunk(part2);
        assert_eq!(events2.len(), 1);
        assert_eq!(events2[0].event.as_deref(), Some("message"));
        assert_eq!(
            events2[0].data.as_deref(),
            Some(Bytes::from("hello\n").as_ref())
        );
    }

    #[test]
    fn test_data_line_across_multiple_chunks_with_preceding_event() {
        let mut parser = SseParser::new();

        let part1 = Bytes::from("event: message\ndata: alpha bra");
        let part2 = Bytes::from("vo charlie del");
        let part3 = Bytes::from("ta echo foxtrot\n\n");

        let events1 = parser.process_new_chunk(part1);
        assert_eq!(events1.len(), 0);

        let events2 = parser.process_new_chunk(part2);
        assert_eq!(events2.len(), 0);

        let events3 = parser.process_new_chunk(part3);
        assert_eq!(events3.len(), 1);
        assert_eq!(events3[0].event.as_deref(), Some("message"));
        assert_eq!(
            events3[0].data.as_deref(),
            Some(Bytes::from("alpha bravo charlie delta echo foxtrot\n").as_ref())
        );
    }

    #[test]
    fn test_multiple_events_split_across_chunks() {
        let mut parser = SseParser::new();

        let part1 = Bytes::from("data: one\n\ndata: t");
        let part2 = Bytes::from("wo\n\n");

        let events1 = parser.process_new_chunk(part1);
        assert_eq!(events1.len(), 1);
        assert_eq!(
            events1[0].data.as_deref(),
            Some(Bytes::from("one\n").as_ref())
        );

        let events2 = parser.process_new_chunk(part2);
        assert_eq!(events2.len(), 1);
        assert_eq!(
            events2[0].data.as_deref(),
            Some(Bytes::from("two\n").as_ref())
        );
    }

    #[test]
    fn test_malformed_lines() {
        let mut parser = SseParser::new();
        let input = Bytes::from("something invalid\ndata: ok\n\n");

        let events = parser.process_new_chunk(input);

        assert_eq!(events.len(), 1);
        assert_eq!(
            events[0].data.as_deref(),
            Some(Bytes::from("ok\n").as_ref())
        );
    }

    #[test]
    fn test_multiple_events_in_one_chunk() {
        let mut parser = SseParser::new();
        let input = Bytes::from("data: first\n\ndata: second\n\n");
        let events = parser.process_new_chunk(input);

        assert_eq!(events.len(), 2);
        assert_eq!(
            events[0].data.as_deref(),
            Some(Bytes::from("first\n").as_ref())
        );
        assert_eq!(
            events[1].data.as_deref(),
            Some(Bytes::from("second\n").as_ref())
        );
    }

    #[test]
    fn test_basic_sse_event() {
        let mut parser = SseParser::new();
        let input = Bytes::from("event: message\ndata: Hello\nid: 1\nretry: 5000\n\n");

        let events = parser.process_new_chunk(input);

        assert_eq!(events.len(), 1);

        let event = &events[0];
        assert_eq!(event.event.as_deref(), Some("message"));
        assert_eq!(event.data.as_deref(), Some(Bytes::from("Hello\n").as_ref()));
        assert_eq!(event.id.as_deref(), Some("1"));
        assert_eq!(event.retry, Some(5000));
    }
}