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(crate) 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
20    loop {
21        let available = reader.fill_buf().await?;
22        if available.is_empty() {
23            // EOF with a partial line left in the buffer: report it as a
24            // (truncated, if it overran) line rather than silently dropping
25            // the data we already read.
26            return Ok(if buf.is_empty() {
27                ReadLineResult::Eof
28            } else if buf.len() <= max_len {
29                ReadLineResult::Line(std::mem::take(buf))
30            } else {
31                ReadLineResult::Truncated(std::mem::take(buf))
32            });
33        }
34
35        if let Some(pos) = memchr::memchr(b'\n', available) {
36            // Found a newline within the available buffer
37            let to_read = pos + 1;
38            let would_be_total = buf.len() + to_read;
39
40            if would_be_total <= max_len {
41                // Line fits within the limit
42                buf.extend_from_slice(&available[..to_read]);
43                reader.consume(to_read);
44                return Ok(ReadLineResult::Line(std::mem::take(buf)));
45            } else {
46                // Line would exceed the limit: fill the remaining space, then
47                // report truncation. Consume the whole line from the reader so
48                // the next read starts at a fresh line.
49                let remaining_space = max_len.saturating_sub(buf.len());
50                if remaining_space > 0 {
51                    buf.extend_from_slice(&available[..remaining_space]);
52                }
53                reader.consume(to_read);
54                return Ok(ReadLineResult::Truncated(std::mem::take(buf)));
55            }
56        }
57
58        // No newline found in current buffer, add what we can
59        let len = available.len();
60        let would_be_total = buf.len() + len;
61
62        if would_be_total <= max_len {
63            // Buffer content fits within the limit
64            buf.extend_from_slice(available);
65            reader.consume(len);
66        } else {
67            // Would exceed the limit, add only what fits and report truncation
68            // immediately: continuing would either grow the buffer past
69            // `max_len` or block forever on a stream that never sends `\n`.
70            let remaining_space = max_len.saturating_sub(buf.len());
71            if remaining_space > 0 {
72                buf.extend_from_slice(&available[..remaining_space]);
73            }
74            reader.consume(len);
75            return Ok(ReadLineResult::Truncated(std::mem::take(buf)));
76        }
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::{ReadLineResult, read_line_with_limit};
83    use tokio::io::BufReader;
84
85    #[tokio::test]
86    async fn read_line_with_limit_truncates() -> std::io::Result<()> {
87        let data = "hello world\n";
88        let mut reader = BufReader::new(data.as_bytes());
89        let mut buf = Vec::new();
90
91        let result = read_line_with_limit(&mut reader, &mut buf, 5).await?;
92        match result {
93            ReadLineResult::Truncated(bytes) => {
94                assert!(!bytes.is_empty());
95            }
96            other => panic!("expected truncation, got {other:?}"),
97        }
98        Ok(())
99    }
100
101    #[tokio::test]
102    async fn read_line_with_limit_no_newline_stream_returns_truncated() -> std::io::Result<()> {
103        // A stream with no '\n' must terminate with Truncated once the limit is
104        // reached instead of growing the buffer or blocking forever.
105        let data = vec![b'a'; 10_000];
106        let mut reader = BufReader::new(data.as_slice());
107        let mut buf = Vec::new();
108
109        let result = read_line_with_limit(&mut reader, &mut buf, 100).await?;
110        match result {
111            ReadLineResult::Truncated(bytes) => {
112                assert_eq!(bytes.len(), 100, "buffer must not exceed the limit");
113            }
114            other => panic!("expected truncation, got {other:?}"),
115        }
116        Ok(())
117    }
118
119    #[tokio::test]
120    async fn read_line_with_limit_reports_partial_line_at_eof() -> std::io::Result<()> {
121        // Data without a trailing newline must be surfaced as a line rather
122        // than silently dropped.
123        let data = "no newline at end";
124        let mut reader = BufReader::new(data.as_bytes());
125        let mut buf = Vec::new();
126
127        let result = read_line_with_limit(&mut reader, &mut buf, 100).await?;
128        match result {
129            ReadLineResult::Line(bytes) => {
130                assert_eq!(bytes, b"no newline at end");
131            }
132            other => panic!("expected line, got {other:?}"),
133        }
134        Ok(())
135    }
136}
vtcode_bash_runner/stream.rs

vtcode_bash_runner/
stream.rs
