allama_http/ndjson.rs
1//! Newline-delimited JSON (NDJSON) for streaming adapters.
2//!
3//! Some APIs stream one JSON value per line instead of SSE (Ollama's native
4//! API, several self-hosted servers). [`NdjsonStream`] frames a response
5//! [`Body`](crate::Body) (or any stream of byte chunks) into its lines, and
6//! [`NdjsonDecoder`] is the same framer without the stream, for callers that
7//! feed bytes themselves. The lines drive an adapter's
8//! [`ChatFrameDecoder<String>`](crate::framing::ChatFrameDecoder) through
9//! the shared [`FramedChatStream`](crate::framing::FramedChatStream), the
10//! same driver SSE adapters use.
11//!
12//! The framing rules:
13//!
14//! - A line ends with `LF`; a `CR` right before it belongs to the line
15//! ending (`CRLF`), also when a chunk boundary splits the two.
16//! - Blank lines (empty, or only JSON whitespace) are skipped: they carry
17//! no value, and servers send them as keep-alives.
18//! - A leading UTF-8 byte order mark is skipped.
19//! - A last line without its `LF` is still a line when the body ends, as
20//! NDJSON allows. A body cut off in the middle of a line therefore
21//! yields a partial line, which the adapter's JSON decoding rejects.
22//! - Each byte is scanned once, however the chunks split the lines, and a
23//! size cap ([`NdjsonDecoder::DEFAULT_MAX_LINE_SIZE`] by default) bounds
24//! what a stream can make the framer hold.
25//!
26//! Two failures end the stream, as in [`sse`](crate::sse): a line that is
27//! not UTF-8 ([`NdjsonError::InvalidUtf8`]: JSON text is UTF-8, and a
28//! replaced byte would change a value silently) and a line past the size cap
29//! ([`NdjsonError::LineTooLarge`]). Neither is retryable: asking again
30//! streams the same bytes. [`NdjsonError::into_error`] maps them, and a
31//! failed read, to the stream errors the variants name.
32//!
33//! Decoding each line's JSON stays with the adapter: the framer splits, it
34//! does not interpret.
35//!
36//! ```
37//! use allama_http::ndjson::NdjsonStream;
38//! use futures_util::{FutureExt, StreamExt, stream};
39//!
40//! // A `CRLF` split by a chunk boundary, a blank keep-alive line, and a last
41//! // line without its line ending.
42//! let body = stream::iter([
43//! Ok::<_, std::io::Error>(&b"{\"text\":\"Hel\"}\r"[..]),
44//! Ok(&b"\n\n{\"text\":\"lo\"}\n{\"done\":true}"[..]),
45//! ]);
46//! let lines: Vec<_> = NdjsonStream::new(body).collect().now_or_never().unwrap();
47//! let lines: Vec<_> = lines.into_iter().map(Result::unwrap).collect();
48//! assert_eq!(lines, [r#"{"text":"Hel"}"#, r#"{"text":"lo"}"#, r#"{"done":true}"#]);
49//! ```
50
51use std::convert::Infallible;
52use std::fmt;
53use std::pin::Pin;
54use std::str::Utf8Error;
55use std::task::{Context, Poll, ready};
56
57use futures_core::Stream;
58
59/// The NDJSON lines of a byte stream, such as an
60/// [`HttpResponse`](crate::HttpResponse) body.
61///
62/// Yields each non-blank line, without its line ending, once the line is
63/// complete (see the [module docs](self) for the framing rules). After the
64/// first error, a failed read ([`NdjsonError::Transport`]), a line that is
65/// not UTF-8 ([`NdjsonError::InvalidUtf8`]) or a line past the size cap
66/// ([`NdjsonError::LineTooLarge`], see [`max_line_size`](Self::max_line_size)),
67/// the stream ends. It polls the body only when no complete line is
68/// buffered, so transport-level idle bounds
69/// ([`StreamTimeouts`](crate::StreamTimeouts)) see the same waits as
70/// before.
71///
72/// Drive an adapter's decoder over it with
73/// [`FramedChatStream::new`](crate::framing::FramedChatStream::new), which
74/// classifies a failure with [`NdjsonError::into_error`].
75#[derive(Debug)]
76pub struct NdjsonStream<S> {
77 /// The byte stream; `None` once it ended or failed.
78 body: Option<S>,
79 decoder: NdjsonDecoder,
80}
81
82impl<S> NdjsonStream<S> {
83 /// Frame `body` into NDJSON lines, with the default size cap
84 /// ([`NdjsonDecoder::DEFAULT_MAX_LINE_SIZE`]).
85 #[must_use]
86 pub fn new(body: S) -> Self {
87 Self {
88 body: Some(body),
89 decoder: NdjsonDecoder::new(),
90 }
91 }
92
93 /// Set the size cap, in bytes: see [`NdjsonDecoder::max_line_size`].
94 /// Provider crates set it from their builder's
95 /// `max_stream_event_size`.
96 #[must_use]
97 pub fn max_line_size(mut self, limit: usize) -> Self {
98 self.decoder = self.decoder.max_line_size(limit);
99 self
100 }
101
102 fn end(&mut self) {
103 self.body = None;
104 self.decoder = NdjsonDecoder::new();
105 }
106}
107
108impl<S, B, E> Stream for NdjsonStream<S>
109where
110 S: Stream<Item = Result<B, E>> + Unpin,
111 B: AsRef<[u8]>,
112{
113 type Item = Result<String, NdjsonError<E>>;
114
115 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
116 let this = self.get_mut();
117 loop {
118 match this.decoder.next_line() {
119 Ok(Some(line)) => return Poll::Ready(Some(Ok(line))),
120 Ok(None) => {}
121 Err(error) => {
122 this.end();
123 return Poll::Ready(Some(Err(error.widen())));
124 }
125 }
126 let Some(body) = this.body.as_mut() else {
127 return Poll::Ready(None);
128 };
129 match ready!(Pin::new(body).poll_next(cx)) {
130 Some(Ok(chunk)) => this.decoder.push(chunk.as_ref()),
131 Some(Err(error)) => {
132 this.end();
133 return Poll::Ready(Some(Err(NdjsonError::Transport(error))));
134 }
135 None => {
136 // What is left is a last line without its line ending.
137 let last = this.decoder.finish();
138 this.end();
139 return Poll::Ready(last.map_err(NdjsonError::widen).transpose());
140 }
141 }
142 }
143 }
144}
145
146/// An incremental NDJSON framer: push bytes, take complete lines.
147///
148/// [`NdjsonStream`] drives one over a byte stream; use it directly when the
149/// bytes do not come as a [`Stream`] (a callback API, a test).
150///
151/// Framing is linear: each byte is scanned for a line ending once, however
152/// the chunks split the lines. A size cap
153/// ([`max_line_size`](Self::max_line_size)) bounds the line still waiting
154/// for its ending, where memory grows from chunk to chunk.
155///
156/// ```
157/// use allama_http::ndjson::NdjsonDecoder;
158///
159/// let mut decoder = NdjsonDecoder::new();
160/// decoder.push(b"{\"a\":1}\r");
161/// assert_eq!(decoder.next_line(), Ok(None));
162/// decoder.push(b"\n\n{\"b\":");
163/// assert_eq!(decoder.next_line(), Ok(Some(r#"{"a":1}"#.to_owned())));
164/// assert_eq!(decoder.next_line(), Ok(None));
165/// decoder.push(b"2}");
166/// // The body ended: the last line needs no line ending.
167/// assert_eq!(decoder.finish(), Ok(Some(r#"{"b":2}"#.to_owned())));
168/// ```
169#[derive(Debug)]
170pub struct NdjsonDecoder {
171 /// Received bytes; those before `read` are consumed.
172 buffer: Vec<u8>,
173 read: usize,
174 /// The scanned offset: the bytes from `read` to here hold no `LF`.
175 /// [`next_line`](Self::next_line) looks for one from here on.
176 scanned: usize,
177 /// The bytes from here on hold no `LF` either: the line still waiting
178 /// for its ending. [`push`](Self::push) scans each chunk back to its
179 /// last `LF` to move it, and `next_line` looks for `LF`s only before
180 /// it, so each byte is scanned once.
181 unterminated: usize,
182 /// Whether the start of the stream (and a byte order mark) is behind.
183 started: bool,
184 max_line_size: usize,
185 /// The line waiting for its ending passed the size cap: the framer
186 /// fails once the complete lines before it are taken.
187 overflowing: bool,
188 /// The size cap was passed: framing cannot go on.
189 overflowed: bool,
190 /// Bytes examined for line endings, for the linearity test.
191 #[cfg(test)]
192 examined: usize,
193}
194
195impl Default for NdjsonDecoder {
196 fn default() -> Self {
197 Self::new()
198 }
199}
200
201impl NdjsonDecoder {
202 /// The default [`max_line_size`](Self::max_line_size): 64 MiB, as
203 /// [`SseDecoder::DEFAULT_MAX_EVENT_SIZE`](crate::sse::SseDecoder::DEFAULT_MAX_EVENT_SIZE),
204 /// for the same reasons: a final line can repeat the whole response,
205 /// and streamed base64 images are large.
206 pub const DEFAULT_MAX_LINE_SIZE: usize = 64 * 1024 * 1024;
207
208 /// A framer at the start of a stream, with the default size cap.
209 #[must_use]
210 pub fn new() -> Self {
211 Self {
212 buffer: Vec::new(),
213 read: 0,
214 scanned: 0,
215 unterminated: 0,
216 started: false,
217 max_line_size: Self::DEFAULT_MAX_LINE_SIZE,
218 overflowing: false,
219 overflowed: false,
220 #[cfg(test)]
221 examined: 0,
222 }
223 }
224
225 /// Set the size cap, in bytes (default
226 /// [`DEFAULT_MAX_LINE_SIZE`](Self::DEFAULT_MAX_LINE_SIZE)).
227 ///
228 /// It bounds a line, its line ending excluded. [`push`](Self::push)
229 /// checks the line still waiting for its ending before it buffers a
230 /// chunk, so that line never grows much past the cap.
231 ///
232 /// Past the cap, [`next_line`](Self::next_line) fails with
233 /// [`NdjsonError::LineTooLarge`], and framing cannot go on: the framer
234 /// drops what it holds, ignores later bytes and keeps returning the
235 /// error.
236 ///
237 /// ```
238 /// use allama_http::ndjson::{NdjsonDecoder, NdjsonError};
239 ///
240 /// let mut decoder = NdjsonDecoder::new().max_line_size(8);
241 /// decoder.push(b"{\"a\":12}");
242 /// assert_eq!(decoder.next_line(), Ok(None));
243 /// // A 9th byte without a line ending.
244 /// decoder.push(b"3");
245 /// assert!(matches!(
246 /// decoder.next_line(),
247 /// Err(NdjsonError::LineTooLarge { limit: 8, .. })
248 /// ));
249 /// ```
250 #[must_use]
251 pub fn max_line_size(mut self, limit: usize) -> Self {
252 self.max_line_size = limit;
253 self
254 }
255
256 /// Append received bytes. Chunk boundaries may fall anywhere, inside a
257 /// line, a `CRLF` or a UTF-8 sequence.
258 ///
259 /// Bytes that would take the line still waiting for its ending past the
260 /// [size cap](Self::max_line_size) are not buffered:
261 /// [`next_line`](Self::next_line) returns the lines completed before
262 /// that line, then fails.
263 pub fn push(&mut self, bytes: &[u8]) {
264 if self.overflowing || self.overflowed {
265 return;
266 }
267 if self.read > 0 {
268 self.buffer.drain(..self.read);
269 self.scanned = self.scanned.saturating_sub(self.read);
270 self.unterminated = self.unterminated.saturating_sub(self.read);
271 self.read = 0;
272 }
273 // Back to the chunk's last `LF`: what follows it starts, or
274 // continues, the line still waiting for its ending.
275 let last_ending = bytes.iter().rposition(|byte| *byte == b'\n');
276 #[cfg(test)]
277 {
278 self.examined += last_ending.map_or(bytes.len(), |end| bytes.len() - end);
279 }
280 let (unterminated, waiting) = match last_ending {
281 Some(end) => (self.buffer.len() + end + 1, bytes.len() - end - 1),
282 None => (
283 self.unterminated,
284 self.buffer.len() - self.unterminated + bytes.len(),
285 ),
286 };
287 // One byte of slack: the line's `CR`, which belongs to its ending.
288 let bytes = if waiting > self.max_line_size.saturating_add(1) {
289 // Keep the complete lines; the rest is never buffered.
290 self.overflowing = true;
291 let Some(end) = last_ending else {
292 return;
293 };
294 &bytes[..=end]
295 } else {
296 bytes
297 };
298 // A chunk without a line ending extends a line scanned to its end.
299 if last_ending.is_none() && self.scanned == self.buffer.len() {
300 self.scanned += bytes.len();
301 }
302 self.buffer.extend_from_slice(bytes);
303 self.unterminated = unterminated.min(self.buffer.len());
304 }
305
306 /// The next complete, non-blank line, without its line ending, or
307 /// `None` until more bytes arrive.
308 ///
309 /// # Errors
310 ///
311 /// - [`NdjsonError::InvalidUtf8`] when a line is not valid UTF-8. The
312 /// line is skipped, and framing can go on, but the stream is corrupt:
313 /// [`NdjsonStream`] ends there.
314 /// - [`NdjsonError::LineTooLarge`] past the
315 /// [size cap](Self::max_line_size). Framing cannot go on.
316 ///
317 /// A framer fails with no other error: it reads no transport.
318 pub fn next_line(&mut self) -> Result<Option<String>, NdjsonError<Infallible>> {
319 if self.overflowed {
320 return Err(self.too_large());
321 }
322 if !self.started {
323 let pending = &self.buffer[self.read..];
324 if !self.overflowing && pending.len() < BOM.len() && BOM.starts_with(pending) {
325 return Ok(None);
326 }
327 self.skip_bom();
328 }
329 loop {
330 // `LF`s lie only between the two scanned offsets.
331 let from = self.scanned.max(self.read);
332 let scan = &self.buffer[from..self.unterminated.max(from)];
333 let found = scan.iter().position(|byte| *byte == b'\n');
334 #[cfg(test)]
335 {
336 self.examined += found.map_or(scan.len(), |offset| offset + 1);
337 }
338 let Some(offset) = found else {
339 self.scanned = self.buffer.len();
340 let waiting = strip_cr(&self.buffer[self.read..]).len();
341 if self.overflowing || waiting > self.max_line_size {
342 self.overflow();
343 return Err(self.too_large());
344 }
345 return Ok(None);
346 };
347 let (start, end) = (self.read, from + offset);
348 self.read = end + 1;
349 self.scanned = self.read;
350 if let Some(line) = self.line(start, end)? {
351 return Ok(Some(line));
352 }
353 }
354 }
355
356 /// The last line, once the body ended: what is left without a line
357 /// ending, unless it is blank. Call it after
358 /// [`next_line`](Self::next_line) returned `None`; the framer is empty
359 /// afterwards.
360 ///
361 /// # Errors
362 ///
363 /// As [`next_line`](Self::next_line).
364 pub fn finish(&mut self) -> Result<Option<String>, NdjsonError<Infallible>> {
365 if self.overflowed || self.overflowing {
366 self.overflow();
367 return Err(self.too_large());
368 }
369 if !self.started {
370 self.skip_bom();
371 }
372 let (start, end) = (self.read, self.buffer.len());
373 self.read = end;
374 self.scanned = end;
375 self.unterminated = end;
376 self.line(start, end)
377 }
378
379 /// The line at `start..end`, its `CR` stripped: `None` when it is
380 /// blank.
381 fn line(
382 &mut self,
383 start: usize,
384 end: usize,
385 ) -> Result<Option<String>, NdjsonError<Infallible>> {
386 let line = strip_cr(&self.buffer[start..end]);
387 if line.len() > self.max_line_size {
388 self.overflow();
389 return Err(self.too_large());
390 }
391 if line.iter().all(|byte| matches!(byte, b' ' | b'\t' | b'\r')) {
392 return Ok(None);
393 }
394 let line = std::str::from_utf8(line).map_err(NdjsonError::InvalidUtf8)?;
395 Ok(Some(line.to_owned()))
396 }
397
398 fn skip_bom(&mut self) {
399 if self.buffer[self.read..].starts_with(BOM) {
400 self.read += BOM.len();
401 self.scanned = self.scanned.max(self.read);
402 }
403 self.started = true;
404 }
405
406 /// Fail for good: drop what is held.
407 fn overflow(&mut self) {
408 self.overflowed = true;
409 self.buffer = Vec::new();
410 self.read = 0;
411 self.scanned = 0;
412 self.unterminated = 0;
413 }
414
415 fn too_large(&self) -> NdjsonError<Infallible> {
416 NdjsonError::LineTooLarge {
417 limit: self.max_line_size,
418 }
419 }
420}
421
422/// Why an [`NdjsonStream`] (or an [`NdjsonDecoder`], with `E` =
423/// [`Infallible`]) failed.
424#[derive(Debug, PartialEq, Eq)]
425#[non_exhaustive]
426pub enum NdjsonError<E> {
427 /// Reading the body failed: the body stream's own error. Map it as you
428 /// would any mid-stream transport failure (a
429 /// [`TransportError`](crate::TransportError) keeps its kind, so a
430 /// stream timeout stays a timeout).
431 Transport(E),
432 /// A line was not valid UTF-8. JSON text is UTF-8, and a replaced byte
433 /// would change a value without anything saying so, which allama never
434 /// does silently (ADR 0003). Invalid UTF-8 means a broken server or
435 /// intermediary, and asking again streams the same bytes: map it to a
436 /// stream decode error ([`allama::Error::stream_decode`]), which is not
437 /// retryable. A chunk boundary inside a UTF-8 sequence is not an error:
438 /// each line is decoded once it is complete.
439 InvalidUtf8(Utf8Error),
440 /// A line passed the framer's size cap
441 /// ([`NdjsonDecoder::max_line_size`]). The stream is not one the caller
442 /// is prepared to hold, and asking again gets the same answer: map it
443 /// to a stream protocol error ([`allama::Error::stream_protocol`]), not
444 /// a retryable transport failure.
445 #[non_exhaustive]
446 LineTooLarge {
447 /// The cap, in bytes.
448 limit: usize,
449 },
450}
451
452impl<E: Into<crate::TransportError>> NdjsonError<E> {
453 /// The portable error for a failed NDJSON stream, as the variants say:
454 ///
455 /// - [`Transport`](Self::Transport): the body error as a
456 /// [`TransportError`](crate::TransportError), mapped by
457 /// [`TransportError::into_stream_error`](crate::TransportError::into_stream_error),
458 /// so a stream timeout stays [`allama::Error::Timeout`] and any other
459 /// failed read is a retryable transport stream error;
460 /// - [`InvalidUtf8`](Self::InvalidUtf8): a stream decode error;
461 /// - [`LineTooLarge`](Self::LineTooLarge): a stream protocol error.
462 ///
463 /// Neither of the last two is retryable, and every message is
464 /// [authored](allama::Error::authored). Add the provider and the request
465 /// id. [`FramedChatStream::new`](crate::framing::FramedChatStream::new)
466 /// maps failures with it.
467 ///
468 /// ```
469 /// use allama::{Error, StreamErrorKind};
470 /// use allama_http::TransportError;
471 /// use allama_http::ndjson::NdjsonStream;
472 /// use futures_util::{FutureExt, StreamExt, stream};
473 ///
474 /// let body = stream::iter([Ok::<_, TransportError>(&b"{\"a\":\"\xFF\"}\n"[..])]);
475 /// let failure = NdjsonStream::new(body).next().now_or_never().flatten().unwrap().unwrap_err();
476 /// let err = failure.into_error().with_provider("acme");
477 /// assert!(matches!(err, Error::Stream { kind: StreamErrorKind::Decode, .. }));
478 /// assert!(!err.is_retryable());
479 /// ```
480 #[must_use]
481 pub fn into_error(self) -> allama::Error {
482 match self {
483 Self::Transport(error) => error.into().into_stream_error(),
484 Self::InvalidUtf8(error) => allama::Error::stream_decode(
485 NdjsonError::<Infallible>::InvalidUtf8(error).to_string(),
486 )
487 .authored(),
488 Self::LineTooLarge { limit } => allama::Error::stream_protocol(
489 NdjsonError::<Infallible>::LineTooLarge { limit }.to_string(),
490 )
491 .authored(),
492 }
493 }
494}
495
496impl<E: Into<crate::TransportError>> From<NdjsonError<E>> for allama::Error {
497 /// The error [`NdjsonError::into_error`] gives: what
498 /// [`FramedChatStream::new`](crate::framing::FramedChatStream::new)
499 /// classifies a failure with.
500 fn from(error: NdjsonError<E>) -> Self {
501 error.into_error()
502 }
503}
504
505impl NdjsonError<Infallible> {
506 /// A framer's error as a stream's.
507 fn widen<E>(self) -> NdjsonError<E> {
508 match self {
509 Self::Transport(never) => match never {},
510 Self::InvalidUtf8(error) => NdjsonError::InvalidUtf8(error),
511 Self::LineTooLarge { limit } => NdjsonError::LineTooLarge { limit },
512 }
513 }
514}
515
516impl<E: fmt::Display> fmt::Display for NdjsonError<E> {
517 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
518 match self {
519 Self::Transport(error) => error.fmt(f),
520 Self::InvalidUtf8(error) => write!(f, "invalid UTF-8 in the NDJSON stream: {error}"),
521 Self::LineTooLarge { limit } => {
522 write!(f, "an NDJSON line passed the size cap of {limit} bytes")
523 }
524 }
525 }
526}
527
528impl<E: std::error::Error + 'static> std::error::Error for NdjsonError<E> {
529 /// A transport failure is transparent: its source is the body error's
530 /// source.
531 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
532 match self {
533 Self::Transport(error) => error.source(),
534 Self::InvalidUtf8(error) => Some(error),
535 Self::LineTooLarge { .. } => None,
536 }
537 }
538}
539
540/// The UTF-8 byte order mark, skipped at the start of a stream.
541const BOM: &[u8] = b"\xEF\xBB\xBF";
542
543/// `line` without the `CR` of a `CRLF` ending.
544fn strip_cr(line: &[u8]) -> &[u8] {
545 line.strip_suffix(b"\r").unwrap_or(line)
546}
547
548#[cfg(test)]
549mod tests;