allama_http/sse.rs
1//! Server-sent events (SSE) for streaming adapters.
2//!
3//! Every streaming HTTP API in the bundled providers speaks SSE, so this is
4//! the framing most adapters need, first-party or not:
5//!
6//! - **Framing.** [`SseStream`] turns a response [`Body`](crate::Body) (or
7//! any stream of byte chunks) into [`SseEvent`]s; [`SseDecoder`] is the
8//! same decoder without the stream, for callers that feed bytes
9//! themselves. Both follow the [WHATWG event stream format]: `LF`, `CR`
10//! and `CRLF` line endings (also when a chunk boundary splits a `CRLF`),
11//! multi-line `data:` fields joined with `\n`, a leading byte order mark,
12//! comments (`: keepalive`) and `retry:` lines skipped, and an event
13//! without a trailing blank line discarded at the end of the body. Each
14//! byte is scanned once, however the chunks split the lines, and a size
15//! cap ([`SseDecoder::DEFAULT_MAX_EVENT_SIZE`] by default) bounds what a
16//! stream can make the decoder hold.
17//! - **Two deliberate failures.** Where the standard decodes a stream with
18//! invalid UTF-8 lossily, these helpers fail
19//! ([`SseError::InvalidUtf8`]), and a line or an event past the size cap
20//! fails too ([`SseError::EventTooLarge`]). Neither is retryable: asking
21//! again streams the same bytes. [`SseError::into_error`] maps them, and a
22//! failed read, to the stream errors the variants name.
23//! - **The driver.** [`SseChatStream`] runs an adapter's
24//! [`ChatFrameDecoder`](crate::framing::ChatFrameDecoder)`<`[`SseEvent`]`>`
25//! over the events: it is the framing-neutral
26//! [`FramedChatStream`](crate::framing::FramedChatStream) over an
27//! [`SseStream`] (a type alias), so it ends the stream after the first error, runs the
28//! decoder's end-of-body or failure hook once, classifies a failed read
29//! (a stream timeout stays a timeout), and never polls the body after
30//! the stream ended.
31//! - **Vendor frames.** [`VendorFrames`] (from [`framing`](crate::framing))
32//! surfaces named events an adapter does not model
33//! (`event: search_results`) as response-level provider data under the
34//! adapter's namespace, with one warning per event type, so they are
35//! never decoded as the adapter's own and never dropped.
36//!
37//! Other wire framings use the same driver and vendor frames:
38//! [`ndjson`](crate::ndjson) frames newline-delimited JSON, and any other
39//! stream of frames drives a
40//! [`ChatFrameDecoder`](crate::framing::ChatFrameDecoder).
41//!
42//! Decoding the events' data stays with each adapter: the helpers frame,
43//! drive and surface, they do not interpret.
44//!
45//! ```
46//! use allama::{ProviderData, StreamEvent};
47//! use allama_http::sse::{SseStream, VendorFrames};
48//! use futures_util::{FutureExt, StreamExt, stream};
49//!
50//! // A body whose second chunk starts in the middle of a `data:` line.
51//! let body = stream::iter([
52//! Ok::<_, std::io::Error>(&b": keepalive\r\nevent: delta\r\ndata: {\"text\":"[..]),
53//! Ok(&b"\"Hi\"}\r\n\r\nevent: search_results\r\ndata: [1]\r\n\r\n"[..]),
54//! ]);
55//! let events: Vec<_> = SseStream::new(body).collect().now_or_never().unwrap();
56//! let mut vendor_frames = VendorFrames::new("demo", "demo");
57//! let mut out = Vec::new();
58//! for event in events {
59//! let event = event.unwrap();
60//! match event.event.as_str() {
61//! "delta" => assert_eq!(event.data, r#"{"text":"Hi"}"#),
62//! name => out.extend(vendor_frames.surface(name, &event.data)),
63//! }
64//! }
65//! // A warning, then the data under the namespace and event type.
66//! assert!(matches!(out[0], StreamEvent::Warning { .. }));
67//! assert_eq!(
68//! out[1],
69//! StreamEvent::response_provider_data(ProviderData::new().with_raw(
70//! "demo",
71//! "search_results",
72//! serde_json::json!([[1]]),
73//! )),
74//! );
75//! ```
76//!
77//! [WHATWG event stream format]: https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream
78
79use std::convert::Infallible;
80use std::fmt;
81use std::pin::Pin;
82use std::str::Utf8Error;
83use std::task::{Context, Poll, ready};
84
85use futures_core::Stream;
86
87/// The SSE events of a byte stream, such as an [`HttpResponse`](crate::HttpResponse)
88/// body.
89///
90/// Yields each event once its terminating blank line arrives (see the
91/// [module docs](self) for the framing rules). After the first error,
92/// a failed read ([`SseError::Transport`]), a line that is not UTF-8
93/// ([`SseError::InvalidUtf8`]) or an event past the size cap
94/// ([`SseError::EventTooLarge`], see [`max_event_size`](Self::max_event_size)),
95/// the stream ends. It polls the body only when no complete event is
96/// buffered, so transport-level idle bounds
97/// ([`StreamTimeouts`](crate::StreamTimeouts)) see the same waits as
98/// before.
99///
100/// ```
101/// use allama_http::sse::SseStream;
102/// use futures_util::{FutureExt, StreamExt, stream};
103///
104/// let body = stream::iter([Ok::<_, std::io::Error>("data: {\"a\":\ndata: 1}\n\n")]);
105/// let events: Vec<_> = SseStream::new(body).collect().now_or_never().unwrap();
106/// let event = events[0].as_ref().unwrap();
107/// assert_eq!(event.event, "message");
108/// assert_eq!(event.data, "{\"a\":\n1}");
109/// ```
110#[derive(Debug)]
111pub struct SseStream<S> {
112 /// The byte stream; `None` once it ended or failed.
113 body: Option<S>,
114 decoder: SseDecoder,
115}
116
117impl<S> SseStream<S> {
118 /// Frame `body` into SSE events, with the default size cap
119 /// ([`SseDecoder::DEFAULT_MAX_EVENT_SIZE`]).
120 #[must_use]
121 pub fn new(body: S) -> Self {
122 Self {
123 body: Some(body),
124 decoder: SseDecoder::new(),
125 }
126 }
127
128 /// Set the size cap, in bytes: see [`SseDecoder::max_event_size`].
129 /// Provider crates set it from their builder's
130 /// `max_stream_event_size`.
131 #[must_use]
132 pub fn max_event_size(mut self, limit: usize) -> Self {
133 self.decoder = self.decoder.max_event_size(limit);
134 self
135 }
136
137 fn end(&mut self) {
138 self.body = None;
139 self.decoder = SseDecoder::new();
140 }
141}
142
143impl<S, B, E> Stream for SseStream<S>
144where
145 S: Stream<Item = Result<B, E>> + Unpin,
146 B: AsRef<[u8]>,
147{
148 type Item = Result<SseEvent, SseError<E>>;
149
150 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
151 let this = self.get_mut();
152 loop {
153 match this.decoder.next_event() {
154 Ok(Some(event)) => return Poll::Ready(Some(Ok(event))),
155 Ok(None) => {}
156 Err(error) => {
157 this.end();
158 return Poll::Ready(Some(Err(error.widen())));
159 }
160 }
161 let Some(body) = this.body.as_mut() else {
162 return Poll::Ready(None);
163 };
164 match ready!(Pin::new(body).poll_next(cx)) {
165 Some(Ok(chunk)) => this.decoder.push(chunk.as_ref()),
166 Some(Err(error)) => {
167 this.end();
168 return Poll::Ready(Some(Err(SseError::Transport(error))));
169 }
170 // What is left is an event without its blank line: discarded.
171 None => {
172 this.end();
173 return Poll::Ready(None);
174 }
175 }
176 }
177 }
178}
179
180/// An incremental SSE decoder: push bytes, take complete events.
181///
182/// [`SseStream`] drives one over a byte stream; use it directly when the
183/// bytes do not come as a [`Stream`] (a callback API, a test).
184///
185/// Decoding is linear: each byte is scanned for a line ending once, however
186/// the chunks split the lines. A size cap
187/// ([`max_event_size`](Self::max_event_size)) bounds the line still waiting
188/// for its ending, where memory grows from chunk to chunk, and the data an
189/// event accumulates.
190///
191/// ```
192/// use allama_http::sse::SseDecoder;
193///
194/// let mut decoder = SseDecoder::new();
195/// decoder.push(b"event: ping\r");
196/// assert_eq!(decoder.next_event(), Ok(None));
197/// decoder.push(b"\ndata: {}\r\n\r\n");
198/// let event = decoder.next_event().unwrap().unwrap();
199/// assert_eq!((event.event.as_str(), event.data.as_str()), ("ping", "{}"));
200/// assert_eq!(decoder.next_event(), Ok(None));
201/// ```
202#[derive(Debug)]
203pub struct SseDecoder {
204 /// Received bytes; those before `read` are consumed.
205 buffer: Vec<u8>,
206 read: usize,
207 /// The scanned offset: the bytes from `read` to here hold no line
208 /// ending. [`next_event`](Self::next_event) looks for one from here on.
209 scanned: usize,
210 /// The bytes from here on hold no line ending either: the line still
211 /// waiting for its ending. [`push`](Self::push) scans each chunk back
212 /// to its last line ending to move it, and `next_event` looks for line
213 /// endings only before it, so each byte is scanned once.
214 unterminated: usize,
215 /// Whether the start of the stream (and a byte order mark) is behind.
216 started: bool,
217 /// The last line ended with `CR`: a `LF` right after it belongs to it.
218 after_cr: bool,
219 max_event_size: usize,
220 /// The line waiting for its ending passed the size cap: the decoder
221 /// fails once the complete lines before it are taken.
222 overflowing: bool,
223 /// The size cap was passed: decoding cannot go on.
224 overflowed: bool,
225 event: PendingEvent,
226 /// Bytes examined for line endings, for the linearity test.
227 #[cfg(test)]
228 examined: usize,
229}
230
231impl Default for SseDecoder {
232 fn default() -> Self {
233 Self::new()
234 }
235}
236
237impl SseDecoder {
238 /// The default [`max_event_size`](Self::max_event_size): 64 MiB.
239 ///
240 /// Generous on purpose: a Responses `response.completed` event repeats
241 /// every output of the response, and streamed base64 images (OpenAI
242 /// partial images, Gemini inline data) can pass 16 MiB.
243 pub const DEFAULT_MAX_EVENT_SIZE: usize = 64 * 1024 * 1024;
244
245 /// A decoder at the start of a stream, with the default size cap.
246 #[must_use]
247 pub fn new() -> Self {
248 Self {
249 buffer: Vec::new(),
250 read: 0,
251 scanned: 0,
252 unterminated: 0,
253 started: false,
254 after_cr: false,
255 max_event_size: Self::DEFAULT_MAX_EVENT_SIZE,
256 overflowing: false,
257 overflowed: false,
258 event: PendingEvent::default(),
259 #[cfg(test)]
260 examined: 0,
261 }
262 }
263
264 /// Set the size cap, in bytes (default
265 /// [`DEFAULT_MAX_EVENT_SIZE`](Self::DEFAULT_MAX_EVENT_SIZE)).
266 ///
267 /// It bounds a line (its line ending excluded) and an event's data (its
268 /// `data:` values, joined). [`push`](Self::push) checks the line still
269 /// waiting for its ending before it buffers a chunk, so that line never
270 /// grows past the cap, not even briefly while the buffer reallocates.
271 ///
272 /// Past the cap, [`next_event`](Self::next_event) fails with
273 /// [`SseError::EventTooLarge`], and decoding cannot go on: the decoder
274 /// drops what it holds, ignores later bytes and keeps returning the
275 /// error.
276 ///
277 /// ```
278 /// use allama_http::sse::{SseDecoder, SseError};
279 ///
280 /// let mut decoder = SseDecoder::new().max_event_size(16);
281 /// decoder.push(b"data: 0123456789");
282 /// assert_eq!(decoder.next_event(), Ok(None));
283 /// // A 17th byte without a line ending.
284 /// decoder.push(b"a");
285 /// assert!(matches!(
286 /// decoder.next_event(),
287 /// Err(SseError::EventTooLarge { limit: 16, .. })
288 /// ));
289 /// ```
290 #[must_use]
291 pub fn max_event_size(mut self, limit: usize) -> Self {
292 self.max_event_size = limit;
293 self
294 }
295
296 /// Append received bytes. Chunk boundaries may fall anywhere, inside a
297 /// line, a `CRLF` or a UTF-8 sequence.
298 ///
299 /// Bytes that would take the line still waiting for its ending past the
300 /// [size cap](Self::max_event_size) are not buffered:
301 /// [`next_event`](Self::next_event) returns the events completed before
302 /// that line, then fails.
303 pub fn push(&mut self, bytes: &[u8]) {
304 if self.overflowing || self.overflowed {
305 return;
306 }
307 if self.read > 0 {
308 self.buffer.drain(..self.read);
309 self.scanned = self.scanned.saturating_sub(self.read);
310 self.unterminated = self.unterminated.saturating_sub(self.read);
311 self.read = 0;
312 }
313 // Back to the chunk's last line ending: what follows it starts, or
314 // continues, the line still waiting for its ending.
315 let last_ending = bytes.iter().rposition(|byte| is_line_ending(*byte));
316 #[cfg(test)]
317 {
318 self.examined += last_ending.map_or(bytes.len(), |end| bytes.len() - end);
319 }
320 let (unterminated, waiting) = match last_ending {
321 Some(end) => (self.buffer.len() + end + 1, bytes.len() - end - 1),
322 None => (
323 self.unterminated,
324 self.buffer.len() - self.unterminated + bytes.len(),
325 ),
326 };
327 let bytes = if waiting > self.max_event_size {
328 // Keep the complete lines; the rest is never buffered.
329 self.overflowing = true;
330 let Some(end) = last_ending else {
331 return;
332 };
333 &bytes[..=end]
334 } else {
335 bytes
336 };
337 let additional = room(
338 self.buffer.len(),
339 self.buffer.capacity(),
340 bytes.len(),
341 self.max_event_size,
342 );
343 self.buffer.reserve_exact(additional);
344 // A chunk without a line ending extends a line scanned to its end.
345 if last_ending.is_none() && self.scanned == self.buffer.len() {
346 self.scanned += bytes.len();
347 }
348 self.buffer.extend_from_slice(bytes);
349 self.unterminated = unterminated.min(self.buffer.len());
350 }
351
352 /// The next complete event, or `None` until more bytes arrive.
353 ///
354 /// # Errors
355 ///
356 /// - [`SseError::InvalidUtf8`] when a line is not valid UTF-8. The line
357 /// is skipped, and decoding can go on, but the stream is corrupt:
358 /// [`SseStream`] ends there.
359 /// - [`SseError::EventTooLarge`] past the
360 /// [size cap](Self::max_event_size). Decoding cannot go on.
361 ///
362 /// A decoder fails with no other error: it reads no transport.
363 pub fn next_event(&mut self) -> Result<Option<SseEvent>, SseError<Infallible>> {
364 if self.overflowed {
365 return Err(self.too_large());
366 }
367 if !self.started {
368 let pending = &self.buffer[self.read..];
369 if !self.overflowing && pending.len() < BOM.len() && BOM.starts_with(pending) {
370 return Ok(None);
371 }
372 if pending.starts_with(BOM) {
373 self.read += BOM.len();
374 self.scanned = self.scanned.max(self.read);
375 }
376 self.started = true;
377 }
378 loop {
379 let pending = &self.buffer[self.read..];
380 if self.after_cr && !pending.is_empty() {
381 self.after_cr = false;
382 if pending[0] == b'\n' {
383 self.read += 1;
384 self.scanned = self.scanned.max(self.read);
385 continue;
386 }
387 }
388 // Line endings lie only between the two scanned offsets.
389 let from = self.scanned.max(self.read);
390 let scan = &self.buffer[from..self.unterminated.max(from)];
391 let found = scan.iter().position(|byte| is_line_ending(*byte));
392 #[cfg(test)]
393 {
394 self.examined += found.map_or(scan.len(), |offset| offset + 1);
395 }
396 let Some(offset) = found else {
397 self.scanned = self.buffer.len();
398 if self.overflowing {
399 self.overflow();
400 return Err(self.too_large());
401 }
402 return Ok(None);
403 };
404 let (start, end) = (self.read, from + offset);
405 if end - start > self.max_event_size {
406 self.overflow();
407 return Err(self.too_large());
408 }
409 self.read = end + 1;
410 self.scanned = self.read;
411 self.after_cr = self.buffer[end] == b'\r';
412 let line =
413 std::str::from_utf8(&self.buffer[start..end]).map_err(SseError::InvalidUtf8)?;
414 match self.event.line(line, self.max_event_size) {
415 Ok(Some(event)) => return Ok(Some(event)),
416 Ok(None) => {}
417 Err(TooLarge) => {
418 self.overflow();
419 return Err(self.too_large());
420 }
421 }
422 }
423 }
424
425 /// Fail for good: drop what is held.
426 fn overflow(&mut self) {
427 self.overflowed = true;
428 self.buffer = Vec::new();
429 self.read = 0;
430 self.scanned = 0;
431 self.unterminated = 0;
432 self.event = PendingEvent::default();
433 }
434
435 fn too_large(&self) -> SseError<Infallible> {
436 SseError::EventTooLarge {
437 limit: self.max_event_size,
438 }
439 }
440}
441
442/// One SSE event.
443#[derive(Debug, Clone, PartialEq, Eq)]
444#[non_exhaustive]
445pub struct SseEvent {
446 /// The event type: the `event:` field, or `"message"` when the event
447 /// has none (as in the SSE standard).
448 pub event: String,
449 /// The `data:` fields, joined with `\n`. An event with no `data:`
450 /// field is never dispatched; one with an empty `data:` field is, with
451 /// empty data.
452 pub data: String,
453 /// The last event id the stream set (`id:`), if any. It carries over to
454 /// later events until another `id:` field changes it.
455 pub id: Option<String>,
456}
457
458impl SseEvent {
459 /// An event with the given type and data, and no event id.
460 ///
461 /// Decoders get their events from [`SseStream`] or [`SseDecoder`]; this
462 /// builds one directly, as when testing an adapter's event mapping. Set
463 /// [`id`](Self::id) afterwards for an event that carries one.
464 ///
465 /// ```
466 /// use allama_http::sse::SseEvent;
467 ///
468 /// let mut event = SseEvent::new("message_delta", r#"{"type":"message_delta"}"#);
469 /// assert_eq!(event.event, "message_delta");
470 /// assert_eq!(event.id, None);
471 ///
472 /// event.id = Some("42".into());
473 /// ```
474 #[must_use]
475 pub fn new(event: impl Into<String>, data: impl Into<String>) -> Self {
476 Self {
477 event: event.into(),
478 data: data.into(),
479 id: None,
480 }
481 }
482}
483
484/// Why an [`SseStream`] (or an [`SseDecoder`], with `E` =
485/// [`Infallible`]) failed.
486#[derive(Debug, PartialEq, Eq)]
487#[non_exhaustive]
488pub enum SseError<E> {
489 /// Reading the body failed: the body stream's own error. Map it as you
490 /// would any mid-stream transport failure (a
491 /// [`TransportError`](crate::TransportError) keeps its kind, so a
492 /// stream timeout stays a timeout).
493 Transport(E),
494 /// A line was not valid UTF-8.
495 ///
496 /// The [WHATWG event stream format] decodes the stream as UTF-8 and
497 /// replaces invalid bytes with U+FFFD. This decoder fails instead, on
498 /// purpose. A browser shows event text to a person; an adapter decodes
499 /// JSON from it, and a replaced byte changes a string value (generated
500 /// text, tool arguments, an opaque signature that is replayed later)
501 /// without anything saying so, which allama never does silently (ADR
502 /// 0003). Invalid UTF-8 means a broken server or intermediary, and
503 /// asking again streams the same bytes: map it to a stream decode
504 /// error ([`allama::Error::stream_decode`]), which is not retryable,
505 /// not to a transport failure. A chunk boundary inside a UTF-8
506 /// sequence is not an error: each line is decoded once it is complete.
507 ///
508 /// [WHATWG event stream format]: https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation
509 InvalidUtf8(Utf8Error),
510 /// A line, or an event's data, passed the decoder's size cap
511 /// ([`SseDecoder::max_event_size`]). The stream is not one the caller
512 /// is prepared to hold, and asking again gets the same answer: map it
513 /// to a stream protocol error ([`allama::Error::stream_protocol`]),
514 /// not a retryable transport failure.
515 #[non_exhaustive]
516 EventTooLarge {
517 /// The cap, in bytes.
518 limit: usize,
519 },
520}
521
522impl<E: Into<crate::TransportError>> SseError<E> {
523 /// The portable error for a failed SSE stream, as the variants say:
524 ///
525 /// - [`Transport`](Self::Transport): the body error as a
526 /// [`TransportError`](crate::TransportError), mapped by
527 /// [`TransportError::into_stream_error`](crate::TransportError::into_stream_error),
528 /// so a stream timeout stays [`allama::Error::Timeout`] and any other
529 /// failed read is a retryable transport stream error, with the
530 /// transport's text as the upstream message;
531 /// - [`InvalidUtf8`](Self::InvalidUtf8): a stream decode error;
532 /// - [`EventTooLarge`](Self::EventTooLarge): a stream protocol error.
533 ///
534 /// Neither of the last two is retryable, and every message is
535 /// [authored](allama::Error::authored). Add the provider and the request
536 /// id. [`SseChatStream::new`] maps failures with it.
537 ///
538 /// ```
539 /// use allama::{Error, StreamErrorKind};
540 /// use allama_http::sse::{SseError, SseStream};
541 /// use allama_http::{TransportError, TransportErrorKind};
542 /// use futures_util::{FutureExt, StreamExt, stream};
543 ///
544 /// let body = stream::iter([Ok::<_, TransportError>(&b"data: \xFF\n\n"[..])]);
545 /// let failure = SseStream::new(body).next().now_or_never().flatten().unwrap().unwrap_err();
546 /// let err = failure.into_error().with_provider("acme");
547 /// assert!(matches!(err, Error::Stream { kind: StreamErrorKind::Decode, .. }));
548 /// assert!(!err.is_retryable());
549 ///
550 /// let timeout = SseError::Transport(TransportError::new(TransportErrorKind::Timeout, "idle"));
551 /// assert!(matches!(timeout.into_error(), Error::Timeout { .. }));
552 /// ```
553 #[must_use]
554 pub fn into_error(self) -> allama::Error {
555 match self {
556 Self::Transport(error) => error.into().into_stream_error(),
557 Self::InvalidUtf8(error) => {
558 allama::Error::stream_decode(SseError::<Infallible>::InvalidUtf8(error).to_string())
559 .authored()
560 }
561 Self::EventTooLarge { limit } => allama::Error::stream_protocol(
562 SseError::<Infallible>::EventTooLarge { limit }.to_string(),
563 )
564 .authored(),
565 }
566 }
567}
568
569impl SseError<Infallible> {
570 /// A decoder's error as a stream's.
571 fn widen<E>(self) -> SseError<E> {
572 match self {
573 Self::Transport(never) => match never {},
574 Self::InvalidUtf8(error) => SseError::InvalidUtf8(error),
575 Self::EventTooLarge { limit } => SseError::EventTooLarge { limit },
576 }
577 }
578}
579
580impl<E: fmt::Display> fmt::Display for SseError<E> {
581 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
582 match self {
583 Self::Transport(error) => error.fmt(f),
584 Self::InvalidUtf8(error) => write!(f, "invalid UTF-8 in the SSE stream: {error}"),
585 Self::EventTooLarge { limit } => {
586 write!(f, "an SSE event passed the size cap of {limit} bytes")
587 }
588 }
589 }
590}
591
592impl<E: std::error::Error + 'static> std::error::Error for SseError<E> {
593 /// A transport failure is transparent: its source is the body error's
594 /// source.
595 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
596 match self {
597 Self::Transport(error) => error.source(),
598 Self::InvalidUtf8(error) => Some(error),
599 Self::EventTooLarge { .. } => None,
600 }
601 }
602}
603
604impl<E: Into<crate::TransportError>> From<SseError<E>> for allama::Error {
605 /// The error [`SseError::into_error`] gives: what
606 /// [`FramedChatStream::new`](crate::framing::FramedChatStream::new)
607 /// classifies a failure with.
608 fn from(error: SseError<E>) -> Self {
609 error.into_error()
610 }
611}
612
613/// The UTF-8 byte order mark, skipped at the start of a stream.
614const BOM: &[u8] = b"\xEF\xBB\xBF";
615
616fn is_line_ending(byte: u8) -> bool {
617 matches!(byte, b'\n' | b'\r')
618}
619
620/// How much to reserve so a buffer of `len` bytes and `capacity` takes
621/// `additional` more: doubling, as `Vec` does, but never past `limit` for
622/// the doubling alone.
623fn room(len: usize, capacity: usize, additional: usize, limit: usize) -> usize {
624 let needed = len + additional;
625 if needed <= capacity {
626 return 0;
627 }
628 capacity.saturating_mul(2).min(limit).max(needed) - len
629}
630
631/// An event's data passed the size cap.
632struct TooLarge;
633
634/// The fields of the event being received.
635#[derive(Debug, Default)]
636struct PendingEvent {
637 event: String,
638 /// `data:` values, each followed by `\n`: empty until a `data:` field.
639 data: String,
640 /// The last event id: it outlives the event.
641 last_id: String,
642}
643
644impl PendingEvent {
645 /// Process one line; a blank line dispatches the event. The event's
646 /// data may hold up to `max_data` bytes.
647 fn line(&mut self, line: &str, max_data: usize) -> Result<Option<SseEvent>, TooLarge> {
648 if line.is_empty() {
649 return Ok(self.dispatch());
650 }
651 if line.starts_with(':') {
652 return Ok(None);
653 }
654 let (field, value) = match line.split_once(':') {
655 Some((field, value)) => (field, value.strip_prefix(' ').unwrap_or(value)),
656 None => (line, ""),
657 };
658 match field {
659 "event" => value.clone_into(&mut self.event),
660 "data" => {
661 // The joined data: the values so far, each with its `\n`,
662 // then this one.
663 if self.data.len() + value.len() > max_data {
664 return Err(TooLarge);
665 }
666 let additional = room(
667 self.data.len(),
668 self.data.capacity(),
669 value.len() + 1,
670 max_data.saturating_add(1),
671 );
672 self.data.reserve_exact(additional);
673 self.data.push_str(value);
674 self.data.push('\n');
675 }
676 "id" if !value.contains('\0') => value.clone_into(&mut self.last_id),
677 // `retry:` sets a reconnection delay; a streamed model response
678 // is never reconnected. Unknown fields are ignored, as the
679 // standard requires.
680 _ => {}
681 }
682 Ok(None)
683 }
684
685 fn dispatch(&mut self) -> Option<SseEvent> {
686 let event = std::mem::take(&mut self.event);
687 let mut data = std::mem::take(&mut self.data);
688 // No `data:` field, no event; otherwise drop the last `\n`.
689 data.pop()?;
690 Some(SseEvent {
691 event: if event.is_empty() {
692 "message".to_owned()
693 } else {
694 event
695 },
696 data,
697 id: (!self.last_id.is_empty()).then(|| self.last_id.clone()),
698 })
699 }
700}
701
702mod chat;
703
704pub use crate::framing::VendorFrames;
705pub use chat::SseChatStream;
706
707#[cfg(test)]
708mod tests;