Skip to main content

ag_ui/encode/
sse.rs

1//! Server-Sent Events framing.
2
3use crate::encode::{EventStreamFormatter, SSE_MEDIA_TYPE};
4use crate::error::Result;
5use crate::event::Event;
6
7/// Encodes events as `text/event-stream` frames.
8///
9/// Each event becomes a single `data:` block holding its JSON, exactly as the
10/// TypeScript SDK writes it:
11///
12/// ```text
13/// data: {"type":"TEXT_MESSAGE_END","messageId":"msg-1"}
14///
15/// ```
16#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
17pub struct SseFormatter;
18
19impl SseFormatter {
20    /// Builds a formatter. It holds no state.
21    pub const fn new() -> Self {
22        Self
23    }
24
25    /// Encodes one event as an SSE frame.
26    pub fn encode_to_string(&self, event: &Event) -> Result<String> {
27        Ok(frame(&serde_json::to_string(event)?))
28    }
29}
30
31impl EventStreamFormatter for SseFormatter {
32    fn content_type(&self) -> &'static str {
33        SSE_MEDIA_TYPE
34    }
35
36    fn encode(&self, event: &Event) -> Result<Vec<u8>> {
37        Ok(self.encode_to_string(event)?.into_bytes())
38    }
39}
40
41/// Wraps an arbitrary payload in an SSE `data:` block.
42///
43/// A payload containing line breaks becomes one `data:` line per line, which is
44/// what the SSE decoder rejoins with `\n`; a payload with none becomes a single
45/// line. Serialized JSON is always single-line — `serde_json` escapes control
46/// characters — so this only matters for callers that frame something else, but
47/// getting it wrong would silently truncate an event at the first newline.
48///
49/// ```
50/// # use ag_ui::encode::sse::frame;
51/// assert_eq!(frame("one\ntwo"), "data: one\ndata: two\n\n");
52/// ```
53pub fn frame(payload: &str) -> String {
54    let mut out = String::with_capacity(payload.len() + 8);
55    let mut rest = payload;
56
57    loop {
58        match rest.find(['\r', '\n']) {
59            Some(index) => {
60                out.push_str("data: ");
61                out.push_str(&rest[..index]);
62                out.push('\n');
63                // A CRLF is one break, not two.
64                let width = if rest[index..].starts_with("\r\n") {
65                    2
66                } else {
67                    1
68                };
69                rest = &rest[index + width..];
70            }
71            None => {
72                out.push_str("data: ");
73                out.push_str(rest);
74                out.push('\n');
75                break;
76            }
77        }
78    }
79
80    out.push('\n');
81    out
82}