rig_http/http_client/
framing.rs

1//! Incremental SSE and newline-delimited framing of response bytes.
2//! SSE dispatch requires a blank line; incomplete events are not flushed at EOF.
3//! NDJSON permits a final unterminated line through [`NdjsonFramer::finish`].
4//!
5//! ```
6//! use rig_http::http_client::framing::NdjsonFramer;
7//!
8//! let mut framer = NdjsonFramer::new();
9//! assert_eq!(framer.push(b"{}\n").collect::<Vec<_>>(), vec![b"{}".to_vec()]);
10//! ```
11
12use std::borrow::Cow;
13
14/// How a reply's bytes are split into wire frames.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum Framing {
17    /// `text/event-stream`: WHATWG event-stream grammar ([`SseFramer`]).
18    Sse,
19    /// Newline-delimited JSON ([`NdjsonFramer`]).
20    Ndjson,
21    /// The whole body is one frame.
22    Whole,
23}
24
25impl Framing {
26    /// The incremental framer of a streamed reply. A whole body is one
27    /// frame, so it has none.
28    pub fn framer(self) -> Option<Framer> {
29        let lines = match self {
30            Self::Sse => Lines::Sse(SseFramer::new()),
31            Self::Ndjson => Lines::Ndjson(NdjsonFramer::new()),
32            Self::Whole => return None,
33        };
34        Some(Framer { lines })
35    }
36
37    /// Split a whole reply body into the frames a decoder reads: one per SSE
38    /// event with a non-blank `data` field, one per NDJSON line, or the
39    /// whole body.
40    ///
41    /// ```
42    /// use rig_http::http_client::framing::{Framing, WireFrame};
43    ///
44    /// let frames = Framing::Sse.split(b"data: {}\n\ndata: [DONE]\n\n");
45    /// assert_eq!(frames.len(), 2);
46    /// assert_eq!(frames[1].as_str(), "[DONE]");
47    /// ```
48    pub fn split(&self, bytes: &[u8]) -> Vec<WireFrame> {
49        let Some(mut framer) = self.framer() else {
50            return (!bytes.is_empty())
51                .then(|| WireFrame::from_bytes(bytes.to_vec()))
52                .into_iter()
53                .collect();
54        };
55        let mut frames: Vec<WireFrame> = framer.push(bytes).collect();
56        frames.extend(framer.finish());
57        frames
58    }
59}
60
61/// One transport frame, after framing but before decoding.
62///
63/// The transport layer (SSE framer, NDJSON splitter, websocket reader) owns
64/// byte splitting and yields these; a decoder never splits bytes.
65#[derive(Debug, Clone, PartialEq, Eq)]
66pub enum WireFrame {
67    /// Text payload, such as an SSE `data:` field or WebSocket message.
68    Text(String),
69    /// Byte payload, such as an NDJSON line or binary frame.
70    Bytes(Vec<u8>),
71}
72
73impl WireFrame {
74    /// The frame payload as text (lossy for byte frames).
75    pub fn as_str(&self) -> Cow<'_, str> {
76        match self {
77            Self::Text(text) => Cow::Borrowed(text),
78            Self::Bytes(bytes) => String::from_utf8_lossy(bytes),
79        }
80    }
81
82    /// The frame payload's bytes.
83    pub fn as_bytes(&self) -> &[u8] {
84        match self {
85            Self::Text(text) => text.as_bytes(),
86            Self::Bytes(bytes) => bytes,
87        }
88    }
89
90    /// A frame of `bytes`: text when they are UTF-8.
91    pub fn from_bytes(bytes: Vec<u8>) -> Self {
92        match String::from_utf8(bytes) {
93            Ok(text) => Self::Text(text),
94            Err(error) => Self::Bytes(error.into_bytes()),
95        }
96    }
97}
98
99/// The incremental framer of one streamed reply, from [`Framing::framer`].
100#[derive(Debug)]
101pub struct Framer {
102    lines: Lines,
103}
104
105#[derive(Debug)]
106enum Lines {
107    Sse(SseFramer),
108    Ndjson(NdjsonFramer),
109}
110
111impl Framer {
112    /// Feed one transport chunk; yields every frame it completed. An SSE
113    /// event whose data is only whitespace is a heartbeat, not a frame.
114    pub fn push(&mut self, chunk: &[u8]) -> impl Iterator<Item = WireFrame> + '_ {
115        let (sse, ndjson) = match &mut self.lines {
116            Lines::Sse(framer) => (Some(framer.push(chunk)), None),
117            Lines::Ndjson(framer) => (None, Some(framer.push(chunk))),
118        };
119        let events = sse.into_iter().flatten();
120        let events = events.filter(|event| !event.data.trim().is_empty());
121        let lines = ndjson.into_iter().flatten().map(WireFrame::from_bytes);
122        events.map(|event| WireFrame::Text(event.data)).chain(lines)
123    }
124
125    /// The stream ended: NDJSON's unterminated last line. SSE dispatches
126    /// only on a blank line, so an unterminated event is never a frame.
127    pub fn finish(&mut self) -> Option<WireFrame> {
128        match &mut self.lines {
129            Lines::Sse(_) => None,
130            Lines::Ndjson(framer) => framer.finish().map(WireFrame::from_bytes),
131        }
132    }
133
134    /// Bytes pushed but not yet part of a frame, for truncation
135    /// diagnostics: [`SseFramer::pending`] or [`NdjsonFramer::pending`].
136    pub fn pending(&self) -> usize {
137        match &self.lines {
138            Lines::Sse(framer) => framer.pending(),
139            Lines::Ndjson(framer) => framer.pending(),
140        }
141    }
142}
143
144/// One dispatched `text/event-stream` event.
145#[derive(Debug, Clone, Default, PartialEq, Eq)]
146pub struct SseEvent {
147    /// The `event:` field, `"message"` when the stream named none.
148    pub event: Cow<'static, str>,
149    /// The joined `data:` fields, without the terminating newline.
150    pub data: String,
151    /// The stream's last event id, empty when none was ever set.
152    pub id: String,
153    /// Most recent parsed `retry:` value since the previous emitted event,
154    /// saturating at `u64::MAX`. This value does not trigger automatic
155    /// reconnection.
156    pub retry: Option<u64>,
157}
158
159/// A field recognized by the SSE grammar.
160#[derive(Debug, Clone, Copy, PartialEq, Eq)]
161pub enum SseField {
162    /// A payload line.
163    Data,
164    /// An event type.
165    Event,
166    /// An accepted event id.
167    Id,
168    /// A blank dispatch line.
169    Blank,
170    /// A comment, retry, ignored id or unknown field.
171    Other,
172}
173
174/// Original byte ranges of one line consumed by [`SseFramer`].
175#[derive(Debug, Clone, PartialEq, Eq)]
176pub struct SseLine {
177    /// Whole line, including its line ending, at saturating stream offsets.
178    pub span: std::ops::Range<usize>,
179    /// Line content without its ending or an initial BOM.
180    pub content: std::ops::Range<usize>,
181    /// Field value after a colon and optional space, when present.
182    pub value: Option<std::ops::Range<usize>>,
183    /// The field accepted by the grammar.
184    pub field: SseField,
185}
186
187/// A parsed event and the byte offset where it became dispatchable.
188#[derive(Debug, Clone, PartialEq, Eq)]
189pub struct RecordedSseEvent {
190    /// The grammar's parsed event.
191    pub event: SseEvent,
192    /// Saturating stream offset of its completing line ending or retained EOF.
193    pub end: usize,
194}
195
196/// Parsed events and source lines consumed by a recording operation.
197#[derive(Debug, Default)]
198pub struct SseBatch {
199    /// Events emitted in order, with their completion offsets.
200    pub events: Vec<RecordedSseEvent>,
201    /// Source lines in order, including non-dispatching lines.
202    pub lines: Vec<SseLine>,
203}
204
205/// The optional leading byte-order mark, which the grammar ignores.
206const BOM: [u8; 3] = [0xef, 0xbb, 0xbf];
207
208/// A push parser for the WHATWG `text/event-stream` grammar.
209#[derive(Debug, Default)]
210pub struct SseFramer {
211    /// Bytes not yet consumed as a complete line.
212    buffer: Vec<u8>,
213    /// How many bytes at the start of `buffer` hold no line end.
214    scanned: usize,
215    /// The last line ended at a CR that was the last byte pushed, so an LF
216    /// opening the next chunk completes that CRLF.
217    after_cr: bool,
218    /// Events completed by the current `push`.
219    ready: Vec<SseEvent>,
220    event_type: String,
221    data: String,
222    last_event_id: String,
223    retry: Option<u64>,
224    /// How many BOM bytes have matched so far, while the prefix is undecided.
225    bom_prefix: usize,
226    bom_done: bool,
227    /// Bytes of complete lines consumed since the last blank line.
228    since_dispatch: usize,
229    /// Absolute byte position after the last consumed line or BOM.
230    offset: usize,
231}
232
233impl SseFramer {
234    /// A framer at the start of a stream.
235    pub fn new() -> Self {
236        Self::default()
237    }
238
239    /// Feed one transport chunk; yields every event it completed.
240    pub fn push(&mut self, chunk: &[u8]) -> std::vec::Drain<'_, SseEvent> {
241        self.push_inner(chunk, None);
242        self.ready.drain(..)
243    }
244
245    /// Feed a chunk and retain parsed events with source ranges for recording.
246    pub fn push_recording(&mut self, chunk: &[u8]) -> SseBatch {
247        let mut batch = SseBatch::default();
248        self.push_inner(chunk, Some(&mut batch));
249        batch
250    }
251
252    /// Retain an incomplete final event for a recording, without changing
253    /// the normal SSE rule that EOF does not dispatch it.
254    pub fn finish_recording(&mut self) -> SseBatch {
255        let mut batch = SseBatch::default();
256        if self.bom_prefix > 0 {
257            self.buffer
258                .extend_from_slice(BOM.get(..self.bom_prefix).unwrap_or_default());
259            self.bom_prefix = 0;
260            self.bom_done = true;
261        }
262        if !self.buffer.is_empty() {
263            let buffer = std::mem::take(&mut self.buffer);
264            let start = self.offset;
265            self.offset = self.offset.saturating_add(buffer.len());
266            let (field, value) = self.line(&String::from_utf8_lossy(&buffer));
267            batch.lines.push(SseLine {
268                span: start..self.offset,
269                content: start..self.offset,
270                value: value.map(|at| start.saturating_add(at)..self.offset),
271                field,
272            });
273        }
274        self.dispatch();
275        if let Some(event) = self.ready.pop() {
276            batch.events.push(RecordedSseEvent {
277                event,
278                end: self.offset,
279            });
280        }
281        batch
282    }
283
284    fn push_inner(&mut self, chunk: &[u8], mut recording: Option<&mut SseBatch>) {
285        let mut chunk = self.strip_bom(chunk);
286        if self.after_cr
287            && let Some((&first, rest)) = chunk.split_first()
288        {
289            self.after_cr = false;
290            if first == b'\n' {
291                chunk = rest;
292                self.offset = self.offset.saturating_add(1);
293                // Pending bytes are counted as if the CRLF had arrived whole.
294                if self.since_dispatch > 0 {
295                    self.since_dispatch = self.since_dispatch.saturating_add(1);
296                }
297            }
298        }
299        self.buffer.extend_from_slice(chunk);
300        // Taken out, so its lines are read in place while `line` changes the rest.
301        let buffer = std::mem::take(&mut self.buffer);
302        let mut rest = buffer.as_slice();
303        while let Some((line, next)) = terminated_line(rest, self.scanned) {
304            self.scanned = 0;
305            let start = self.offset;
306            self.offset = self.offset.saturating_add(rest.len() - next.len());
307            self.since_dispatch = self.since_dispatch.saturating_add(rest.len() - next.len());
308            self.after_cr = next.is_empty() && rest.last() == Some(&b'\r');
309            let (field, value) = self.line(&String::from_utf8_lossy(line));
310            if let Some(batch) = recording.as_deref_mut() {
311                batch.lines.push(SseLine {
312                    span: start..self.offset,
313                    content: start..start.saturating_add(line.len()),
314                    value: value
315                        .map(|at| start.saturating_add(at)..start.saturating_add(line.len())),
316                    field,
317                });
318                if let Some(event) = self.ready.pop() {
319                    batch.events.push(RecordedSseEvent {
320                        event,
321                        end: self.offset,
322                    });
323                }
324            }
325            rest = next;
326        }
327        self.scanned = rest.len();
328        let consumed = buffer.len() - rest.len();
329        self.buffer = buffer;
330        self.buffer.drain(..consumed);
331    }
332
333    /// Saturating count of buffered bytes and complete lines since the last
334    /// blank line. Excludes any undecided leading BOM prefix; exposes no frame data.
335    pub fn pending(&self) -> usize {
336        self.since_dispatch.saturating_add(self.buffer.len())
337    }
338
339    /// Consume the stream's optional leading BOM, which may itself be split
340    /// across chunks. Returns the chunk with any BOM bytes removed.
341    fn strip_bom<'a>(&mut self, chunk: &'a [u8]) -> &'a [u8] {
342        if self.bom_done {
343            return chunk;
344        }
345        let mut rest = chunk;
346        while let Some(&byte) = rest.first() {
347            if BOM.get(self.bom_prefix) != Some(&byte) {
348                // Not a BOM after all: replay the bytes that matched so far.
349                let matched = std::mem::take(&mut self.bom_prefix);
350                self.bom_done = true;
351                self.buffer
352                    .extend_from_slice(BOM.get(..matched).unwrap_or_default());
353                return rest;
354            }
355            self.bom_prefix += 1;
356            rest = rest.get(1..).unwrap_or_default();
357            if self.bom_prefix == BOM.len() {
358                self.offset = self.offset.saturating_add(BOM.len());
359                self.bom_prefix = 0;
360                self.bom_done = true;
361                return rest;
362            }
363        }
364        rest
365    }
366
367    /// Process one complete line of the grammar.
368    fn line(&mut self, line: &str) -> (SseField, Option<usize>) {
369        if line.is_empty() {
370            self.dispatch();
371            return (SseField::Blank, None);
372        }
373        // A line starting with a colon is a comment.
374        if line.starts_with(':') {
375            return (SseField::Other, None);
376        }
377        let (field, value, at) = match line.split_once(':') {
378            Some((field, value)) => {
379                let trimmed = value.strip_prefix(' ').unwrap_or(value);
380                (
381                    field,
382                    trimmed,
383                    Some(field.len() + 1 + value.len() - trimmed.len()),
384                )
385            }
386            None => (line, "", None),
387        };
388        let field = match field {
389            "event" => {
390                self.event_type.clear();
391                self.event_type.push_str(value);
392                SseField::Event
393            }
394            "data" => {
395                self.data.push_str(value);
396                self.data.push('\n');
397                SseField::Data
398            }
399            // A NUL in an id is ignored outright, per the grammar.
400            "id" if !value.contains('\0') => {
401                self.last_event_id.clear();
402                self.last_event_id.push_str(value);
403                SseField::Id
404            }
405            "retry" if !value.is_empty() && value.bytes().all(|b| b.is_ascii_digit()) => {
406                // Digits only is a valid field, so a value past `u64` saturates.
407                self.retry = Some(value.parse().unwrap_or(u64::MAX));
408                SseField::Other
409            }
410            _ => SseField::Other,
411        };
412        (field, at)
413    }
414
415    /// A blank line ends an event. An event with no data is not dispatched:
416    /// it only resets the event type, as the grammar says.
417    fn dispatch(&mut self) {
418        self.since_dispatch = 0;
419        if self.data.is_empty() {
420            self.event_type.clear();
421            return;
422        }
423        self.data.pop();
424        let event = if self.event_type.is_empty() {
425            Cow::Borrowed("message")
426        } else {
427            Cow::Owned(std::mem::take(&mut self.event_type))
428        };
429        self.ready.push(SseEvent {
430            event,
431            data: std::mem::take(&mut self.data),
432            id: self.last_event_id.clone(),
433            retry: self.retry.take(),
434        });
435        self.event_type.clear();
436    }
437}
438
439/// Splits nonempty newline-delimited payloads without validating JSON.
440/// A nonempty unterminated last line is available through [`Self::finish`].
441#[derive(Debug, Default)]
442pub struct NdjsonFramer {
443    buffer: Vec<u8>,
444    /// How many bytes at the start of `buffer` hold no line end.
445    scanned: usize,
446    ready: Vec<Vec<u8>>,
447}
448
449impl NdjsonFramer {
450    /// A framer at the start of a stream.
451    pub fn new() -> Self {
452        Self::default()
453    }
454
455    /// Feed one transport chunk; yields every line it completed. Blank lines
456    /// are not frames.
457    pub fn push(&mut self, chunk: &[u8]) -> std::vec::Drain<'_, Vec<u8>> {
458        self.buffer.extend_from_slice(chunk);
459        let mut rest = self.buffer.as_slice();
460        while let Some((line, end)) = split_line(rest, self.scanned, |byte| *byte == b'\n') {
461            self.scanned = 0;
462            let line = line.strip_suffix(b"\r").unwrap_or(line);
463            if !line.is_empty() {
464                self.ready.push(line.to_vec());
465            }
466            rest = end.get(1..).unwrap_or_default();
467        }
468        self.scanned = rest.len();
469        let consumed = self.buffer.len() - rest.len();
470        self.buffer.drain(..consumed);
471        self.ready.drain(..)
472    }
473
474    /// The unterminated last line, if the stream ended with a nonempty one.
475    /// A CR ending it is part of the line end, as in [`Self::push`].
476    pub fn finish(&mut self) -> Option<Vec<u8>> {
477        let mut line = std::mem::take(&mut self.buffer);
478        if line.last() == Some(&b'\r') {
479            line.pop();
480        }
481        (!line.is_empty()).then_some(line)
482    }
483
484    /// Bytes of an unterminated trailing line, for truncation diagnostics.
485    pub fn pending(&self) -> usize {
486        self.buffer.len()
487    }
488}
489
490/// The first complete line in `buffer`, without its line end, and the
491/// bytes after that line end; the first `scanned` bytes hold no line end.
492///
493/// A CR ends a line at once, as the grammar says; an LF right after it in
494/// `buffer` belongs to the same line end.
495fn terminated_line(buffer: &[u8], scanned: usize) -> Option<(&[u8], &[u8])> {
496    let (line, end) = split_line(buffer, scanned, |byte| *byte == b'\n' || *byte == b'\r')?;
497    let rest = end.strip_prefix(b"\r\n").or_else(|| end.get(1..));
498    Some((line, rest.unwrap_or_default()))
499}
500
501/// `buffer` split before its first byte that `ends` a line, if it has one;
502/// the first `scanned` bytes hold none.
503fn split_line(buffer: &[u8], scanned: usize, ends: impl Fn(&u8) -> bool) -> Option<(&[u8], &[u8])> {
504    buffer.split_at_checked(scanned + buffer.iter().skip(scanned).position(ends)?)
505}
506
507#[cfg(test)]
508mod tests;
rig_http/http_client/framing.rs

rig_http/http_client/
framing.rs
