From 58c921967aafdb6787da81fe9bfcb623f0ff5cc3 Mon Sep 17 00:00:00 2001 From: "Xinyao (Morry) Niu" Date: Wed, 22 Jul 2026 17:50:55 +0800 Subject: [PATCH 1/3] refactor(storage): decode checkpoints incrementally Stream complete physical records through one shared checkpoint decoder so trajectory loading, append recovery, and future transcript sources agree on committed groups and byte offsets. Preserve exact raw lines while buffering at most one incomplete checkpoint. Keep append recovery bounded by counting committed records without collecting the full trajectory, and cover torn tails, malformed groups, arbitrary tail offsets, exact tool arguments, and large-log lazy reads. --- docs/architecture.md | 8 + docs/source-map.md | 5 + src/storage/message_log.rs | 113 +++----- src/storage/message_log/decoder.rs | 263 +++++++++++++++++++ src/storage/message_log/decoder/tests.rs | 313 +++++++++++++++++++++++ 5 files changed, 623 insertions(+), 79 deletions(-) create mode 100644 src/storage/message_log/decoder.rs create mode 100644 src/storage/message_log/decoder/tests.rs 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..fea757a --- /dev/null +++ b/src/storage/message_log/decoder.rs @@ -0,0 +1,263 @@ +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, + 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, + 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")?; + 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 - 1) + .context("message sequence overflow inside 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() { + 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; + return Ok(DecodeResult::Checkpoint(CommittedCheckpoint { + records: vec![record], + committed_end: line_end, + })); + }; + let count = checkpoint.count; + 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 + .take() + .context("checkpoint decoder has no pending checkpoint to commit")?; + ensure!( + pending.records.len() as u64 == pending.metadata.count, + "checkpoint decoder committed an incomplete checkpoint" + ); + self.next_seq = self + .next_seq + .checked_add(pending.metadata.count) + .context("message sequence overflow after checkpoint")?; + 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..1719166 --- /dev/null +++ b/src/storage/message_log/decoder/tests.rs @@ -0,0 +1,313 @@ +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_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, +} + +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); +} From 6bed8deb54cc59ae2826cc1af60eef15a0bbadaa Mon Sep 17 00:00:00 2001 From: "Xinyao (Morry) Niu" Date: Wed, 22 Jul 2026 18:06:09 +0800 Subject: [PATCH 2/3] fix(storage): validate checkpoint source offsets Require every decoded physical line to begin exactly where the previous line ended so tail candidates cannot publish raw records with gaps or overlaps. Cover deterministic chunk boundaries, large declared counts, and multi-record checkpoints that complete across repeated EOF observations. --- src/storage/message_log/decoder.rs | 10 ++ src/storage/message_log/decoder/tests.rs | 198 +++++++++++++++++++++++ 2 files changed, 208 insertions(+) diff --git a/src/storage/message_log/decoder.rs b/src/storage/message_log/decoder.rs index fea757a..78ef3b1 100644 --- a/src/storage/message_log/decoder.rs +++ b/src/storage/message_log/decoder.rs @@ -45,6 +45,7 @@ struct PendingCheckpoint { pub(crate) struct CheckpointDecoder { next_seq: u64, committed_end: u64, + next_line_offset: u64, pending: Option, } @@ -59,6 +60,7 @@ impl CheckpointDecoder { Self { next_seq, committed_end, + next_line_offset: committed_end, pending: None, } } @@ -82,6 +84,11 @@ impl CheckpointDecoder { 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 @@ -137,6 +144,7 @@ impl CheckpointDecoder { }; 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); @@ -152,12 +160,14 @@ impl CheckpointDecoder { .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], diff --git a/src/storage/message_log/decoder/tests.rs b/src/storage/message_log/decoder/tests.rs index 1719166..fc3c7cb 100644 --- a/src/storage/message_log/decoder/tests.rs +++ b/src/storage/message_log/decoder/tests.rs @@ -121,6 +121,26 @@ fn decoder_buffers_only_the_current_checkpoint_and_preserves_raw_line() { 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 decoder_validates_a_candidate_group_from_a_tail_offset() { let first = line("m900", 0, 2, "tail call"); @@ -243,6 +263,65 @@ struct CountingReader { 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>, @@ -311,3 +390,122 @@ async fn partial_line_can_complete_after_a_later_read() { }; 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() + ) + ); +} From c9a00e09108894a4e08bd35685691a9a3a17676c Mon Sep 17 00:00:00 2001 From: "Xinyao (Morry) Niu" Date: Wed, 22 Jul 2026 18:12:21 +0800 Subject: [PATCH 3/3] fix(storage): reject checkpoint sequence overflow early Validate the sequence after a complete checkpoint before retaining its first record. This prevents an unrepresentable next sequence from surfacing only after the checkpoint has reached its commit boundary. Keep commit_pending transactional as a defensive invariant and cover the boundary with a retry test that proves pending and byte offsets remain untouched. --- src/storage/message_log/decoder.rs | 13 +++++++--- src/storage/message_log/decoder/tests.rs | 33 ++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/src/storage/message_log/decoder.rs b/src/storage/message_log/decoder.rs index 78ef3b1..435d863 100644 --- a/src/storage/message_log/decoder.rs +++ b/src/storage/message_log/decoder.rs @@ -129,8 +129,8 @@ impl CheckpointDecoder { usize::try_from(checkpoint.count) .context("message checkpoint count does not fit in memory")?; self.next_seq - .checked_add(checkpoint.count - 1) - .context("message sequence overflow inside checkpoint")?; + .checked_add(checkpoint.count) + .context("message sequence overflow after checkpoint")?; } let expected_seq = self .next_seq @@ -182,16 +182,21 @@ impl CheckpointDecoder { fn commit_pending(&mut self, committed_end: u64) -> Result { let pending = self .pending - .take() + .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" ); - self.next_seq = self + 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, diff --git a/src/storage/message_log/decoder/tests.rs b/src/storage/message_log/decoder/tests.rs index fc3c7cb..5c43b6c 100644 --- a/src/storage/message_log/decoder/tests.rs +++ b/src/storage/message_log/decoder/tests.rs @@ -141,6 +141,39 @@ fn decoder_does_not_reserve_the_declared_checkpoint_count() { 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");