Skip to main content

agent_sessions/
raw.rs

1//! Exact JSONL framing without interpreting JSON or changing bytes.
2use crate::reader::input::read_record;
3use crate::{LineErrorKind, ReadError, ReadOptions, ReadStatus, ReadSummary, StreamError};
4use std::{
5    fs::File,
6    io::{BufRead, BufReader, Seek, SeekFrom},
7    path::Path,
8};
9
10#[derive(Debug, Clone)]
11pub struct RawReadOptions {
12    pub start_offset: u64,
13    pub stop_at_byte: Option<u64>,
14    pub max_read_bytes: Option<u64>,
15    pub max_line_bytes: Option<usize>,
16}
17impl Default for RawReadOptions {
18    fn default() -> Self {
19        Self {
20            start_offset: 0,
21            stop_at_byte: None,
22            max_read_bytes: Some(200 * 1024 * 1024),
23            max_line_bytes: Some(8 * 1024 * 1024),
24        }
25    }
26}
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct RawRecord {
29    pub byte_start: u64,
30    pub byte_end: u64,
31    pub line_no: u64,
32    /// Exact bytes, including CR/LF delimiter when present.
33    pub bytes: Vec<u8>,
34    pub terminated: bool,
35}
36#[derive(Debug, Clone)]
37pub struct RawSummary {
38    pub status: ReadStatus,
39    pub records: u64,
40    pub bytes_read: u64,
41    /// End of the contiguous delivered raw prefix, not a JSON-validity promise.
42    pub delivered_through: u64,
43    pub line_errors: u64,
44}
45pub struct RawReader<R> {
46    reader: R,
47    options: ReadOptions,
48    summary: ReadSummary,
49    start: u64,
50    ended: bool,
51}
52pub fn read_raw_file(
53    path: &Path,
54    options: &RawReadOptions,
55) -> Result<RawReader<BufReader<File>>, ReadError> {
56    let mut file = File::open(path).map_err(ReadError::Io)?;
57    let mut options = options.clone();
58    if options.stop_at_byte.is_none() {
59        options.stop_at_byte = Some(file.metadata().map_err(ReadError::Io)?.len());
60    }
61    file.seek(SeekFrom::Start(options.start_offset))
62        .map_err(ReadError::Io)?;
63    read_raw_from(BufReader::new(file), &options)
64}
65/// `source` must already be positioned at options.start_offset.
66pub fn read_raw_from<R: BufRead>(
67    source: R,
68    options: &RawReadOptions,
69) -> Result<RawReader<R>, ReadError> {
70    if options
71        .stop_at_byte
72        .is_some_and(|end| end < options.start_offset)
73    {
74        return Err(ReadError::InvalidOptions(
75            "raw boundary precedes start offset",
76        ));
77    }
78    let limit = options
79        .max_read_bytes
80        .map(|n| {
81            options
82                .start_offset
83                .checked_add(n)
84                .ok_or(ReadError::InvalidOptions(
85                    "raw read budget overflows offset",
86                ))
87        })
88        .transpose()?;
89    Ok(RawReader {
90        reader: source,
91        options: ReadOptions {
92            stop_at_byte: options.stop_at_byte,
93            max_file_bytes: limit,
94            max_line_bytes: options.max_line_bytes,
95            ..Default::default()
96        },
97        summary: ReadSummary {
98            bytes_read: options.start_offset,
99            last_complete_byte: options.start_offset,
100            ..Default::default()
101        },
102        start: options.start_offset,
103        ended: false,
104    })
105}
106impl<R: BufRead> Iterator for RawReader<R> {
107    type Item = Result<RawRecord, StreamError>;
108    fn next(&mut self) -> Option<Self::Item> {
109        if self.ended {
110            return None;
111        }
112        let start = self.summary.bytes_read;
113        match read_record(
114            &mut self.reader,
115            &self.options,
116            &mut self.summary,
117            Vec::new(),
118        ) {
119            Ok(Some(record)) => {
120                self.summary.lines += 1;
121                if record.too_long {
122                    self.summary.line_errors += 1;
123                    return Some(Err(StreamError::Line {
124                        line_no: self.summary.lines,
125                        byte_start: start,
126                        kind: LineErrorKind::TooLong,
127                    }));
128                }
129                if self.summary.line_errors == 0 {
130                    self.summary.last_complete_byte = self.summary.bytes_read;
131                }
132                Some(Ok(RawRecord {
133                    byte_start: start,
134                    byte_end: self.summary.bytes_read,
135                    line_no: self.summary.lines,
136                    bytes: record.bytes,
137                    terminated: record.newline,
138                }))
139            }
140            Ok(None) => {
141                self.ended = true;
142                self.summary.status = if self.summary.line_errors == 0 {
143                    ReadStatus::Complete
144                } else {
145                    ReadStatus::CompleteWithErrors
146                };
147                None
148            }
149            Err(e) => {
150                self.ended = true;
151                self.summary.status = ReadStatus::Failed;
152                Some(Err(e))
153            }
154        }
155    }
156}
157impl<R: BufRead> RawReader<R> {
158    /// Next physical record number, including blank records. After an I/O error
159    /// this identifies the record whose read failed.
160    pub fn next_line_no(&self) -> u64 {
161        self.summary.lines.saturating_add(1)
162    }
163    pub fn finish(self) -> RawSummary {
164        RawSummary {
165            status: self.summary.status,
166            records: self.summary.lines,
167            bytes_read: self.summary.bytes_read - self.start,
168            delivered_through: self.summary.last_complete_byte,
169            line_errors: self.summary.line_errors,
170        }
171    }
172}