vtcode_bash_runner/
stream.rs

1use std::io;
2use tokio::io::{AsyncBufRead, AsyncBufReadExt};
3
4/// Result of a bounded line read.
5#[derive(Debug)]
6pub enum ReadLineResult {
7    Line(Vec<u8>),
8    Truncated(Vec<u8>),
9    Eof,
10}
11
12/// Read a line with a size limit, preventing unbounded memory growth.
13pub async fn read_line_with_limit<R: AsyncBufRead + Unpin>(
14    reader: &mut R,
15    buf: &mut Vec<u8>,
16    max_len: usize,
17) -> io::Result<ReadLineResult> {
18    buf.clear();
19    let mut total_read = 0;
20
21    loop {
22        let available = reader.fill_buf().await?;
23        if available.is_empty() {
24            return Ok(ReadLineResult::Eof);
25        }
26
27        if let Some(pos) = available.iter().position(|&b| b == b'\n') {
28            // Found a newline within the available buffer
29            let to_read = pos + 1;
30            let would_be_total = total_read + to_read;
31
32            if would_be_total <= max_len {
33                // Line fits within the limit
34                buf.extend_from_slice(&available[..to_read]);
35                reader.consume(to_read);
36                return Ok(ReadLineResult::Line(buf.clone()));
37            } else {
38                // Line would exceed the limit, so we truncate
39                let remaining_space = max_len.saturating_sub(total_read);
40                if remaining_space > 0 {
41                    buf.extend_from_slice(&available[..remaining_space]);
42                }
43                reader.consume(to_read); // Still consume the whole line from the reader
44                return Ok(ReadLineResult::Truncated(buf.clone()));
45            }
46        }
47
48        // No newline found in current buffer, add what we can
49        let len = available.len();
50        let would_be_total = total_read + len;
51
52        if would_be_total <= max_len {
53            // Buffer content fits within the limit
54            buf.extend_from_slice(available);
55            total_read = would_be_total;
56            reader.consume(len);
57        } else {
58            // Would exceed the limit, add only what fits
59            let remaining_space = max_len.saturating_sub(total_read);
60            if remaining_space > 0 {
61                buf.extend_from_slice(&available[..remaining_space]);
62            }
63            reader.consume(len);
64        }
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::{ReadLineResult, read_line_with_limit};
71    use tokio::io::BufReader;
72
73    #[tokio::test]
74    async fn read_line_with_limit_truncates() {
75        let data = "hello world\n";
76        let mut reader = BufReader::new(data.as_bytes());
77        let mut buf = Vec::new();
78
79        let result = read_line_with_limit(&mut reader, &mut buf, 5).await.unwrap();
80        match result {
81            ReadLineResult::Truncated(bytes) => {
82                assert!(!bytes.is_empty());
83            }
84            other => panic!("expected truncation, got {:?}", other),
85        }
86    }
87}
vtcode_bash_runner/stream.rs

vtcode_bash_runner/
stream.rs
