use super::{args::FILE_READ_MAX_BYTES, read_selector::LineSelector};
use crate::cancellation::AgentCancellation;
use std::{
fs::File,
io::{BufRead, BufReader, Read},
};
pub(super) const SCAN_MAX_BYTES: u64 = 16 * FILE_READ_MAX_BYTES;
const EXCERPT_MAX_BYTES: usize = FILE_READ_MAX_BYTES as usize - 256;
pub(super) struct RawExcerpt {
pub text: String,
pub request_complete: bool,
pub bytes_scanned: u64,
pub output_limited: bool,
}
pub(super) fn read_raw_excerpt(
file: File,
selector: LineSelector,
cancellation: &AgentCancellation,
) -> anyhow::Result<RawExcerpt> {
let (start, end) = match selector {
LineSelector::Single(line) => (line, line),
LineSelector::Range { start, end } => (start, end),
LineSelector::FromStart { end } => (1, end),
LineSelector::FromLine { start } => (start, usize::MAX),
LineSelector::Count { start, count } => (start, start.saturating_add(count - 1)),
};
let mut reader = BufReader::new(file.take(SCAN_MAX_BYTES));
let mut bytes_scanned = 0u64;
let mut line_number = 1usize;
let mut line = Vec::new();
let mut output = String::new();
loop {
cancellation.check()?;
let chunk = reader.fill_buf()?;
if chunk.is_empty() {
if bytes_scanned == SCAN_MAX_BYTES {
return Ok(RawExcerpt {
text: output,
request_complete: false,
bytes_scanned,
output_limited: false,
});
}
if !line.is_empty() && line_number >= start {
append_line(&mut output, &line)?;
}
return Ok(RawExcerpt {
text: output,
request_complete: true,
bytes_scanned,
output_limited: false,
});
}
let consumed = chunk
.iter()
.position(|byte| *byte == b'\n')
.map_or(chunk.len(), |n| n + 1);
let ends_line = chunk[consumed - 1] == b'\n';
if line_number >= start {
let body = &chunk[..consumed - usize::from(ends_line)];
if output
.len()
.saturating_add(line.len())
.saturating_add(body.len())
.saturating_add(1)
> EXCERPT_MAX_BYTES
{
return Ok(RawExcerpt {
text: output,
request_complete: false,
bytes_scanned,
output_limited: true,
});
}
line.extend_from_slice(body);
}
reader.consume(consumed);
bytes_scanned += consumed as u64;
if ends_line {
if line_number >= start {
append_line(&mut output, &line)?;
line.clear();
}
if line_number == end {
return Ok(RawExcerpt {
text: output,
request_complete: true,
bytes_scanned,
output_limited: false,
});
}
line_number = line_number.saturating_add(1);
}
}
}
fn append_line(output: &mut String, line: &[u8]) -> anyhow::Result<()> {
let line = line.strip_suffix(b"\r").unwrap_or(line);
output.push_str(std::str::from_utf8(line)?);
output.push('\n');
Ok(())
}