jules_api/streaming/
buffer.rs

1//! Chunk buffer for streaming responses.
2
3/// Error indicating that the `ChunkBuffer` capacity has been exceeded.
4#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
5#[error("Chunk buffer capacity exceeded")]
6pub struct BufferCapacityError;
7
8/// Buffer that accumulates chunks and enforces a maximum byte capacity.
9#[derive(Debug, Default)]
10pub struct ChunkBuffer {
11    buffer: String,
12    max_capacity: usize,
13}
14
15impl ChunkBuffer {
16    /// Creates a new `ChunkBuffer` with the specified maximum capacity.
17    #[must_use]
18    pub fn new(max_capacity: usize) -> Self {
19        Self {
20            buffer: String::with_capacity(max_capacity.min(8192)),
21            max_capacity,
22        }
23    }
24
25    /// Pushes a string chunk into the buffer.
26    ///
27    /// # Errors
28    ///
29    /// Returns `BufferCapacityError` if pushing the chunk would exceed the maximum capacity.
30    pub fn push(&mut self, chunk: &str) -> Result<(), BufferCapacityError> {
31        if self.buffer.len() + chunk.len() > self.max_capacity {
32            return Err(BufferCapacityError);
33        }
34        self.buffer.push_str(chunk);
35        Ok(())
36    }
37
38    /// Drains up to `max_bytes` from the beginning of the buffer.
39    pub fn drain(&mut self, max_bytes: usize) -> String {
40        let drain_len = self.buffer.len().min(max_bytes);
41
42        // Ensure we split at a char boundary.
43        let mut split_idx = drain_len;
44        while split_idx > 0 && !self.buffer.is_char_boundary(split_idx) {
45            split_idx -= 1;
46        }
47
48        let drained = self.buffer[..split_idx].to_string();
49        self.buffer.drain(..split_idx);
50        drained
51    }
52
53    /// Returns the current number of bytes in the buffer.
54    #[must_use]
55    pub fn len(&self) -> usize {
56        self.buffer.len()
57    }
58
59    /// Returns whether the buffer is empty.
60    #[must_use]
61    pub fn is_empty(&self) -> bool {
62        self.buffer.is_empty()
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    #[test]
71    fn test_push_and_drain() {
72        let mut buffer = ChunkBuffer::new(100);
73        buffer.push("hello").unwrap();
74        buffer.push(" world").unwrap();
75
76        assert_eq!(buffer.len(), 11);
77
78        let drained = buffer.drain(5);
79        assert_eq!(drained, "hello");
80        assert_eq!(buffer.len(), 6);
81
82        let drained2 = buffer.drain(10);
83        assert_eq!(drained2, " world");
84        assert!(buffer.is_empty());
85    }
86
87    #[test]
88    fn test_capacity_exceeded() {
89        let mut buffer = ChunkBuffer::new(10);
90        buffer.push("12345").unwrap();
91
92        let res = buffer.push("678901");
93        assert_eq!(res.unwrap_err(), BufferCapacityError);
94        assert_eq!(buffer.len(), 5); // Unchanged
95    }
96
97    #[test]
98    fn test_char_boundary() {
99        let mut buffer = ChunkBuffer::new(100);
100        // '🚀' is 4 bytes
101        buffer.push("a🚀b").unwrap();
102
103        // Try to drain 2 bytes, should split before the rocket.
104        let drained = buffer.drain(2);
105        assert_eq!(drained, "a");
106
107        let drained2 = buffer.drain(10);
108        assert_eq!(drained2, "🚀b");
109    }
110}
jules_api/streaming/buffer.rs

jules_api/streaming/
buffer.rs
