crank/crates/crank-adapter-rest/src/sse.rs

use std::collections::BTreeMap;

use futures_util::StreamExt;
use reqwest::header::HeaderMap;
use serde_json::{Value, json};
use tokio::time::{Duration, Instant, timeout_at};

use crate::RestAdapterError;

pub async fn collect_sse_window(
    response: reqwest::Response,
    window_duration_ms: u64,
    max_items: Option<u32>,
) -> Result<(u16, BTreeMap<String, String>, Value), RestAdapterError> {
    let status = response.status();
    let headers = normalize_headers(response.headers());
    let deadline = Instant::now() + Duration::from_millis(window_duration_ms);
    let mut stream = response.bytes_stream();
    let mut buffer = String::new();
    let mut items = Vec::new();
    let mut done = true;

    loop {
        if max_items.is_some_and(|limit| items.len() >= limit as usize) {
            done = false;
            break;
        }

        let next_chunk = match timeout_at(deadline, stream.next()).await {
            Ok(next_chunk) => next_chunk,
            Err(_) => break,
        };

        let Some(next_chunk) = next_chunk else {
            break;
        };

        let chunk = next_chunk?;
        buffer.push_str(&String::from_utf8_lossy(&chunk));

        while let Some(event_end) = find_event_boundary(&buffer) {
            let event = buffer[..event_end].to_owned();
            let boundary_len = boundary_length(&buffer[event_end..]);
            buffer = buffer[event_end + boundary_len..].to_owned();

            if let Some(item) = parse_sse_event(&event)? {
                items.push(item);

                if max_items.is_some_and(|limit| items.len() >= limit as usize) {
                    done = false;
                    break;
                }
            }
        }

        if !done {
            break;
        }
    }

    Ok((
        status.as_u16(),
        headers,
        json!({
            "items": items,
            "done": done
        }),
    ))
}

fn parse_sse_event(raw: &str) -> Result<Option<Value>, RestAdapterError> {
    let mut data_lines = Vec::new();

    for line in raw.lines() {
        if line.is_empty() || line.starts_with(':') {
            continue;
        }

        if let Some(value) = line.strip_prefix("data:") {
            data_lines.push(value.trim_start().to_owned());
        }
    }

    if data_lines.is_empty() {
        return Ok(None);
    }

    let payload = data_lines.join("\n");
    if payload.is_empty() {
        return Ok(None);
    }

    serde_json::from_str::<Value>(&payload)
        .map(Some)
        .or_else(|_| {
            if payload.starts_with('{') || payload.starts_with('[') {
                Err(RestAdapterError::InvalidSseEvent)
            } else {
                Ok(Some(Value::String(payload)))
            }
        })
}

fn find_event_boundary(buffer: &str) -> Option<usize> {
    buffer
        .find("\r\n\r\n")
        .or_else(|| buffer.find("\n\n"))
        .or_else(|| buffer.find("\r\r"))
}

fn boundary_length(boundary: &str) -> usize {
    if boundary.starts_with("\r\n\r\n") {
        4
    } else {
        2
    }
}

fn normalize_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
    headers
        .iter()
        .filter_map(|(name, value)| {
            value
                .to_str()
                .ok()
                .map(|value| (name.as_str().to_owned(), value.to_owned()))
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::parse_sse_event;
    use crate::RestAdapterError;

    #[test]
    fn parses_json_event_payload() {
        let event = "event: message\ndata: {\"message\":\"ok\"}\n\n";

        let parsed = parse_sse_event(event).unwrap();

        assert_eq!(parsed, Some(json!({ "message": "ok" })));
    }

    #[test]
    fn ignores_comment_only_events() {
        let parsed = parse_sse_event(": keepalive\n\n").unwrap();

        assert_eq!(parsed, None);
    }

    #[test]
    fn rejects_malformed_json_like_payload() {
        let error = parse_sse_event("data: {broken-json}\n\n").unwrap_err();

        assert!(matches!(error, RestAdapterError::InvalidSseEvent));
    }
}