Add a server-sent-events reader to loopd

Implemented-By: Laguna S 2.1 (OpenCode)

crates/loopd/src/lib.rs

@@ -2,3 +2,4 @@
 
 pub mod config;
 pub mod http;
+pub mod sse;

crates/loopd/src/sse.rs

@@ -0,0 +1,140 @@
+//! A server-sent-events reader: turns a streamed body into `data:` payloads.
+//!
+//! `llama-server` streams a completion as SSE — one JSON object per `data:` line, then
+//! `data: [DONE]`. Lines end in `\n` or `\r\n`. Blank lines, comment lines (`:`) and any field
+//! other than `data` carry nothing and are skipped. After `data:` one leading space is dropped, and
+//! only one. Every byte the peer sends is read in bounded chunks; a line is only UTF-8 decoded once
+//! its whole length has arrived.
+
+use std::io::{ErrorKind, Read};
+
+use crate::http::is_timeout;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum SseItem {
+    /// The text after `data:`, without the line ending.
+    Data(String),
+    Done,
+}
+
+#[derive(Debug)]
+pub enum SseError {
+    Io(std::io::Error),
+    Timeout,
+    Truncated,
+    LineTooLong,
+    NotUtf8,
+}
+
+impl std::fmt::Display for SseError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            SseError::Io(e) => write!(f, "io error: {e}"),
+            SseError::Timeout => write!(f, "read timed out before the line arrived"),
+            SseError::Truncated => write!(f, "the body ended in the middle of a line"),
+            SseError::LineTooLong => write!(f, "a line exceeded the limit"),
+            SseError::NotUtf8 => write!(f, "a line was not valid UTF-8"),
+        }
+    }
+}
+
+impl std::error::Error for SseError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match self {
+            SseError::Io(e) => Some(e),
+            _ => None,
+        }
+    }
+}
+
+pub struct Events<R> {
+    reader: R,
+    buf: Vec<u8>,
+    max_line: usize,
+    ended: bool,
+}
+
+impl<R: Read> Events<R> {
+    pub fn new(reader: R, max_line: usize) -> Self {
+        Events {
+            reader,
+            buf: Vec::new(),
+            max_line,
+            ended: false,
+        }
+    }
+
+    /// The next item, or `None` when the stream has ended cleanly between lines.
+    pub fn next_item(&mut self) -> Result<Option<SseItem>, SseError> {
+        loop {
+            if self.ended {
+                return Ok(None);
+            }
+            let line = match self.read_line() {
+                Ok(None) => {
+                    self.ended = true;
+                    return Ok(None);
+                }
+                Ok(Some(line)) => line,
+                Err(e) => return Err(e),
+            };
+            match self.process_line(line)? {
+                Some(item) => return Ok(Some(item)),
+                None => continue,
+            }
+        }
+    }
+
+    /// Reads one complete line, without its ending. Returns `None` on a clean end with nothing
+    /// buffered, and `Truncated` on a clean end with a partial line left over.
+    fn read_line(&mut self) -> Result<Option<Vec<u8>>, SseError> {
+        loop {
+            if let Some(pos) = self.buf.iter().position(|&b| b == b'\n') {
+                let mut line: Vec<u8> = self.buf.drain(..=pos).collect();
+                line.pop(); // drop the '\n'
+                if line.last() == Some(&b'\r') {
+                    line.pop();
+                }
+                if line.len() > self.max_line {
+                    return Err(SseError::LineTooLong);
+                }
+                return Ok(Some(line));
+            }
+            if self.buf.len() > self.max_line {
+                return Err(SseError::LineTooLong);
+            }
+            let mut tmp = [0u8; 4096];
+            let n = match self.reader.read(&mut tmp) {
+                Ok(0) => {
+                    if self.buf.is_empty() {
+                        return Ok(None);
+                    }
+                    return Err(SseError::Truncated);
+                }
+                Ok(n) => n,
+                Err(e) if e.kind() == ErrorKind::Interrupted => continue,
+                Err(e) if is_timeout(&e) => return Err(SseError::Timeout),
+                Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Err(SseError::Truncated),
+                Err(e) => return Err(SseError::Io(e)),
+            };
+            self.buf.extend_from_slice(&tmp[..n]);
+        }
+    }
+
+    fn process_line(&self, line: Vec<u8>) -> Result<Option<SseItem>, SseError> {
+        let rest = match line.strip_prefix(b"data:") {
+            Some(rest) => rest,
+            None => return Ok(None),
+        };
+        let rest = match rest.split_first() {
+            Some((&b' ', rest)) => rest,
+            _ => rest,
+        };
+        let text = std::str::from_utf8(rest).map_err(|_| SseError::NotUtf8)?;
+        if text == "[DONE]" {
+            Ok(Some(SseItem::Done))
+        } else {
+            Ok(Some(SseItem::Data(text.to_string())))
+        }
+    }
+}

crates/loopd/tests/sse.rs

@@ -0,0 +1,165 @@
+//! Tests for the SSE reader. Do not edit.
+
+mod support;
+
+use loopd::http::{Connection, Request};
+use loopd::sse::{Events, SseError, SseItem};
+use std::io::{Cursor, Read};
+use std::time::Duration;
+use support::{FakeServer, Reply};
+
+fn data(s: &str) -> Option<SseItem> {
+    Some(SseItem::Data(s.to_string()))
+}
+
+/// Hands out one byte per call.
+struct OneByte(Cursor<Vec<u8>>);
+
+impl Read for OneByte {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        let end = buf.len().min(1);
+        self.0.read(&mut buf[..end])
+    }
+}
+
+/// Serves `bytes`, then fails every read with `kind`.
+struct ThenError(Cursor<Vec<u8>>, std::io::ErrorKind);
+
+impl Read for ThenError {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        match self.0.read(buf)? {
+            0 => Err(std::io::Error::from(self.1)),
+            n => Ok(n),
+        }
+    }
+}
+
+const STREAM: &str =
+    "data: {\"a\":1}\n\n: a comment\ndata:{\"b\":2}\r\n\r\nevent: x\nid: 7\n\ndata: [DONE]\n\n";
+
+#[test]
+fn yields_data_lines_and_done() {
+    let mut ev = Events::new(Cursor::new(STREAM), 1024);
+    assert_eq!(ev.next_item().unwrap(), data("{\"a\":1}"));
+    assert_eq!(
+        ev.next_item().unwrap(),
+        data("{\"b\":2}"),
+        "no space after the colon, and CRLF"
+    );
+    assert_eq!(ev.next_item().unwrap(), Some(SseItem::Done));
+    assert_eq!(ev.next_item().unwrap(), None);
+    assert_eq!(
+        ev.next_item().unwrap(),
+        None,
+        "asking again after the end is harmless"
+    );
+}
+
+#[test]
+fn the_result_does_not_depend_on_how_the_bytes_arrive() {
+    let mut ev = Events::new(OneByte(Cursor::new(STREAM.as_bytes().to_vec())), 1024);
+    assert_eq!(ev.next_item().unwrap(), data("{\"a\":1}"));
+    assert_eq!(ev.next_item().unwrap(), data("{\"b\":2}"));
+    assert_eq!(ev.next_item().unwrap(), Some(SseItem::Done));
+    assert_eq!(ev.next_item().unwrap(), None);
+}
+
+#[test]
+fn only_a_leading_space_is_stripped_and_multibyte_text_survives() {
+    let mut ev = Events::new(
+        OneByte(Cursor::new(
+            "data:  two spaces\ndata: caf\u{e9} \u{1f4e6}\n".into(),
+        )),
+        1024,
+    );
+    assert_eq!(ev.next_item().unwrap(), data(" two spaces"));
+    assert_eq!(ev.next_item().unwrap(), data("caf\u{e9} \u{1f4e6}"));
+}
+
+#[test]
+fn a_stream_that_ends_mid_line_is_truncated() {
+    let mut ev = Events::new(Cursor::new("data: {\"a\":1}\n\ndata: {\"b\""), 1024);
+    assert_eq!(ev.next_item().unwrap(), data("{\"a\":1}"));
+    assert!(matches!(ev.next_item(), Err(SseError::Truncated)));
+}
+
+#[test]
+fn reader_errors_are_classified() {
+    use std::io::ErrorKind::{ConnectionReset, TimedOut, UnexpectedEof, WouldBlock};
+    let first = "data: 1\n";
+    for (kind, want) in [
+        (WouldBlock, "timeout"),
+        (TimedOut, "timeout"),
+        (UnexpectedEof, "truncated"),
+        (ConnectionReset, "io"),
+    ] {
+        let mut ev = Events::new(ThenError(Cursor::new(first.into()), kind), 1024);
+        assert_eq!(ev.next_item().unwrap(), data("1"));
+        let got = match ev.next_item() {
+            Err(SseError::Timeout) => "timeout",
+            Err(SseError::Truncated) => "truncated",
+            Err(SseError::Io(_)) => "io",
+            other => panic!("{kind:?}: {other:?}"),
+        };
+        assert_eq!(got, want, "{kind:?}");
+    }
+}
+
+#[test]
+fn long_lines_and_bad_text_are_errors() {
+    let long = format!("data: {}\n", "x".repeat(5000));
+    let mut ev = Events::new(Cursor::new(long.clone()), 100);
+    assert!(matches!(ev.next_item(), Err(SseError::LineTooLong)));
+    let mut ev = Events::new(OneByte(Cursor::new("x".repeat(5000).into())), 100);
+    assert!(
+        matches!(ev.next_item(), Err(SseError::LineTooLong)),
+        "with no newline at all"
+    );
+    let mut ev = Events::new(Cursor::new(long), 10_000);
+    assert!(
+        matches!(ev.next_item(), Ok(Some(SseItem::Data(_)))),
+        "under the limit it is fine"
+    );
+
+    let mut ev = Events::new(Cursor::new(b"data: \xff\xfe\n".to_vec()), 1024);
+    assert!(matches!(ev.next_item(), Err(SseError::NotUtf8)));
+    let e: Box<dyn std::error::Error> = Box::new(SseError::Truncated);
+    assert!(!e.to_string().is_empty());
+}
+
+/// The recorded streams, read through the HTTP client as `loopd` will read them.
+#[test]
+fn recorded_streams_have_the_expected_shape() {
+    for (name, want_data) in [("plain", 21), ("tool_call", 27), ("turn1", 6), ("turn2", 7)] {
+        for piece in [usize::MAX, 1, 13] {
+            let server = FakeServer::start();
+            server.route(
+                "/v1/chat/completions",
+                vec![Reply::fixture(name).trickle(piece, 0)],
+            );
+            let mut conn = Connection::open(&server.socket).unwrap();
+            conn.set_read_timeout(Duration::from_secs(5)).unwrap();
+            conn.send(&Request {
+                method: "POST",
+                path: "/v1/chat/completions",
+                body: Some(b"{}"),
+            })
+            .unwrap();
+            let head = conn.read_head().unwrap();
+            let mut ev = Events::new(conn.body(&head).unwrap(), 1024 * 1024);
+            let mut items = Vec::new();
+            while let Some(item) = ev.next_item().unwrap() {
+                items.push(item);
+            }
+            assert_eq!(items.last(), Some(&SseItem::Done), "{name}");
+            assert_eq!(items.len() - 1, want_data, "{name}: data items");
+            for item in &items[..items.len() - 1] {
+                let SseItem::Data(text) = item else {
+                    panic!("{name}: Done in the middle")
+                };
+                let json: serde_json::Value = serde_json::from_str(text).unwrap();
+                assert!(json["choices"].is_array(), "{name}");
+            }
+        }
+    }
+}