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, ) -> Result<(u16, BTreeMap, 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, 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::(&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 { 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 { 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)); } }