rustapi_core/

sse.rs

//! Server-Sent Events (SSE) response types for RustAPI
//!
//! This module provides types for streaming Server-Sent Events to clients.
//! SSE is ideal for real-time updates like notifications, live feeds, and progress updates.
//!
//! # Example
//!
//! ```rust,ignore
//! use rustapi_core::sse::{Sse, SseEvent, KeepAlive};
//! use futures_util::stream;
//! use std::time::Duration;
//!
//! async fn events() -> Sse<impl Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
//!     let stream = stream::iter(vec![
//!         Ok(SseEvent::new("Hello")),
//!         Ok(SseEvent::new("World").event("greeting")),
//!     ]);
//!     Sse::new(stream)
//!         .keep_alive(KeepAlive::new().interval(Duration::from_secs(15)))
//! }
//! ```
//!
//! # Keep-Alive Support
//!
//! SSE connections can be kept alive by sending periodic comments:
//!
//! ```rust,ignore
//! use rustapi_core::sse::{Sse, SseEvent, KeepAlive};
//! use std::time::Duration;
//!
//! async fn events() -> impl IntoResponse {
//!     let stream = async_stream::stream! {
//!         for i in 0..10 {
//!             yield Ok::<_, std::convert::Infallible>(
//!                 SseEvent::new(format!("Event {}", i))
//!             );
//!             tokio::time::sleep(Duration::from_secs(1)).await;
//!         }
//!     };
//!
//!     Sse::new(stream)
//!         .keep_alive(KeepAlive::new()
//!             .interval(Duration::from_secs(30))
//!             .text("ping"))
//! }
//! ```

use bytes::Bytes;
use futures_util::Stream;
use http::{header, StatusCode};

use pin_project_lite::pin_project;
use rustapi_openapi::{MediaType, Operation, ResponseModifier, ResponseSpec, SchemaRef};
use std::fmt::Write;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;

use crate::response::{IntoResponse, Response};

/// A Server-Sent Event
///
/// SSE events follow the format specified in the W3C Server-Sent Events specification.
/// Each event can have:
/// - `data`: The event data (required)
/// - `event`: The event type/name (optional)
/// - `id`: The event ID for reconnection (optional)
/// - `retry`: Reconnection time in milliseconds (optional)
/// - `comment`: A comment line (optional, not visible to most clients)
#[derive(Debug, Clone, Default)]
pub struct SseEvent {
    /// The event data
    pub data: String,
    /// The event type/name
    pub event: Option<String>,
    /// The event ID
    pub id: Option<String>,
    /// Reconnection time in milliseconds
    pub retry: Option<u64>,
    /// Comment line
    comment: Option<String>,
}

impl SseEvent {
    /// Create a new SSE event with the given data
    pub fn new(data: impl Into<String>) -> Self {
        Self {
            data: data.into(),
            event: None,
            id: None,
            retry: None,
            comment: None,
        }
    }

    /// Create an SSE comment (keep-alive)
    ///
    /// Comments are lines starting with `:` and are typically used for keep-alive.
    pub fn comment(text: impl Into<String>) -> Self {
        Self {
            data: String::new(),
            event: None,
            id: None,
            retry: None,
            comment: Some(text.into()),
        }
    }

    /// Set the event type/name
    pub fn event(mut self, event: impl Into<String>) -> Self {
        self.event = Some(event.into());
        self
    }

    /// Set the event ID
    pub fn id(mut self, id: impl Into<String>) -> Self {
        self.id = Some(id.into());
        self
    }

    /// Set the retry time in milliseconds
    pub fn retry(mut self, retry: u64) -> Self {
        self.retry = Some(retry);
        self
    }

    /// Set JSON data (serializes the value)
    pub fn json_data<T: serde::Serialize>(data: &T) -> Result<Self, serde_json::Error> {
        Ok(Self::new(serde_json::to_string(data)?))
    }

    /// Format the event as an SSE message
    ///
    /// The format follows the SSE specification:
    /// - Lines starting with "event:" specify the event type
    /// - Lines starting with "id:" specify the event ID
    /// - Lines starting with "retry:" specify the reconnection time
    /// - Lines starting with "data:" contain the event data
    /// - Lines starting with ":" are comments
    /// - Events are terminated with a blank line
    pub fn to_sse_string(&self) -> String {
        let mut output = String::new();

        // Comment (for keep-alive)
        if let Some(ref comment) = self.comment {
            writeln!(output, ": {}", comment).unwrap();
            output.push('\n');
            return output;
        }

        // Event type
        if let Some(ref event) = self.event {
            writeln!(output, "event: {}", event).unwrap();
        }

        // Event ID
        if let Some(ref id) = self.id {
            writeln!(output, "id: {}", id).unwrap();
        }

        // Retry time
        if let Some(retry) = self.retry {
            writeln!(output, "retry: {}", retry).unwrap();
        }

        // Data - handle multi-line data by prefixing each line with "data: "
        for line in self.data.lines() {
            writeln!(output, "data: {}", line).unwrap();
        }

        // If data is empty, still send an empty data line
        if self.data.is_empty() && self.comment.is_none() {
            writeln!(output, "data:").unwrap();
        }

        // Empty line to terminate the event
        output.push('\n');

        output
    }

    /// Convert the event to bytes
    pub fn to_bytes(&self) -> Bytes {
        Bytes::from(self.to_sse_string())
    }
}

/// Keep-alive configuration for SSE connections
///
/// Keep-alive sends periodic comments to prevent connection timeouts.
///
/// # Example
///
/// ```rust,ignore
/// use rustapi_core::sse::KeepAlive;
/// use std::time::Duration;
///
/// let keep_alive = KeepAlive::new()
///     .interval(Duration::from_secs(30))
///     .text("ping");
/// ```
#[derive(Debug, Clone)]
pub struct KeepAlive {
    /// Interval between keep-alive messages
    interval: Duration,
    /// Text to send as keep-alive comment
    text: String,
}

impl Default for KeepAlive {
    fn default() -> Self {
        Self {
            interval: Duration::from_secs(15),
            text: "keep-alive".to_string(),
        }
    }
}

impl KeepAlive {
    /// Create a new keep-alive configuration with default settings
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the keep-alive interval
    pub fn interval(mut self, interval: Duration) -> Self {
        self.interval = interval;
        self
    }

    /// Set the keep-alive text
    pub fn text(mut self, text: impl Into<String>) -> Self {
        self.text = text.into();
        self
    }

    /// Get the interval
    pub fn get_interval(&self) -> Duration {
        self.interval
    }

    /// Create the keep-alive event
    pub fn event(&self) -> SseEvent {
        SseEvent::comment(&self.text)
    }
}

/// Server-Sent Events response wrapper
///
/// Wraps a stream of `SseEvent` items and converts them to an SSE response.
///
/// # Example
///
/// ```rust,ignore
/// use rustapi_core::sse::{Sse, SseEvent, KeepAlive};
/// use futures_util::stream;
/// use std::time::Duration;
///
/// async fn events() -> Sse<impl Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
///     let stream = stream::iter(vec![
///         Ok(SseEvent::new("Hello")),
///         Ok(SseEvent::new("World").event("greeting")),
///     ]);
///     Sse::new(stream)
///         .keep_alive(KeepAlive::new().interval(Duration::from_secs(30)))
/// }
/// ```
pub struct Sse<S> {
    stream: S,
    keep_alive: Option<KeepAlive>,
}

impl<S> Sse<S> {
    /// Create a new SSE response from a stream
    pub fn new(stream: S) -> Self {
        Self {
            stream,
            keep_alive: None,
        }
    }

    /// Set the keep-alive configuration
    ///
    /// When set, the server will send periodic comments to keep the connection alive.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use rustapi_core::sse::{Sse, KeepAlive};
    /// use std::time::Duration;
    ///
    /// Sse::new(stream)
    ///     .keep_alive(KeepAlive::new().interval(Duration::from_secs(30)))
    /// ```
    pub fn keep_alive(mut self, config: KeepAlive) -> Self {
        self.keep_alive = Some(config);
        self
    }

    /// Get the keep-alive configuration
    pub fn get_keep_alive(&self) -> Option<&KeepAlive> {
        self.keep_alive.as_ref()
    }
}

// Stream that merges SSE events with keep-alive events
pin_project! {
    /// A stream that combines SSE events with keep-alive messages
    pub struct SseStream<S> {
        #[pin]
        inner: S,
        keep_alive: Option<KeepAlive>,
        #[pin]
        keep_alive_timer: Option<tokio::time::Interval>,
    }
}

impl<S, E> Stream for SseStream<S>
where
    S: Stream<Item = Result<SseEvent, E>>,
{
    type Item = Result<Bytes, E>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.project();

        // First, check if there's an event ready from the inner stream
        match this.inner.poll_next(cx) {
            Poll::Ready(Some(Ok(event))) => {
                return Poll::Ready(Some(Ok(event.to_bytes())));
            }
            Poll::Ready(Some(Err(e))) => {
                return Poll::Ready(Some(Err(e)));
            }
            Poll::Ready(None) => {
                return Poll::Ready(None);
            }
            Poll::Pending => {}
        }

        // Check keep-alive timer
        if let Some(mut timer) = this.keep_alive_timer.as_pin_mut() {
            if timer.poll_tick(cx).is_ready() {
                if let Some(keep_alive) = this.keep_alive {
                    let event = keep_alive.event();
                    return Poll::Ready(Some(Ok(event.to_bytes())));
                }
            }
        }

        Poll::Pending
    }
}

// For now, we'll implement IntoResponse by collecting the stream into a single response
// This is a simplified implementation that works with the current Response type (Full<Bytes>)
// A full streaming implementation would require changes to the Response type
impl<S, E> IntoResponse for Sse<S>
where
    S: Stream<Item = Result<SseEvent, E>> + Send + 'static,
    E: std::error::Error + Send + Sync + 'static,
{
    fn into_response(self) -> Response {
        let timer = self.keep_alive.as_ref().map(|k| {
            let mut interval = tokio::time::interval(k.interval);
            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
            interval
        });

        let stream = SseStream {
            inner: self.stream,
            keep_alive: self.keep_alive,
            keep_alive_timer: timer,
        };

        use futures_util::StreamExt;
        let stream =
            stream.map(|res| res.map_err(|e| crate::error::ApiError::internal(e.to_string())));
        let body = crate::response::Body::from_stream(stream);

        http::Response::builder()
            .status(StatusCode::OK)
            .header(header::CONTENT_TYPE, "text/event-stream")
            .header(header::CACHE_CONTROL, "no-cache")
            .header(header::CONNECTION, "keep-alive")
            .header("X-Accel-Buffering", "no") // Disable nginx buffering
            .body(body)
            .unwrap()
    }
}

// OpenAPI support: ResponseModifier for SSE streams
impl<S> ResponseModifier for Sse<S> {
    fn update_response(op: &mut Operation) {
        let mut content = std::collections::HashMap::new();
        content.insert(
            "text/event-stream".to_string(),
            MediaType {
                schema: SchemaRef::Inline(serde_json::json!({
                    "type": "string",
                    "description": "Server-Sent Events stream. Events follow the SSE format: 'event: <type>\\ndata: <json>\\n\\n'",
                    "example": "event: message\ndata: {\"id\": 1, \"text\": \"Hello\"}\n\n"
                })),
            },
        );

        let response = ResponseSpec {
            description: "Server-Sent Events stream for real-time updates".to_string(),
            content: Some(content),
        };
        op.responses.insert("200".to_string(), response);
    }
}

/// Collect all SSE events from a stream into a single response body
///
/// This is useful for testing or when you know the stream is finite.
pub async fn collect_sse_events<S, E>(stream: S) -> Result<Bytes, E>
where
    S: Stream<Item = Result<SseEvent, E>> + Send,
{
    use futures_util::StreamExt;

    let mut buffer = Vec::new();
    futures_util::pin_mut!(stream);

    while let Some(result) = stream.next().await {
        let event = result?;
        buffer.extend_from_slice(&event.to_bytes());
    }

    Ok(Bytes::from(buffer))
}

/// Create an SSE response from a synchronous iterator of events
///
/// This is a convenience function for simple cases with pre-computed events.
///
/// # Example
///
/// ```rust,ignore
/// use rustapi_core::sse::{sse_response, SseEvent};
///
/// async fn handler() -> Response {
///     sse_response(vec![
///         SseEvent::new("Hello"),
///         SseEvent::new("World").event("greeting"),
///     ])
/// }
/// ```
pub fn sse_response<I>(events: I) -> Response
where
    I: IntoIterator<Item = SseEvent>,
{
    let mut buffer = String::new();
    for event in events {
        buffer.push_str(&event.to_sse_string());
    }

    http::Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "text/event-stream")
        .header(header::CACHE_CONTROL, "no-cache")
        .header(header::CONNECTION, "keep-alive")
        .header("X-Accel-Buffering", "no")
        .body(crate::response::Body::from(buffer))
        .unwrap()
}

/// Helper function to create an SSE response from an iterator of events
///
/// This is useful for simple cases where you have a fixed set of events.
pub fn sse_from_iter<I, E>(
    events: I,
) -> Sse<futures_util::stream::Iter<std::vec::IntoIter<Result<SseEvent, E>>>>
where
    I: IntoIterator<Item = Result<SseEvent, E>>,
{
    use futures_util::stream;
    let vec: Vec<_> = events.into_iter().collect();
    Sse::new(stream::iter(vec))
}

#[cfg(test)]
mod tests {
    use super::*;
    use proptest::prelude::*;

    #[test]
    fn test_sse_event_basic() {
        let event = SseEvent::new("Hello, World!");
        let output = event.to_sse_string();
        assert_eq!(output, "data: Hello, World!\n\n");
    }

    #[test]
    fn test_sse_event_with_event_type() {
        let event = SseEvent::new("Hello").event("greeting");
        let output = event.to_sse_string();
        assert!(output.contains("event: greeting\n"));
        assert!(output.contains("data: Hello\n"));
    }

    #[test]
    fn test_sse_event_with_id() {
        let event = SseEvent::new("Hello").id("123");
        let output = event.to_sse_string();
        assert!(output.contains("id: 123\n"));
        assert!(output.contains("data: Hello\n"));
    }

    #[test]
    fn test_sse_event_with_retry() {
        let event = SseEvent::new("Hello").retry(5000);
        let output = event.to_sse_string();
        assert!(output.contains("retry: 5000\n"));
        assert!(output.contains("data: Hello\n"));
    }

    #[test]
    fn test_sse_event_multiline_data() {
        let event = SseEvent::new("Line 1\nLine 2\nLine 3");
        let output = event.to_sse_string();
        assert!(output.contains("data: Line 1\n"));
        assert!(output.contains("data: Line 2\n"));
        assert!(output.contains("data: Line 3\n"));
    }

    #[test]
    fn test_sse_event_full() {
        let event = SseEvent::new("Hello").event("message").id("1").retry(3000);
        let output = event.to_sse_string();

        // Check all fields are present
        assert!(output.contains("event: message\n"));
        assert!(output.contains("id: 1\n"));
        assert!(output.contains("retry: 3000\n"));
        assert!(output.contains("data: Hello\n"));

        // Check it ends with double newline
        assert!(output.ends_with("\n\n"));
    }

    #[test]
    fn test_sse_response_headers() {
        use futures_util::stream;

        let events: Vec<Result<SseEvent, std::convert::Infallible>> =
            vec![Ok(SseEvent::new("test"))];
        let sse = Sse::new(stream::iter(events));
        let response = sse.into_response();

        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response.headers().get(header::CONTENT_TYPE).unwrap(),
            "text/event-stream"
        );
        assert_eq!(
            response.headers().get(header::CACHE_CONTROL).unwrap(),
            "no-cache"
        );
        assert_eq!(
            response.headers().get(header::CONNECTION).unwrap(),
            "keep-alive"
        );
    }

    // **Feature: phase3-batteries-included, Property 20: SSE response format**
    //
    // For any stream of SseEvent items, `Sse<S>` SHALL produce a response with
    // `Content-Type: text/event-stream` and body formatted according to SSE specification.
    //
    // **Validates: Requirements 6.3**
    proptest! {
        #![proptest_config(ProptestConfig::with_cases(100))]

        #[test]
        fn prop_sse_response_format(
            // Generate random event data (alphanumeric to avoid special chars)
            data in "[a-zA-Z0-9 ]{1,50}",
            // Optional event type
            event_type in proptest::option::of("[a-zA-Z][a-zA-Z0-9_]{0,20}"),
            // Optional event ID
            event_id in proptest::option::of("[a-zA-Z0-9]{1,10}"),
            // Optional retry time
            retry_time in proptest::option::of(1000u64..60000u64),
        ) {
            use futures_util::stream;

            // Build the SSE event with optional fields
            let mut event = SseEvent::new(data.clone());
            if let Some(ref et) = event_type {
                event = event.event(et.clone());
            }
            if let Some(ref id) = event_id {
                event = event.id(id.clone());
            }
            if let Some(retry) = retry_time {
                event = event.retry(retry);
            }

            // Verify the SSE string format
            let sse_string = event.to_sse_string();

            // Property 1: SSE string must end with double newline (event terminator)
            prop_assert!(
                sse_string.ends_with("\n\n"),
                "SSE event must end with double newline, got: {:?}",
                sse_string
            );

            // Property 2: Data must be present with "data: " prefix
            prop_assert!(
                sse_string.contains(&format!("data: {}", data)),
                "SSE event must contain data field with 'data: ' prefix"
            );

            // Property 3: If event type is set, it must be present with "event: " prefix
            if let Some(ref et) = event_type {
                prop_assert!(
                    sse_string.contains(&format!("event: {}", et)),
                    "SSE event must contain event type with 'event: ' prefix"
                );
            }

            // Property 4: If ID is set, it must be present with "id: " prefix
            if let Some(ref id) = event_id {
                prop_assert!(
                    sse_string.contains(&format!("id: {}", id)),
                    "SSE event must contain ID with 'id: ' prefix"
                );
            }

            // Property 5: If retry is set, it must be present with "retry: " prefix
            if let Some(retry) = retry_time {
                prop_assert!(
                    sse_string.contains(&format!("retry: {}", retry)),
                    "SSE event must contain retry with 'retry: ' prefix"
                );
            }

            // Property 6: Verify response headers are correct
            let events: Vec<Result<SseEvent, std::convert::Infallible>> = vec![Ok(event)];
            let sse = Sse::new(stream::iter(events));
            let response = sse.into_response();

            prop_assert_eq!(
                response.headers().get(header::CONTENT_TYPE).map(|v| v.to_str().unwrap()),
                Some("text/event-stream"),
                "SSE response must have Content-Type: text/event-stream"
            );

            prop_assert_eq!(
                response.headers().get(header::CACHE_CONTROL).map(|v| v.to_str().unwrap()),
                Some("no-cache"),
                "SSE response must have Cache-Control: no-cache"
            );

            prop_assert_eq!(
                response.headers().get(header::CONNECTION).map(|v| v.to_str().unwrap()),
                Some("keep-alive"),
                "SSE response must have Connection: keep-alive"
            );
        }

        #[test]
        fn prop_sse_multiline_data_format(
            // Generate multiple lines of data
            lines in proptest::collection::vec("[a-zA-Z0-9 ]{1,30}", 1..5),
        ) {
            let data = lines.join("\n");
            let event = SseEvent::new(data.clone());
            let sse_string = event.to_sse_string();

            // Property: Each line of data must be prefixed with "data: "
            for line in lines.iter() {
                prop_assert!(
                    sse_string.contains(&format!("data: {}", line)),
                    "Each line of multiline data must be prefixed with 'data: '"
                );
            }

            // Property: Must end with double newline
            prop_assert!(
                sse_string.ends_with("\n\n"),
                "SSE event must end with double newline"
            );
        }
    }
}