wsb 0.124.3689

use axum::response::sse::{Event, KeepAlive, Sse};
use futures::stream::Stream;
use std::convert::Infallible;
use tokio::sync::broadcast;

/// Create an SSE response from a broadcast receiver.
/// Each broadcast message is serialized as JSON in the "data" field
/// with the given event name.
pub fn sse_from_broadcast<T: serde::Serialize + Clone + Send + 'static>(
    event_name: &'static str,
    mut rx: broadcast::Receiver<T>,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let stream = async_stream::stream! {
        loop {
            match rx.recv().await {
                Ok(msg) => {
                    match serde_json::to_string(&msg) {
                        Ok(json) => {
                            yield Ok(Event::default().event(event_name).data(json));
                        }
                        Err(e) => {
                            yield Ok(Event::default().event("error").data(e.to_string()));
                        }
                    }
                }
                Err(broadcast::error::RecvError::Lagged(n)) => {
                    yield Ok(Event::default().event("warning").data(format!("lagged {} messages", n)));
                }
                Err(broadcast::error::RecvError::Closed) => {
                    break;
                }
            }
        }
    };

    Sse::new(stream).keep_alive(KeepAlive::default())
}

/// Create a simple SSE event with JSON data.
pub fn sse_event<T: serde::Serialize>(
    event_name: &str,
    data: &T,
) -> Result<Event, serde_json::Error> {
    let json = serde_json::to_string(data)?;
    Ok(Event::default().event(event_name).data(json))
}

/// Create an SSE event with a plain string.
pub fn sse_event_str(event_name: &str, data: &str) -> Event {
    Event::default().event(event_name).data(data.to_string())
}