magi-code 0.96.2

use crate::mcp::{McpError, McpResult};
use std::io::BufRead;

pub(crate) const DEFAULT_MAX_SSE_EVENT_BYTES: usize = 1_048_576;

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct SseEvent {
    pub(crate) id: Option<String>,
    pub(crate) event_type: Option<String>,
    pub(crate) retry: Option<u64>,
    pub(crate) data: String,
}

#[derive(Debug, Default)]
pub(crate) struct SseParser {
    data_lines: Vec<String>,
    id: Option<String>,
    event_type: Option<String>,
    retry: Option<u64>,
}

impl SseParser {
    pub(crate) fn push_line(&mut self, line: &str) -> Option<SseEvent> {
        let line = line.strip_suffix('\r').unwrap_or(line);
        if line.is_empty() {
            return self.dispatch();
        }
        if line.starts_with(':') {
            return None;
        }
        let (field, value) = match line.split_once(':') {
            Some((field, value)) => (field, value.strip_prefix(' ').unwrap_or(value)),
            None => (line, ""),
        };
        match field {
            "data" => self.data_lines.push(value.to_string()),
            "id" => self.id = Some(value.to_string()),
            "event" => self.event_type = Some(value.to_string()),
            "retry" => {
                if let Ok(retry) = value.parse::<u64>() {
                    self.retry = Some(retry);
                }
            }
            _ => {}
        }
        None
    }

    fn dispatch(&mut self) -> Option<SseEvent> {
        if self.data_lines.is_empty()
            && self.id.is_none()
            && self.event_type.is_none()
            && self.retry.is_none()
        {
            return None;
        }
        Some(SseEvent {
            id: self.id.take(),
            event_type: self.event_type.take(),
            retry: self.retry.take(),
            data: std::mem::take(&mut self.data_lines).join("\n"),
        })
    }
}

const SSE_BUFFER_COMPACTION_THRESHOLD: usize = 8 * 1024;

#[derive(Debug, Default)]
pub(crate) struct SseDecoder {
    buffer: Vec<u8>,
    cursor: usize,
    parser: SseParser,
    total: usize,
}

impl SseDecoder {
    pub(crate) fn push_chunk(
        &mut self,
        chunk: &[u8],
        max_bytes: usize,
    ) -> McpResult<Vec<SseEvent>> {
        self.buffer.extend_from_slice(chunk);
        let mut events = Vec::new();
        while let Some(relative_newline) = self.buffer[self.cursor..]
            .iter()
            .position(|byte| *byte == b'\n')
        {
            let line_start = self.cursor;
            let newline = line_start + relative_newline;
            self.cursor = newline + 1;
            self.total = self
                .total
                .saturating_add(newline.saturating_sub(line_start) + 1);
            if self.total > max_bytes {
                return Err(McpError::Transport(format!(
                    "MCP SSE event exceeded {max_bytes} bytes"
                )));
            }
            let line = std::str::from_utf8(&self.buffer[line_start..newline])
                .map_err(McpError::transport)?;
            if let Some(event) = self.parser.push_line(line) {
                self.total = 0;
                events.push(event);
            }
        }
        self.compact_buffer();
        if self.pending_len().saturating_add(self.total) > max_bytes {
            return Err(McpError::Transport(format!(
                "MCP SSE event exceeded {max_bytes} bytes"
            )));
        }
        Ok(events)
    }

    pub(crate) fn finish(&mut self, max_bytes: usize) -> McpResult<Option<SseEvent>> {
        if self.pending_len() == 0 {
            return Ok(self.parser.dispatch());
        }
        if self.pending_len().saturating_add(self.total) > max_bytes {
            return Err(McpError::Transport(format!(
                "MCP SSE event exceeded {max_bytes} bytes"
            )));
        }
        let line = std::str::from_utf8(&self.buffer[self.cursor..]).map_err(McpError::transport)?;
        let event = self
            .parser
            .push_line(line)
            .or_else(|| self.parser.dispatch());
        self.buffer.clear();
        self.cursor = 0;
        Ok(event)
    }

    fn pending_len(&self) -> usize {
        self.buffer.len().saturating_sub(self.cursor)
    }

    fn compact_buffer(&mut self) {
        if self.cursor == 0 {
            return;
        }
        if self.cursor == self.buffer.len() {
            self.buffer.clear();
            self.cursor = 0;
        } else if self.cursor >= SSE_BUFFER_COMPACTION_THRESHOLD
            && self.cursor >= self.buffer.len() / 2
        {
            self.buffer.drain(..self.cursor);
            self.cursor = 0;
        }
    }
}

pub(crate) fn read_sse_event(
    reader: &mut impl BufRead,
    max_bytes: usize,
) -> McpResult<Option<SseEvent>> {
    let mut parser = SseParser::default();
    let mut line = Vec::new();
    let mut total = 0usize;
    loop {
        if !read_sse_line(reader, &mut line, &mut total, max_bytes)? {
            return Ok(parser.dispatch());
        }
        while line.last() == Some(&b'\r') {
            line.pop();
        }
        let line = std::str::from_utf8(&line).map_err(McpError::transport)?;
        if let Some(event) = parser.push_line(line) {
            return Ok(Some(event));
        }
    }
}

fn read_sse_line(
    reader: &mut impl BufRead,
    line: &mut Vec<u8>,
    total: &mut usize,
    max_bytes: usize,
) -> McpResult<bool> {
    line.clear();
    loop {
        let available = reader.fill_buf().map_err(McpError::transport)?;
        if available.is_empty() {
            return Ok(!line.is_empty());
        }

        let mut consumed = 0;
        let mut complete = false;
        for byte in available {
            if *total == max_bytes {
                break;
            }
            *total += 1;
            consumed += 1;
            if *byte == b'\n' {
                complete = true;
                break;
            }
            line.push(*byte);
        }
        reader.consume(consumed);

        if complete {
            return Ok(true);
        }
        if *total == max_bytes {
            if reader.fill_buf().map_err(McpError::transport)?.is_empty() {
                return Ok(!line.is_empty());
            }
            return Err(McpError::Transport(format!(
                "MCP SSE event exceeded {max_bytes} bytes"
            )));
        }
    }
}

pub(crate) fn jsonrpc_message_from_event(
    event: &SseEvent,
) -> McpResult<Option<crate::mcp::jsonrpc::JsonRpcMessage>> {
    if event.data.trim().is_empty() {
        return Ok(None);
    }
    serde_json::from_str(&event.data)
        .map(Some)
        .map_err(|error| McpError::Transport(format!("MCP SSE JSON parse failed: {error}")))
}