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}