diff --git a/docs/architecture.md b/docs/architecture.md index 0378b23..e0b28bc 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -173,6 +173,14 @@ The writer trims an incomplete tail group before resuming appends. Malformed completed records and out-of-sequence refs fail loading. The current pre-release format intentionally does not load older run-record versions. +One incremental checkpoint decoder owns this definition for every reader. Its +synchronous core validates complete physical lines from any known sequence and +byte offset, retains at most one incomplete checkpoint, and returns exact raw +lines only after the whole group commits. The async file wrapper retains a torn +final line across EOF and streams forward without reading the whole file; +trajectory loading collects its committed records, while append recovery only +counts them and uses the same committed byte offset. + The message file is created and directory-synced with the run. The writer's cached next sequence is invalidated before cancellable I/O and restored only after a complete record has synced. Multiple independent writers for one run diff --git a/docs/source-map.md b/docs/source-map.md index b6bbb7a..04581d8 100644 --- a/docs/source-map.md +++ b/docs/source-map.md @@ -50,6 +50,11 @@ spilled result. - `src/storage/mod.rs`: run directories, metadata, events, and shared JSON persistence helpers. +- `src/storage/message_log.rs`: message checkpoint append, committed-prefix + loading, and writer-tail recovery. +- `src/storage/message_log/decoder.rs`: synchronous incremental checkpoint + validation plus the forward async reader; it preserves exact raw lines and + committed byte offsets without exposing partial groups. - `src/storage/trajectory.rs`: classified append-only messages and compacted-history loading. - `src/skills/mod.rs`: Agent Skills metadata discovery and body/path loading; diff --git a/src/storage/message_log.rs b/src/storage/message_log.rs index 8f221c1..8eb4238 100644 --- a/src/storage/message_log.rs +++ b/src/storage/message_log.rs @@ -5,7 +5,7 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use tokio::{ fs::OpenOptions, - io::{AsyncSeekExt, AsyncWriteExt}, + io::{AsyncSeekExt, AsyncWriteExt, BufReader}, }; use crate::{ @@ -13,6 +13,10 @@ use crate::{ trajectory::{CompactionMessage, TrajectoryMessage, message_ref, message_ref_seq}, }; +pub(crate) mod decoder; + +use decoder::{CommittedCheckpointReader, CommittedRecord, DecodeResult}; + /// One durable transcript line. The provider-neutral content blocks are kept /// directly in the record so readers do not need a second file to reconstruct /// the message that the runner will replay. @@ -62,6 +66,7 @@ struct MessageCheckpoint { struct JsonlFile { records: Vec, + record_count: u64, original_len: u64, committed_end: u64, } @@ -131,13 +136,13 @@ pub(super) async fn append_checkpoint(path: &Path, records: &[TrajectoryMessage] /// newline-terminated message lines are present, so viewers can ignore both a /// crash-torn line and complete lines from an incomplete final checkpoint. pub(super) async fn load(path: &Path) -> Result> { - Ok(read_jsonl(path).await?.records) + Ok(read_jsonl(path, true).await?.records) } /// Validate the committed prefix and remove an uncommitted tail before the /// sole writer resumes appending. pub(super) async fn prepare_append(path: &Path) -> Result { - let log = read_jsonl(path).await?; + let log = read_jsonl(path, false).await?; if log.original_len != log.committed_end { let mut file = OpenOptions::new() .write(true) @@ -154,92 +159,42 @@ pub(super) async fn prepare_append(path: &Path) -> Result { .await .with_context(|| format!("sync recovered trajectory file {}", path.display()))?; } - Ok(log.records.len() as u64) + Ok(log.record_count) } -async fn read_jsonl(path: &Path) -> Result { - let bytes = tokio::fs::read(path) +async fn read_jsonl(path: &Path, collect_records: bool) -> Result { + let file = tokio::fs::File::open(path) .await .with_context(|| format!("read initialized message log {}", path.display()))?; - let original_len = bytes.len() as u64; - let complete_lines = bytes - .split_inclusive(|byte| *byte == b'\n') - .scan(0_usize, |end, line| { - if !line.ends_with(b"\n") { - return None; - } - *end += line.len(); - Some((line, *end)) - }) - .collect::>(); + let mut reader = CommittedCheckpointReader::new(BufReader::new(file), path.to_owned()); let mut records = Vec::new(); - let mut committed_end = 0_usize; - let mut line_index = 0_usize; - - while line_index < complete_lines.len() { - let first = parse_stored_line(path, complete_lines[line_index].0)?; - let Some(checkpoint) = first.local.checkpoint.clone() else { - // Pre-checkpoint logs are singleton records. New appends always - // persist explicit count=1 checkpoint metadata. - let record = trajectory_record(first, records.len() as u64 + 1)?; - records.push(record); - committed_end = complete_lines[line_index].1; - line_index += 1; - continue; - }; - ensure!( - checkpoint.count > 0, - "message checkpoint count must be positive" - ); - ensure!( - checkpoint.index == 0, - "message checkpoint `{}` starts at index {} instead of 0", - checkpoint.first_message_ref, - checkpoint.index - ); - ensure!( - checkpoint.first_message_ref == first.message_ref, - "message checkpoint `{}` starts with message `{}`", - checkpoint.first_message_ref, - first.message_ref - ); - let checkpoint_count = usize::try_from(checkpoint.count) - .context("message checkpoint count does not fit in memory")?; - let available = complete_lines.len().saturating_sub(line_index); - let inspected = checkpoint_count.min(available); - let mut checkpoint_records = Vec::with_capacity(inspected); - for offset in 0..inspected { - let stored = parse_stored_line(path, complete_lines[line_index + offset].0)?; - let actual = stored.local.checkpoint.as_ref().with_context(|| { - format!( - "message `{}` is missing checkpoint metadata inside group `{}`", - stored.message_ref, checkpoint.first_message_ref - ) - })?; - ensure!( - actual.first_message_ref == checkpoint.first_message_ref - && actual.count == checkpoint.count - && actual.index == offset as u64, - "message `{}` has inconsistent checkpoint metadata", - stored.message_ref - ); - checkpoint_records.push(trajectory_record( - stored, - records.len() as u64 + offset as u64 + 1, - )?); - } - if inspected < checkpoint_count { - break; + let mut record_count = 0_u64; + let mut committed_end = 0; + while let DecodeResult::Checkpoint(checkpoint) = reader.next_checkpoint().await? { + committed_end = checkpoint.committed_end; + record_count = record_count + .checked_add(checkpoint.records.len() as u64) + .context("committed message count overflow")?; + for record in checkpoint.records { + let CommittedRecord { + trajectory, + raw, + source_offset, + } = record; + drop(raw); + let _ = source_offset; + if collect_records { + records.push(trajectory); + } } - line_index += checkpoint_count; - committed_end = complete_lines[line_index - 1].1; - records.extend(checkpoint_records); } + debug_assert_eq!(committed_end, reader.committed_end()); Ok(JsonlFile { records, - original_len, - committed_end: committed_end as u64, + record_count, + original_len: reader.bytes_read(), + committed_end, }) } diff --git a/src/storage/message_log/decoder.rs b/src/storage/message_log/decoder.rs new file mode 100644 index 0000000..435d863 --- /dev/null +++ b/src/storage/message_log/decoder.rs @@ -0,0 +1,278 @@ +use std::{mem, path::PathBuf}; + +use anyhow::{Context, Result, ensure}; +use tokio::io::{AsyncBufRead, AsyncBufReadExt}; + +use crate::trajectory::TrajectoryMessage; + +use super::{MessageCheckpoint, parse_stored_line, trajectory_record}; + +/// One committed physical message record and its exact source representation. +/// +/// `raw` includes the terminating LF and remains byte-for-byte identical to the +/// durable log, so an embedded viewer or NDJSON sink never needs to serialize a +/// parsed [`TrajectoryMessage`] back into JSON or recreate record delimiters. +#[derive(Debug)] +pub(crate) struct CommittedRecord { + pub(crate) trajectory: TrajectoryMessage, + pub(crate) raw: Vec, + pub(crate) source_offset: u64, +} + +/// A complete logical commit group. No record is returned before the declared +/// checkpoint count, indexes, refs, and newline boundaries all validate. +#[derive(Debug)] +pub(crate) struct CommittedCheckpoint { + pub(crate) records: Vec, + pub(crate) committed_end: u64, +} + +#[derive(Debug)] +pub(crate) enum DecodeResult { + Checkpoint(CommittedCheckpoint), + NeedMore, +} + +#[derive(Debug)] +struct PendingCheckpoint { + metadata: MessageCheckpoint, + records: Vec, +} + +/// Incrementally validates the shared physical-record and logical-checkpoint +/// contract. At most one not-yet-committed checkpoint is retained. +#[derive(Debug)] +pub(crate) struct CheckpointDecoder { + next_seq: u64, + committed_end: u64, + next_line_offset: u64, + pending: Option, +} + +impl Default for CheckpointDecoder { + fn default() -> Self { + Self::new(1, 0) + } +} + +impl CheckpointDecoder { + pub(crate) fn new(next_seq: u64, committed_end: u64) -> Self { + Self { + next_seq, + committed_end, + next_line_offset: committed_end, + pending: None, + } + } + + pub(crate) fn committed_end(&self) -> u64 { + self.committed_end + } + + pub(crate) fn push_complete_line( + &mut self, + path: &std::path::Path, + line_with_newline: Vec, + line_end: u64, + ) -> Result { + ensure!( + line_with_newline.ends_with(b"\n"), + "checkpoint decoder requires a newline-terminated record" + ); + let line_len = u64::try_from(line_with_newline.len()) + .context("message record length does not fit in u64")?; + let source_offset = line_end + .checked_sub(line_len) + .context("message record offset precedes the start of the file")?; + ensure!( + source_offset == self.next_line_offset, + "message record starts at byte {source_offset}, expected {}", + self.next_line_offset + ); + let stored = parse_stored_line(path, &line_with_newline)?; + let checkpoint = stored.local.checkpoint.clone(); + let expected_index = self + .pending + .as_ref() + .map_or(0, |pending| pending.records.len() as u64); + if let Some(pending) = self.pending.as_ref() { + let actual = checkpoint.as_ref().with_context(|| { + format!( + "message `{}` is missing checkpoint metadata inside group `{}`", + stored.message_ref, pending.metadata.first_message_ref + ) + })?; + ensure!( + actual.first_message_ref == pending.metadata.first_message_ref + && actual.count == pending.metadata.count + && actual.index == expected_index, + "message `{}` has inconsistent checkpoint metadata", + stored.message_ref + ); + } else if let Some(checkpoint) = checkpoint.as_ref() { + ensure!( + checkpoint.count > 0, + "message checkpoint count must be positive" + ); + ensure!( + checkpoint.index == 0, + "message checkpoint `{}` starts at index {} instead of 0", + checkpoint.first_message_ref, + checkpoint.index + ); + ensure!( + checkpoint.first_message_ref == stored.message_ref, + "message checkpoint `{}` starts with message `{}`", + checkpoint.first_message_ref, + stored.message_ref + ); + usize::try_from(checkpoint.count) + .context("message checkpoint count does not fit in memory")?; + self.next_seq + .checked_add(checkpoint.count) + .context("message sequence overflow after checkpoint")?; + } + let expected_seq = self + .next_seq + .checked_add(expected_index) + .context("message sequence overflow inside checkpoint")?; + let trajectory = trajectory_record(stored, expected_seq)?; + let record = CommittedRecord { + trajectory, + raw: line_with_newline, + source_offset, + }; + + if let Some(pending) = self.pending.as_mut() { + self.next_line_offset = line_end; + pending.records.push(record); + if expected_index + 1 < pending.metadata.count { + return Ok(DecodeResult::NeedMore); + } + return self.commit_pending(line_end); + } + + let Some(checkpoint) = checkpoint else { + // Pre-checkpoint logs are singleton records. New appends always + // persist explicit count=1 checkpoint metadata. + self.next_seq = self + .next_seq + .checked_add(1) + .context("message sequence overflow after singleton checkpoint")?; + self.committed_end = line_end; + self.next_line_offset = line_end; + return Ok(DecodeResult::Checkpoint(CommittedCheckpoint { + records: vec![record], + committed_end: line_end, + })); + }; + let count = checkpoint.count; + self.next_line_offset = line_end; + self.pending = Some(PendingCheckpoint { + metadata: checkpoint, + records: vec![record], + }); + if count == 1 { + self.commit_pending(line_end) + } else { + Ok(DecodeResult::NeedMore) + } + } + + fn commit_pending(&mut self, committed_end: u64) -> Result { + let pending = self + .pending + .as_ref() + .context("checkpoint decoder has no pending checkpoint to commit")?; + ensure!( + pending.records.len() as u64 == pending.metadata.count, + "checkpoint decoder committed an incomplete checkpoint" + ); + let next_seq = self + .next_seq + .checked_add(pending.metadata.count) + .context("message sequence overflow after checkpoint")?; + let pending = self + .pending + .take() + .context("checkpoint decoder lost its pending checkpoint")?; + self.next_seq = next_seq; + self.committed_end = committed_end; + Ok(DecodeResult::Checkpoint(CommittedCheckpoint { + records: pending.records, + committed_end, + })) + } +} + +/// Reads newline-delimited physical records lazily and feeds them to the +/// checkpoint decoder. A partial final line remains buffered so the same reader +/// can continue after a live writer appends its remainder. +pub(crate) struct CommittedCheckpointReader { + reader: R, + path: PathBuf, + decoder: CheckpointDecoder, + partial_line: Vec, + bytes_read: u64, +} + +impl CommittedCheckpointReader { + pub(crate) fn new(reader: R, path: PathBuf) -> Self { + Self::with_position(reader, path, 1, 0) + } + + pub(crate) fn with_position( + reader: R, + path: PathBuf, + next_seq: u64, + source_offset: u64, + ) -> Self { + Self { + reader, + path, + decoder: CheckpointDecoder::new(next_seq, source_offset), + partial_line: Vec::new(), + bytes_read: source_offset, + } + } + + pub(crate) fn bytes_read(&self) -> u64 { + self.bytes_read + } + + pub(crate) fn committed_end(&self) -> u64 { + self.decoder.committed_end() + } +} + +impl CommittedCheckpointReader { + pub(crate) async fn next_checkpoint(&mut self) -> Result { + loop { + let read = self + .reader + .read_until(b'\n', &mut self.partial_line) + .await + .with_context(|| format!("read initialized message log {}", self.path.display()))?; + self.bytes_read = self + .bytes_read + .checked_add(read as u64) + .context("message log byte offset overflow")?; + if !self.partial_line.ends_with(b"\n") { + return Ok(DecodeResult::NeedMore); + } + let line = mem::take(&mut self.partial_line); + match self + .decoder + .push_complete_line(&self.path, line, self.bytes_read)? + { + DecodeResult::Checkpoint(checkpoint) => { + return Ok(DecodeResult::Checkpoint(checkpoint)); + } + DecodeResult::NeedMore => {} + } + } + } +} + +#[cfg(test)] +mod tests; diff --git a/src/storage/message_log/decoder/tests.rs b/src/storage/message_log/decoder/tests.rs new file mode 100644 index 0000000..5c43b6c --- /dev/null +++ b/src/storage/message_log/decoder/tests.rs @@ -0,0 +1,544 @@ +use std::{ + io::Cursor, + path::{Path, PathBuf}, + pin::Pin, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + task::{Context, Poll}, +}; + +use serde_json::json; +use tokio::io::{AsyncRead, BufReader, ReadBuf}; + +use crate::model::MessageContent; + +use super::*; + +fn line(message_ref: &str, index: u64, count: u64, text: &str) -> Vec { + let first = message_ref + .strip_prefix('m') + .and_then(|seq| seq.parse::().ok()) + .unwrap() + .saturating_sub(index); + let mut bytes = serde_json::to_vec(&json!({ + "ref": message_ref, + "created_at": "2026-07-22T00:00:00Z", + "role": "user", + "content": [{"type": "text", "text": text}], + "_fiasco": {"checkpoint": { + "first_message_ref": format!("m{first}"), + "index": index, + "count": count, + }}, + })) + .unwrap(); + bytes.push(b'\n'); + bytes +} + +async fn collect_visible(bytes: Vec) -> Vec { + let mut reader = CommittedCheckpointReader::new( + BufReader::with_capacity(31, Cursor::new(bytes)), + PathBuf::from("messages.jsonl"), + ); + let mut refs = Vec::new(); + loop { + match reader.next_checkpoint().await.unwrap() { + DecodeResult::Checkpoint(checkpoint) => refs.extend( + checkpoint + .records + .into_iter() + .map(|record| record.trajectory.message_ref), + ), + DecodeResult::NeedMore => return refs, + } + } +} + +#[tokio::test] +async fn every_torn_tail_cut_hides_the_incomplete_checkpoint() { + let prefix = line("m1", 0, 1, "prefix"); + let group = [ + line("m2", 0, 3, "assistant"), + line("m3", 1, 3, "result one"), + line("m4", 2, 3, "result two"), + ] + .concat(); + let complete = [prefix.as_slice(), group.as_slice()].concat(); + + for cut in 0..=complete.len() { + let visible = collect_visible(complete[..cut].to_vec()).await; + let expected = if cut < prefix.len() { + Vec::new() + } else if cut < complete.len() { + vec!["m1".to_owned()] + } else { + vec![ + "m1".to_owned(), + "m2".to_owned(), + "m3".to_owned(), + "m4".to_owned(), + ] + }; + assert_eq!(visible, expected, "unexpected visible prefix at cut {cut}"); + } +} + +#[test] +fn decoder_buffers_only_the_current_checkpoint_and_preserves_raw_line() { + let first = line("m1", 0, 2, "first"); + let second = line("m2", 1, 2, "second"); + let mut decoder = CheckpointDecoder::default(); + + assert!(matches!( + decoder + .push_complete_line( + Path::new("messages.jsonl"), + first.clone(), + first.len() as u64 + ) + .unwrap(), + DecodeResult::NeedMore + )); + let pending = decoder.pending.as_ref().unwrap(); + assert_eq!(pending.records.len(), 1); + assert_eq!(pending.records[0].raw, first); + assert_eq!(pending.records[0].source_offset, 0); + + let end = (first.len() + second.len()) as u64; + let DecodeResult::Checkpoint(checkpoint) = decoder + .push_complete_line(Path::new("messages.jsonl"), second.clone(), end) + .unwrap() + else { + panic!("second record should complete the checkpoint"); + }; + assert!(decoder.pending.is_none()); + assert_eq!(checkpoint.records.len(), 2); + assert_eq!(checkpoint.records[1].raw, second); + assert_eq!(checkpoint.records[1].source_offset, first.len() as u64); + assert_eq!(checkpoint.committed_end, end); +} + +#[test] +fn decoder_does_not_reserve_the_declared_checkpoint_count() { + let first = line("m1", 0, 1_000_000, "first"); + let mut decoder = CheckpointDecoder::default(); + + assert!(matches!( + decoder + .push_complete_line( + Path::new("messages.jsonl"), + first.clone(), + first.len() as u64, + ) + .unwrap(), + DecodeResult::NeedMore + )); + let pending = decoder.pending.as_ref().unwrap(); + assert_eq!(pending.records.len(), 1); + assert!(pending.records.capacity() < 16); +} + +#[test] +fn sequence_overflow_rejects_checkpoint_before_advancing_decoder_state() { + let initial_seq = u64::MAX - 1; + let initial_offset = 97_u64; + let overflowing = line(&format!("m{initial_seq}"), 0, 2, "first"); + let overflowing_end = initial_offset + overflowing.len() as u64; + let mut decoder = CheckpointDecoder::new(initial_seq, initial_offset); + + let error = decoder + .push_complete_line(Path::new("messages.jsonl"), overflowing, overflowing_end) + .unwrap_err(); + assert!( + error + .to_string() + .contains("message sequence overflow after checkpoint") + ); + assert_eq!(decoder.next_seq, initial_seq); + assert_eq!(decoder.committed_end, initial_offset); + assert_eq!(decoder.next_line_offset, initial_offset); + assert!(decoder.pending.is_none()); + + let valid = line(&format!("m{initial_seq}"), 0, 1, "retry"); + let valid_end = initial_offset + valid.len() as u64; + let DecodeResult::Checkpoint(checkpoint) = decoder + .push_complete_line(Path::new("messages.jsonl"), valid, valid_end) + .unwrap() + else { + panic!("decoder should remain reusable after rejecting the overflow"); + }; + assert_eq!(checkpoint.committed_end, valid_end); + assert_eq!(decoder.next_seq, u64::MAX); +} + +#[test] +fn decoder_validates_a_candidate_group_from_a_tail_offset() { + let first = line("m900", 0, 2, "tail call"); + let second = line("m901", 1, 2, "tail result"); + let source_offset = 8_000_000_u64; + let mut decoder = CheckpointDecoder::new(900, source_offset); + + assert!(matches!( + decoder + .push_complete_line( + Path::new("messages.jsonl"), + first.clone(), + source_offset + first.len() as u64, + ) + .unwrap(), + DecodeResult::NeedMore + )); + let end = source_offset + (first.len() + second.len()) as u64; + let DecodeResult::Checkpoint(checkpoint) = decoder + .push_complete_line(Path::new("messages.jsonl"), second, end) + .unwrap() + else { + panic!("candidate tail group should validate independently"); + }; + assert_eq!( + checkpoint + .records + .iter() + .map(|record| record.trajectory.message_ref.as_str()) + .collect::>(), + ["m900", "m901"] + ); + assert_eq!(checkpoint.records[0].source_offset, source_offset); + assert_eq!(checkpoint.committed_end, end); +} + +#[test] +fn malformed_or_inconsistent_completed_lines_fail_without_publishing_the_group() { + let first = line("m1", 0, 2, "first"); + let mut decoder = CheckpointDecoder::default(); + assert!(matches!( + decoder + .push_complete_line( + Path::new("messages.jsonl"), + first.clone(), + first.len() as u64, + ) + .unwrap(), + DecodeResult::NeedMore + )); + + let error = decoder + .push_complete_line( + Path::new("messages.jsonl"), + b"{not-json}\n".to_vec(), + first.len() as u64 + 11, + ) + .unwrap_err(); + assert!(format!("{error:#}").contains("parse completed message")); + assert_eq!(decoder.pending.as_ref().unwrap().records.len(), 1); + + let inconsistent = line("m2", 0, 1, "wrong checkpoint"); + let error = decoder + .push_complete_line( + Path::new("messages.jsonl"), + inconsistent.clone(), + (first.len() + inconsistent.len()) as u64, + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("inconsistent checkpoint metadata") + ); + assert_eq!(decoder.pending.as_ref().unwrap().records.len(), 1); +} + +#[test] +fn committed_record_keeps_exact_json_and_tool_argument_string() { + let arguments = "{\n \"cmd\": \"printf 'a b'\"\n}"; + let mut raw = serde_json::to_vec(&json!({ + "ref": "m1", + "created_at": "2026-07-22T00:00:00Z", + "role": "assistant", + "content": [{ + "type": "tool_call", + "id": "call_1", + "name": "bash", + "arguments": arguments, + }], + "_fiasco": {"checkpoint": { + "first_message_ref": "m1", + "index": 0, + "count": 1, + }}, + })) + .unwrap(); + raw.push(b'\n'); + let expected_raw = raw.clone(); + let mut decoder = CheckpointDecoder::default(); + + let DecodeResult::Checkpoint(checkpoint) = decoder + .push_complete_line(Path::new("messages.jsonl"), raw.clone(), raw.len() as u64) + .unwrap() + else { + panic!("singleton checkpoint should commit"); + }; + assert_eq!(checkpoint.records[0].raw, expected_raw); + let MessageContent::ToolCall { + arguments: decoded, .. + } = &checkpoint.records[0].trajectory.message.content[0] + else { + panic!("stored content should remain a tool call"); + }; + assert_eq!(decoded.as_raw(), arguments); +} + +struct CountingReader { + inner: R, + reads: Arc, +} + +struct DeterministicChunkReader { + bytes: Vec, + offset: usize, + step: usize, +} + +impl AsyncRead for DeterministicChunkReader { + fn poll_read( + mut self: Pin<&mut Self>, + _context: &mut Context<'_>, + buffer: &mut ReadBuf<'_>, + ) -> Poll> { + if self.offset == self.bytes.len() { + return Poll::Ready(Ok(())); + } + let chunk = (self.step.wrapping_mul(17).wrapping_add(5) % 23) + 1; + self.step = self.step.wrapping_add(1); + let end = self + .offset + .saturating_add(chunk.min(buffer.remaining())) + .min(self.bytes.len()); + buffer.put_slice(&self.bytes[self.offset..end]); + self.offset = end; + Poll::Ready(Ok(())) + } +} + +#[tokio::test] +async fn deterministic_chunk_boundaries_preserve_checkpoint_and_raw_order() { + let expected = [ + line("m1", 0, 1, "prefix"), + line("m2", 0, 3, "assistant"), + line("m3", 1, 3, "first result"), + line("m4", 2, 3, "second result"), + line("m5", 0, 1, "suffix"), + ] + .concat(); + let source = DeterministicChunkReader { + bytes: expected.clone(), + offset: 0, + step: 0, + }; + let mut reader = CommittedCheckpointReader::new( + BufReader::with_capacity(29, source), + PathBuf::from("messages.jsonl"), + ); + let mut refs = Vec::new(); + let mut raw = Vec::new(); + while let DecodeResult::Checkpoint(checkpoint) = reader.next_checkpoint().await.unwrap() { + for record in checkpoint.records { + refs.push(record.trajectory.message_ref); + raw.extend(record.raw); + } + } + + assert_eq!(refs, ["m1", "m2", "m3", "m4", "m5"]); + assert_eq!(raw, expected); +} + +impl AsyncRead for CountingReader { + fn poll_read( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &mut ReadBuf<'_>, + ) -> Poll> { + let before = buffer.filled().len(); + let poll = Pin::new(&mut self.inner).poll_read(context, buffer); + if let Poll::Ready(Ok(())) = &poll { + self.reads + .fetch_add(buffer.filled().len() - before, Ordering::Relaxed); + } + poll + } +} + +#[tokio::test] +async fn first_checkpoint_does_not_read_the_rest_of_a_large_log() { + let mut bytes = Vec::new(); + for seq in 1..=20_000_u64 { + bytes.extend(line(&format!("m{seq}"), 0, 1, "payload")); + } + let total = bytes.len(); + let reads = Arc::new(AtomicUsize::new(0)); + let source = CountingReader { + inner: Cursor::new(bytes), + reads: reads.clone(), + }; + let mut reader = CommittedCheckpointReader::new( + BufReader::with_capacity(256, source), + PathBuf::from("messages.jsonl"), + ); + + let DecodeResult::Checkpoint(first) = reader.next_checkpoint().await.unwrap() else { + panic!("first complete line should produce a checkpoint"); + }; + assert_eq!(first.records[0].trajectory.message_ref, "m1"); + assert!(reads.load(Ordering::Relaxed) < total / 100); + assert!(reader.bytes_read() < (total / 100) as u64); +} + +#[tokio::test] +async fn partial_line_can_complete_after_a_later_read() { + let complete = line("m1", 0, 1, "later"); + let split = complete.len() / 2; + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("messages.jsonl"); + tokio::fs::write(&path, &complete[..split]).await.unwrap(); + let file = tokio::fs::File::open(&path).await.unwrap(); + let mut reader = CommittedCheckpointReader::new(BufReader::new(file), path.clone()); + + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); + use tokio::io::AsyncWriteExt; + let mut writer = tokio::fs::OpenOptions::new() + .append(true) + .open(&path) + .await + .unwrap(); + writer.write_all(&complete[split..]).await.unwrap(); + writer.flush().await.unwrap(); + let DecodeResult::Checkpoint(checkpoint) = reader.next_checkpoint().await.unwrap() else { + panic!("completed line should become visible"); + }; + assert_eq!(checkpoint.records[0].raw, complete); +} + +#[tokio::test] +async fn partial_multi_record_checkpoint_commits_once_after_multiple_eofs() { + let prefix = line("m1", 0, 1, "prefix"); + let first = line("m2", 0, 3, "assistant"); + let second = line("m3", 1, 3, "first result"); + let third = line("m4", 2, 3, "second result"); + let second_split = second.len() / 2; + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("messages.jsonl"); + let initial = [prefix.as_slice(), first.as_slice(), &second[..second_split]].concat(); + tokio::fs::write(&path, &initial).await.unwrap(); + let file = tokio::fs::File::open(&path).await.unwrap(); + let mut reader = CommittedCheckpointReader::new(BufReader::new(file), path.clone()); + + let DecodeResult::Checkpoint(committed_prefix) = reader.next_checkpoint().await.unwrap() else { + panic!("the singleton prefix should be committed"); + }; + assert_eq!(committed_prefix.records[0].trajectory.message_ref, "m1"); + assert_eq!(reader.committed_end(), prefix.len() as u64); + + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); + assert_eq!(reader.committed_end(), prefix.len() as u64); + + use tokio::io::AsyncWriteExt; + let mut writer = tokio::fs::OpenOptions::new() + .append(true) + .open(&path) + .await + .unwrap(); + writer.write_all(&second[second_split..]).await.unwrap(); + writer.flush().await.unwrap(); + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); + assert_eq!(reader.committed_end(), prefix.len() as u64); + + writer.write_all(&third).await.unwrap(); + writer.flush().await.unwrap(); + let DecodeResult::Checkpoint(checkpoint) = reader.next_checkpoint().await.unwrap() else { + panic!("the completed group should commit after the final append"); + }; + assert_eq!( + checkpoint + .records + .iter() + .map(|record| record.trajectory.message_ref.as_str()) + .collect::>(), + ["m2", "m3", "m4"] + ); + assert_eq!(checkpoint.records[0].source_offset, prefix.len() as u64); + assert_eq!( + checkpoint.committed_end, + (prefix.len() + first.len() + second.len() + third.len()) as u64 + ); + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); +} + +#[test] +fn decoder_rejects_non_contiguous_and_underflowing_line_offsets() { + let first = line("m1", 0, 1, "first"); + let mut decoder = CheckpointDecoder::default(); + let error = decoder + .push_complete_line( + Path::new("messages.jsonl"), + first.clone(), + first.len() as u64 + 7, + ) + .unwrap_err(); + assert!(error.to_string().contains("starts at byte 7, expected 0")); + + let error = decoder + .push_complete_line(Path::new("messages.jsonl"), first.clone(), 1) + .unwrap_err(); + assert!( + error + .to_string() + .contains("record offset precedes the start of the file") + ); + + let DecodeResult::Checkpoint(_) = decoder + .push_complete_line( + Path::new("messages.jsonl"), + first.clone(), + first.len() as u64, + ) + .unwrap() + else { + panic!("the contiguous first line should commit"); + }; + let second = line("m2", 0, 1, "second"); + let overlap_end = first.len() as u64 + second.len() as u64 - 1; + let error = decoder + .push_complete_line(Path::new("messages.jsonl"), second, overlap_end) + .unwrap_err(); + assert_eq!( + error.to_string(), + format!( + "message record starts at byte {}, expected {}", + first.len() - 1, + first.len() + ) + ); +}