rustlavel_client/
stream.rs

1//! Streaming response bodies, and the server-sent events AI providers use.
2
3use crate::{Connection, find_head_end, parse_head};
4use rustlavel_core::{Error, Json, Result};
5use rustlavel_http::{Headers, Status};
6use std::time::Duration;
7
8/// A response whose body is read as it arrives.
9pub struct Body {
10    pub status: Status,
11    pub headers: Headers,
12    source: Source,
13    buffer: Vec<u8>,
14    /// Set once the transfer is complete, so `chunk` stops asking for more.
15    finished: bool,
16    /// Bytes still expected, when the body is chunked.
17    chunk_remaining: Option<usize>,
18    chunked: bool,
19    timeout: Duration,
20}
21
22enum Source {
23    Live(Box<Connection>),
24    /// A faked body, already complete.
25    Memory,
26}
27
28/// Send a request and hand back the body before it has finished arriving.
29pub(crate) async fn open(
30    mut connection: Connection,
31    request: Vec<u8>,
32    timeout: Duration,
33) -> Result<Body> {
34    connection.write_all(&request).await.map_err(Error::Io)?;
35    connection.flush().await.map_err(Error::Io)?;
36
37    let mut buffer = Vec::with_capacity(8 * 1024);
38    let head_end = loop {
39        if let Some(at) = find_head_end(&buffer) {
40            break at;
41        }
42        let mut chunk = [0u8; 4096];
43        let read = tokio::time::timeout(timeout, connection.read(&mut chunk))
44            .await
45            .map_err(|_| Error::msg("timed out waiting for response headers"))?
46            .map_err(Error::Io)?;
47        if read == 0 {
48            return Err(Error::Protocol("the server closed before sending headers".into()));
49        }
50        buffer.extend_from_slice(&chunk[..read]);
51    };
52
53    let (status, headers) = parse_head(&buffer[..head_end])?;
54    let rest = buffer.split_off(head_end);
55    let chunked = headers.get("transfer-encoding").is_some_and(|te| te.contains("chunked"));
56
57    Ok(Body {
58        status,
59        headers,
60        source: Source::Live(Box::new(connection)),
61        buffer: rest,
62        finished: false,
63        chunk_remaining: None,
64        chunked,
65        timeout,
66    })
67}
68
69impl Body {
70    /// A body that is already in memory, for fakes and tests.
71    pub fn from_bytes(status: Status, headers: Headers, body: Vec<u8>) -> Body {
72        Body {
73            status,
74            headers,
75            source: Source::Memory,
76            buffer: body,
77            finished: true,
78            chunk_remaining: None,
79            chunked: false,
80            timeout: Duration::from_secs(30),
81        }
82    }
83
84    /// The next piece of the body, or `None` when it is complete.
85    pub async fn chunk(&mut self) -> Result<Option<Vec<u8>>> {
86        loop {
87            if self.chunked {
88                if let Some(chunk) = self.take_chunked()? {
89                    return Ok(Some(chunk));
90                }
91            } else if !self.buffer.is_empty() {
92                return Ok(Some(std::mem::take(&mut self.buffer)));
93            }
94
95            if self.finished {
96                return Ok(None);
97            }
98            if !self.fill().await? {
99                self.finished = true;
100                if !self.chunked && !self.buffer.is_empty() {
101                    return Ok(Some(std::mem::take(&mut self.buffer)));
102                }
103                return Ok(None);
104            }
105        }
106    }
107
108    /// Read the rest of the body into memory.
109    pub async fn bytes(mut self) -> Result<Vec<u8>> {
110        let mut out = Vec::new();
111        while let Some(chunk) = self.chunk().await? {
112            out.extend_from_slice(&chunk);
113        }
114        Ok(out)
115    }
116
117    pub async fn text(self) -> Result<String> {
118        Ok(String::from_utf8_lossy(&self.bytes().await?).into_owned())
119    }
120
121    /// Read the body as a server-sent event stream.
122    pub fn events(self) -> SseReader {
123        SseReader { body: self, pending: String::new() }
124    }
125
126    /// Pull one chunk out of a chunked transfer, if a whole one is buffered.
127    fn take_chunked(&mut self) -> Result<Option<Vec<u8>>> {
128        loop {
129            match self.chunk_remaining {
130                None => {
131                    let Some(line_end) = find_crlf(&self.buffer) else { return Ok(None) };
132                    let header: Vec<u8> = self.buffer.drain(..line_end + 2).collect();
133                    let text = String::from_utf8_lossy(&header[..line_end]);
134                    let size =
135                        usize::from_str_radix(text.split(';').next().unwrap_or("").trim(), 16)
136                            .map_err(|_| Error::Protocol("invalid chunk size".into()))?;
137                    if size == 0 {
138                        self.finished = true;
139                        return Ok(None);
140                    }
141                    self.chunk_remaining = Some(size);
142                }
143                Some(size) => {
144                    if self.buffer.len() < size + 2 {
145                        return Ok(None);
146                    }
147                    let chunk: Vec<u8> = self.buffer.drain(..size).collect();
148                    self.buffer.drain(..2);
149                    self.chunk_remaining = None;
150                    return Ok(Some(chunk));
151                }
152            }
153        }
154    }
155
156    async fn fill(&mut self) -> Result<bool> {
157        let Source::Live(connection) = &mut self.source else {
158            return Ok(false);
159        };
160
161        let mut chunk = [0u8; 8192];
162        let read = tokio::time::timeout(self.timeout, connection.read(&mut chunk))
163            .await
164            .map_err(|_| Error::msg("timed out reading the response body"))?
165            .map_err(Error::Io)?;
166
167        self.buffer.extend_from_slice(&chunk[..read]);
168        Ok(read > 0)
169    }
170}
171
172fn find_crlf(buffer: &[u8]) -> Option<usize> {
173    buffer.windows(2).position(|w| w == b"\r\n")
174}
175
176/// One server-sent event.
177#[derive(Debug, Clone, PartialEq, Eq, Default)]
178pub struct ServerSentEvent {
179    pub event: Option<String>,
180    pub data: String,
181    pub id: Option<String>,
182}
183
184impl ServerSentEvent {
185    /// The data parsed as JSON, which is how every AI provider sends it.
186    pub fn json(&self) -> Result<Json> {
187        Json::parse(&self.data)
188    }
189
190    /// Whether this is the conventional end-of-stream marker.
191    pub fn is_done(&self) -> bool {
192        self.data.trim() == "[DONE]"
193    }
194}
195
196/// Reads a body as a sequence of server-sent events.
197pub struct SseReader {
198    body: Body,
199    /// Text received but not yet ending in a blank line.
200    pending: String,
201}
202
203impl SseReader {
204    /// The next event, or `None` at the end of the stream.
205    pub async fn next(&mut self) -> Result<Option<ServerSentEvent>> {
206        loop {
207            if let Some(event) = self.take_event() {
208                return Ok(Some(event));
209            }
210            match self.body.chunk().await? {
211                Some(bytes) => self.pending.push_str(&String::from_utf8_lossy(&bytes)),
212                None => return Ok(self.take_event().or_else(|| self.take_remainder())),
213            }
214        }
215    }
216
217    /// Collect every event, for a caller that does not need them as they arrive.
218    pub async fn collect(mut self) -> Result<Vec<ServerSentEvent>> {
219        let mut events = Vec::new();
220        while let Some(event) = self.next().await? {
221            events.push(event);
222        }
223        Ok(events)
224    }
225
226    fn take_event(&mut self) -> Option<ServerSentEvent> {
227        // An event ends at a blank line; both line endings appear in the wild.
228        let end = self
229            .pending
230            .find("\n\n")
231            .map(|at| (at, 2))
232            .or_else(|| self.pending.find("\r\n\r\n").map(|at| (at, 4)))?;
233
234        let block: String = self.pending.drain(..end.0 + end.1).collect();
235        parse_event(&block)
236    }
237
238    fn take_remainder(&mut self) -> Option<ServerSentEvent> {
239        if self.pending.trim().is_empty() {
240            return None;
241        }
242        let block = std::mem::take(&mut self.pending);
243        parse_event(&block)
244    }
245}
246
247fn parse_event(block: &str) -> Option<ServerSentEvent> {
248    let mut event = ServerSentEvent::default();
249    let mut data_lines: Vec<&str> = Vec::new();
250
251    for line in block.lines() {
252        let line = line.trim_end_matches('\r');
253        // A line starting with `:` is a comment, used as a keep-alive.
254        if line.is_empty() || line.starts_with(':') {
255            continue;
256        }
257        let (field, value) = match line.split_once(':') {
258            Some((field, value)) => (field, value.strip_prefix(' ').unwrap_or(value)),
259            None => (line, ""),
260        };
261        match field {
262            "event" => event.event = Some(value.to_string()),
263            "data" => data_lines.push(value),
264            "id" => event.id = Some(value.to_string()),
265            _ => {}
266        }
267    }
268
269    if data_lines.is_empty() && event.event.is_none() {
270        return None;
271    }
272    event.data = data_lines.join("\n");
273    Some(event)
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279
280    fn body(text: &str) -> Body {
281        Body::from_bytes(Status::OK, Headers::new(), text.as_bytes().to_vec())
282    }
283
284    #[tokio::test]
285    async fn reads_a_stream_of_events() {
286        let mut events = body("data: one\n\ndata: two\n\ndata: [DONE]\n\n").events();
287
288        assert_eq!(events.next().await.unwrap().unwrap().data, "one");
289        assert_eq!(events.next().await.unwrap().unwrap().data, "two");
290        assert!(events.next().await.unwrap().unwrap().is_done());
291        assert!(events.next().await.unwrap().is_none());
292    }
293
294    #[tokio::test]
295    async fn keeps_named_events_and_ids() {
296        let events = body("event: message_start\nid: 42\ndata: {\"a\":1}\n\n")
297            .events()
298            .collect()
299            .await
300            .unwrap();
301
302        assert_eq!(events.len(), 1);
303        assert_eq!(events[0].event.as_deref(), Some("message_start"));
304        assert_eq!(events[0].id.as_deref(), Some("42"));
305        assert_eq!(events[0].json().unwrap().get("a").unwrap().as_i64(), Some(1));
306    }
307
308    #[tokio::test]
309    async fn joins_multi_line_data_and_skips_comments() {
310        let events = body(": keep-alive\n\ndata: first\ndata: second\n\n").events().collect().await.unwrap();
311
312        assert_eq!(events.len(), 1);
313        assert_eq!(events[0].data, "first\nsecond");
314    }
315
316    #[tokio::test]
317    async fn handles_crlf_line_endings() {
318        let events = body("data: one\r\n\r\ndata: two\r\n\r\n").events().collect().await.unwrap();
319
320        assert_eq!(events.len(), 2);
321        assert_eq!(events[1].data, "two");
322    }
323
324    #[tokio::test]
325    async fn an_unterminated_final_event_is_still_delivered() {
326        let events = body("data: one\n\ndata: trailing").events().collect().await.unwrap();
327
328        assert_eq!(events.len(), 2);
329        assert_eq!(events[1].data, "trailing");
330    }
331
332    #[tokio::test]
333    async fn reads_a_whole_body_as_text() {
334        assert_eq!(body("hello").text().await.unwrap(), "hello");
335    }
336}
rustlavel_client/stream.rs

rustlavel_client/
stream.rs
