sz_rust_http_facade/
sse.rs

1//! Server-Sent Events (SSE) 支持
2//!
3//! 基于 axum 0.8 的 `axum::response::sse` 模块，提供轻量级服务器推送能力。
4//!
5//! ## 与 WebSocket 的区别
6//!
7//! | 维度 | SSE | WebSocket |
8//! |------|-----|-----------|
9//! | 方向 | 服务器 → 客户端（单向） | 双向 |
10//! | 协议 | HTTP | WS |
11//! | 重连 | 自动重连 | 需手动 |
12//! | 浏览器支持 | EventSource API | WebSocket API |
13//! | 适用场景 | 通知/日志流/进度 | 聊天/实时交互 |
14//!
15//! ## 用法
16//!
17//! ```ignore
18//! use sz_rust_http_facade::sse::{SseEvent, sse_response};
19//! use futures::stream::{self, StreamExt};
20//!
21//! async fn events_handler() -> impl IntoResponse {
22//!     let stream = stream::iter(vec![
23//!         SseEvent::data("hello").event("greeting"),
24//!         SseEvent::data("world").event("message"),
25//!     ])
26//!     .map(Ok);
27//!     sse_response(stream)
28//! }
29//! ```
30
31use axum::response::sse::{Event as AxumEvent, KeepAlive, Sse};
32use axum::response::IntoResponse;
33use core::convert::Infallible;
34use futures::stream::{Stream, StreamExt};
35
36/// SSE 事件构建器
37///
38/// 对 `axum::response::sse::Event` 的封装，提供更简洁的 API。
39#[derive(Debug, Clone)]
40pub struct SseEvent {
41    data: String,
42    event: Option<String>,
43    id: Option<String>,
44    retry: Option<u64>,
45}
46
47impl SseEvent {
48    /// 创建数据事件
49    pub fn data(data: impl Into<String>) -> Self {
50        Self {
51            data: data.into(),
52            event: None,
53            id: None,
54            retry: None,
55        }
56    }
57
58    /// 设置事件类型
59    pub fn event(mut self, event: impl Into<String>) -> Self {
60        self.event = Some(event.into());
61        self
62    }
63
64    /// 设置事件 ID
65    pub fn id(mut self, id: impl Into<String>) -> Self {
66        self.id = Some(id.into());
67        self
68    }
69
70    /// 设置重连等待时间（毫秒）
71    pub fn retry(mut self, retry_ms: u64) -> Self {
72        self.retry = Some(retry_ms);
73        self
74    }
75
76    /// 转换为 axum Event
77    pub fn into_axum_event(self) -> Result<AxumEvent, Infallible> {
78        let mut event = AxumEvent::default().data(&self.data);
79        if let Some(name) = self.event {
80            event = event.event(name);
81        }
82        if let Some(id) = self.id {
83            event = event.id(id);
84        }
85        if let Some(retry) = self.retry {
86            event = event.retry(std::time::Duration::from_millis(retry));
87        }
88        Ok(event)
89    }
90}
91
92/// 创建 SSE 响应，带 KeepAlive
93///
94/// 接受 `Stream<Item = Result<SseEvent, Infallible>>`，内部转换为 axum Event 流。
95pub fn sse_response<S>(stream: S) -> impl IntoResponse
96where
97    S: Stream<Item = Result<SseEvent, Infallible>> + Send + 'static,
98{
99    let axum_stream = stream.map(|item| item.and_then(|e| e.into_axum_event()));
100    Sse::new(axum_stream).keep_alive(KeepAlive::default())
101}
102
103/// 创建 SSE 响应，自定义 KeepAlive 间隔
104pub fn sse_response_with_interval<S>(stream: S, interval_secs: u64) -> impl IntoResponse
105where
106    S: Stream<Item = Result<SseEvent, Infallible>> + Send + 'static,
107{
108    let axum_stream = stream.map(|item| item.and_then(|e| e.into_axum_event()));
109    Sse::new(axum_stream).keep_alive(
110        KeepAlive::new()
111            .interval(std::time::Duration::from_secs(interval_secs))
112            .text("keep-alive"),
113    )
114}
115
116/// 从 Vec 创建有限 SSE 流（发送完所有事件后关闭）
117pub fn sse_from_events(events: Vec<SseEvent>) -> impl Stream<Item = Result<SseEvent, Infallible>> {
118    futures::stream::iter(events.into_iter().map(Ok))
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn test_sse_event_data() {
127        let event = SseEvent::data("hello");
128        assert_eq!(event.data, "hello");
129        assert!(event.event.is_none());
130        assert!(event.id.is_none());
131        assert!(event.retry.is_none());
132    }
133
134    #[test]
135    fn test_sse_event_builder() {
136        let event = SseEvent::data("payload")
137            .event("update")
138            .id("123")
139            .retry(5000);
140        assert_eq!(event.data, "payload");
141        assert_eq!(event.event.as_deref(), Some("update"));
142        assert_eq!(event.id.as_deref(), Some("123"));
143        assert_eq!(event.retry, Some(5000));
144    }
145
146    #[test]
147    fn test_sse_event_to_axum() {
148        let event = SseEvent::data("test").event("ping");
149        let axum_event = event.into_axum_event();
150        assert!(axum_event.is_ok());
151    }
152
153    #[test]
154    fn test_sse_from_events() {
155        let events = vec![SseEvent::data("first"), SseEvent::data("second")];
156        let stream = sse_from_events(events);
157        let collected: Vec<_> = futures::executor::block_on(stream.collect());
158        assert_eq!(collected.len(), 2);
159    }
160}
sz_rust_http_facade/sse.rs

sz_rust_http_facade/
sse.rs
