eric / rustwarden

//! Helper to turn a [`NotificationBus`](super::NotificationBus) subscription
//! into an HTTP Server-Sent Events response.
//!
//! Wraps the usual boilerplate (JSON-serialising each event, attaching a
//! keep-alive, converting errors into client-friendly ends of stream) in a
//! single call so handlers stay trivial:
//!
//! ```no_run
//! use axum::response::sse::Sse;
//! use axum::response::IntoResponse;
//! use std::sync::Arc;
//! use rustwarden_core::notifications::{InMemoryBus, NotificationBus, sse::sse_response};
//!
//! #[derive(Clone, serde::Serialize)]
//! struct MyEvent { msg: String }
//!
//! async fn stream_handler(
//!     bus: axum::extract::Extension<Arc<dyn NotificationBus<MyEvent>>>,
//! ) -> impl IntoResponse {
//!     let filter = Arc::new(|_: &MyEvent| true);
//!     sse_response(bus.subscribe(filter))
//! }
//! ```

use axum::response::sse::{Event, KeepAlive, Sse};
use futures::stream::{Stream, StreamExt};
use serde::Serialize;
use std::convert::Infallible;
use std::time::Duration;

/// Default keep-alive interval (30 s) — matches most reverse-proxy idle timeouts.
pub const DEFAULT_SSE_KEEPALIVE: Duration = Duration::from_secs(30);

/// Wrap a stream of serializable events into an axum [`Sse`] response with
/// JSON payloads and a 30-second keep-alive.
///
/// Serialization errors for individual events are logged and the event is
/// dropped from the output stream; the connection stays open for subsequent
/// events.
pub fn sse_response<T, S>(stream: S) -> Sse<impl Stream<Item = Result<Event, Infallible>>>
where
    T: Serialize + Send + 'static,
    S: Stream<Item = T> + Send + 'static,
{
    let events = stream.filter_map(|event| async move {
        match Event::default().json_data(&event) {
            Ok(e) => Some(Ok::<_, Infallible>(e)),
            Err(err) => {
                tracing::warn!("Dropping SSE event that failed to serialize: {}", err);
                None
            }
        }
    });
    Sse::new(events).keep_alive(KeepAlive::new().interval(DEFAULT_SSE_KEEPALIVE))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::notifications::{EventFilter, InMemoryBus, NotificationBus};
    use axum::body::to_bytes;
    use axum::response::IntoResponse;
    use serde::{Deserialize, Serialize};
    use std::sync::Arc;

    #[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
    struct Event {
        tag: String,
        count: u32,
    }

    fn pass_all<T>() -> EventFilter<T> {
        Arc::new(|_| true)
    }

    #[tokio::test]
    async fn sse_response_emits_serialized_events_then_keepalive_stops_stream() {
        let bus: InMemoryBus<Event> = InMemoryBus::new(4);
        let stream = bus.subscribe(pass_all());

        bus.publish(Event {
            tag: "a".into(),
            count: 1,
        })
        .await;
        bus.publish(Event {
            tag: "b".into(),
            count: 2,
        })
        .await;

        // Close the bus so the subscription stream terminates naturally.
        drop(bus);

        let response = sse_response(stream).into_response();
        let status = response.status();
        assert!(status.is_success(), "expected 2xx, got {status}");

        // Read the whole body — bounded because the stream terminates once
        // the bus is dropped.
        let body = response.into_body();
        let bytes = to_bytes(body, 64 * 1024).await.unwrap();
        let text = std::str::from_utf8(&bytes).unwrap();

        assert!(
            text.contains(r#""tag":"a""#),
            "body must contain first event, got: {text}"
        );
        assert!(
            text.contains(r#""tag":"b""#),
            "body must contain second event, got: {text}"
        );
    }
}