From 78dc2503840d05f4411ae184bb5560ef327fce3c Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 14 Oct 2026 10:39:18 +0000 Subject: [PATCH] Stop rollout reads when the calling task is cancelled (#53947) ## Why Blocking rollout reads can continue decompressing and scanning history after the calling async task is dropped. ## What changed Propagate caller cancellation through rollout readers, prefix validation, model-context scans, history materialization, and working-directory lookup. Check cancellation during compressed input and output, bounded reads, seeks, and between history records. Return an error on cancellation so partial history cannot be mistaken for a complete read. ## Testing Add regression coverage for dropping a caller during decompression, cancelling compressed input and output (including skippable frames), cancelling reads and seeks on plain and decoded snapshots, and cancelling history materialization between records. GitOrigin-RevId: e3dac952789b53402f2a5f17c2a2232506feb9d6 --- codex-rs/exec/src/lib.rs | 10 +- codex-rs/memories/write/src/rollout_input.rs | 2 +- codex-rs/rollout/src/compression_tests.rs | 2 +- codex-rs/rollout/src/seekable_reader.rs | 86 +++++++++++++-- codex-rs/rollout/src/seekable_reader_tests.rs | 100 ++++++++++++++++-- .../thread-store/src/local/model_context.rs | 18 +++- .../thread-store/src/local/rollout_lineage.rs | 4 +- .../local/thread_history_materialization.rs | 27 +++-- .../thread_history_materialization_tests.rs | 15 +++ codex-rs/tui/src/worktree_startup.rs | 10 +- 10 files changed, 236 insertions(+), 38 deletions(-) diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index f97ce42737..0a40529c82 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -1826,8 +1826,10 @@ async fn latest_thread_cwd(thread: &AppServerThread) -> PathBuf { async fn parse_latest_turn_context_cwd(path: &Path) -> Option { let path = path.to_path_buf(); - tokio::task::spawn_blocking(move || { - let reader = codex_rollout::open_rollout_seekable_reader(&path).ok()?; + let (stop, keep_running) = tokio::sync::oneshot::channel::<()>(); + let result = tokio::task::spawn_blocking(move || { + let reader = + codex_rollout::open_rollout_seekable_reader(&path, || stop.is_closed()).ok()?; let mut scanner = codex_rollout::ReverseJsonlScanner::new(reader).ok()?; while let Some(outcome) = scanner.scan_next_rollout_line().ok()? { if let codex_rollout::ScanOutcome::Parsed(RolloutLine { @@ -1842,7 +1844,9 @@ async fn parse_latest_turn_context_cwd(path: &Path) -> Option { }) .await .ok() - .flatten() + .flatten(); + drop(keep_running); + result } fn cwds_match(current_cwd: &Path, session_cwd: &Path) -> bool { diff --git a/codex-rs/memories/write/src/rollout_input.rs b/codex-rs/memories/write/src/rollout_input.rs index 65b4461e85..e0f9e643d1 100644 --- a/codex-rs/memories/write/src/rollout_input.rs +++ b/codex-rs/memories/write/src/rollout_input.rs @@ -36,7 +36,7 @@ pub(crate) async fn read_recent_rollout_items(path: &Path) -> anyhow::Result(); let task = tokio::task::spawn_blocking(move || -> anyhow::Result> { - let mut reader = open_rollout_seekable_reader(&path)?; + let mut reader = open_rollout_seekable_reader(&path, || stop.is_closed())?; // Validate this session's format from its header, not metadata copied into fork history. // Bound the header read too: session metadata may contain large instructions or schemas. let mut header = String::new(); diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index a4171801dd..40ec1ceb64 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -519,7 +519,7 @@ async fn worker_compresses_archived_fork_chain() -> anyhow::Result<()> { ] { assert!(!path.exists(), "compression should replace {path:?}"); let mut restored = Vec::new(); - crate::open_rollout_seekable_reader(path)?.read_to_end(&mut restored)?; + crate::open_rollout_seekable_reader(path, || false)?.read_to_end(&mut restored)?; assert_eq!(restored, original); } Ok(()) diff --git a/codex-rs/rollout/src/seekable_reader.rs b/codex-rs/rollout/src/seekable_reader.rs index 17b7e4a022..a72617aeaa 100644 --- a/codex-rs/rollout/src/seekable_reader.rs +++ b/codex-rs/rollout/src/seekable_reader.rs @@ -2,11 +2,13 @@ //! //! Offsets always address the original JSONL bytes. Readers retain an open file, or an anonymous //! decoded snapshot, so a concurrent compression cannot invalidate an in-progress scan. +//! Cancellation is checked during decompression and subsequent bounded reads and seeks. use std::fs::File; use std::io; use std::io::Read; use std::io::Seek; +use std::io::SeekFrom; use std::path::Path; use std::time::Duration; @@ -67,15 +69,68 @@ pub fn read_rollout_prefix(path: &Path, max_bytes: usize) -> io::Result io::Result { - match RolloutReader::open(path)? { - RolloutReader::Plain(file) => Ok(file), +/// `is_cancelled` must observe the lifetime of the calling task. Cancellation returns an error, +/// never EOF, so callers cannot mistake a partial read for a complete history. +pub fn open_rollout_seekable_reader( + path: &Path, + is_cancelled: impl Fn() -> bool, +) -> io::Result { + if is_cancelled() { + return Err(io::Error::other("rollout read cancelled")); + } + let file = match RolloutReader::open(path)? { + RolloutReader::Plain(file) => file, RolloutReader::Compressed(file) => { let mut decoded = tempfile::tempfile()?; - io::copy(&mut zstd::stream::read::Decoder::new(file)?, &mut decoded)?; + // Empty and skippable frames can consume input without producing decoded bytes. + let file = CancellableReader { + inner: file, + is_cancelled: &is_cancelled, + }; + let mut reader = CancellableReader { + inner: zstd::stream::read::Decoder::new(file)?, + is_cancelled: &is_cancelled, + }; + io::copy(&mut reader, &mut decoded)?; decoded.rewind()?; - Ok(decoded) + decoded } + }; + Ok(CancellableReader { + inner: file, + is_cancelled, + }) +} + +struct CancellableReader { + inner: R, + is_cancelled: F, +} + +impl bool> CancellableReader { + fn check_cancelled(&self) -> io::Result<()> { + if (self.is_cancelled)() { + // Interrupted is retried by io::copy and read_exact, defeating cancellation. + return Err(io::Error::other("rollout read cancelled")); + } + Ok(()) + } +} + +impl bool> Read for CancellableReader { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + self.check_cancelled()?; + let limit = buffer.len().min(64 * 1024); + let read = self.inner.read(&mut buffer[..limit])?; + self.check_cancelled()?; + Ok(read) + } +} + +impl bool> Seek for CancellableReader { + fn seek(&mut self, position: SeekFrom) -> io::Result { + self.check_cancelled()?; + self.inner.seek(position) } } @@ -84,10 +139,21 @@ pub fn open_rollout_seekable_reader(path: &Path) -> io::Result { /// A known first-frame size is a lower bound even for concatenated zstd frames. New compressed /// rollouts include that size, so ordinary lineage validation only reads the header. Older frames /// without a size, and prefixes extending beyond the first frame, are decoded up to the bound. -pub fn rollout_contains_prefix(path: &Path, end_byte_offset: u64) -> io::Result { +pub fn rollout_contains_prefix( + path: &Path, + end_byte_offset: u64, + is_cancelled: impl Fn() -> bool, +) -> io::Result { + if is_cancelled() { + return Err(io::Error::other("rollout read cancelled")); + } match RolloutReader::open(path)? { RolloutReader::Plain(file) => Ok(end_byte_offset <= file.metadata()?.len()), - RolloutReader::Compressed(mut file) => { + RolloutReader::Compressed(file) => { + let mut file = CancellableReader { + inner: file, + is_cancelled: &is_cancelled, + }; // A zstd frame header occupies at most 18 bytes. let mut header = [0; 18]; let read = file.read(&mut header)?; @@ -99,7 +165,11 @@ pub fn rollout_contains_prefix(path: &Path, end_byte_offset: u64) -> io::Result< return Ok(true); } file.rewind()?; - let mut prefix = zstd::stream::read::Decoder::new(file)?.take(end_byte_offset); + let mut prefix = CancellableReader { + inner: zstd::stream::read::Decoder::new(file)?, + is_cancelled: &is_cancelled, + } + .take(end_byte_offset); Ok(io::copy(&mut prefix, &mut io::sink())? == end_byte_offset) } } diff --git a/codex-rs/rollout/src/seekable_reader_tests.rs b/codex-rs/rollout/src/seekable_reader_tests.rs index abc2866490..5156598810 100644 --- a/codex-rs/rollout/src/seekable_reader_tests.rs +++ b/codex-rs/rollout/src/seekable_reader_tests.rs @@ -1,3 +1,4 @@ +use std::cell::Cell; use std::fs; use std::io::Read; use std::io::Seek; @@ -8,8 +9,78 @@ use pretty_assertions::assert_eq; use super::*; +#[tokio::test] +async fn dropping_the_caller_stops_compressed_materialization() -> anyhow::Result<()> { + let home = tempfile::tempdir()?; + let path = home.path().join("rollout.jsonl.zst"); + let original = vec![b'x'; 1024 * 1024]; + let compressed = zstd::stream::encode_all(original.as_slice(), /*level*/ 3)?; + fs::write(&path, &compressed)?; + + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (resume_tx, resume_rx) = std::sync::mpsc::channel(); + let (finished_tx, finished_rx) = tokio::sync::oneshot::channel(); + let read_path = path.clone(); + let task = tokio::spawn(async move { + let (stop, keep_running) = tokio::sync::oneshot::channel::<()>(); + let worker = tokio::task::spawn_blocking(move || { + let checks = Cell::new(0); + let started = Cell::new(Some(started_tx)); + let result = open_rollout_seekable_reader(&read_path, || { + checks.set(checks.get() + 1); + // Check on open and around the first source read and decoded chunk. + // Pause after that chunk so a check only before decompression cannot pass. + if checks.get() >= 5 + && let Some(started) = started.take() + { + let _ = started.send(()); + let _ = resume_rx.recv(); + } + stop.is_closed() + }); + let _ = finished_tx.send(result.err().map(|error| error.kind())); + }); + let result = worker.await; + drop(keep_running); + result + }); + + started_rx.await?; + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + resume_tx.send(())?; + assert_eq!(finished_rx.await?, Some(io::ErrorKind::Other)); + assert_eq!(fs::read(&path)?, compressed); + assert!(!path.with_extension("").exists()); + Ok(()) +} + #[test] -fn compressed_offsets_preserve_exact_bytes_and_open_snapshots() -> io::Result<()> { +fn cancellation_stops_compressed_input_and_output() -> io::Result<()> { + let home = tempfile::tempdir()?; + let path = home.path().join("rollout.jsonl.zst"); + let original = vec![b'x'; 1024 * 1024]; + let mut skippable = 0x184D2A50_u32.to_le_bytes().to_vec(); + skippable.extend((original.len() as u32).to_le_bytes()); + skippable.extend(&original); + // Exercise many input reads with no output, and many output reads from little input. + for bytes in [ + skippable, + zstd::stream::encode_all(original.as_slice(), /*level*/ 3)?, + ] { + fs::write(&path, bytes)?; + drop(open_rollout_seekable_reader(&path, || false)?); + let checks = Cell::new(0); + let cancelled = || checks.replace(checks.get() + 1) >= 10; + assert!(open_rollout_seekable_reader(&path, cancelled).is_err()); + checks.set(0); + assert!(rollout_contains_prefix(&path, original.len() as u64, cancelled).is_err()); + } + Ok(()) +} + +#[test] +fn read_snapshots_preserve_bytes_and_report_cancellation() -> io::Result<()> { let home = tempfile::tempdir()?; let path = home.path().join("rollout.jsonl"); let compressed = path.with_extension("jsonl.zst"); @@ -19,7 +90,9 @@ fn compressed_offsets_preserve_exact_bytes_and_open_snapshots() -> io::Result<() zstd::stream::encode_all(original, /*level*/ 3)?, )?; - let mut reader = open_rollout_seekable_reader(&path)?; + let cancelled = Cell::new(false); + let is_cancelled = || cancelled.get(); + let mut reader = open_rollout_seekable_reader(&path, is_cancelled)?; // The retained snapshot must remain valid after another caller materializes the logical path. fs::write(&path, b"replacement\n")?; fs::remove_file(&compressed)?; @@ -29,10 +102,10 @@ fn compressed_offsets_preserve_exact_bytes_and_open_snapshots() -> io::Result<() assert_eq!(suffix, original[7..]); let mut current = Vec::new(); - open_rollout_seekable_reader(&compressed)?.read_to_end(&mut current)?; + open_rollout_seekable_reader(&compressed, || false)?.read_to_end(&mut current)?; assert_eq!(current, b"replacement\n"); - let mut plain_snapshot = open_rollout_seekable_reader(&path)?; + let mut plain_snapshot = open_rollout_seekable_reader(&path, is_cancelled)?; fs::write( &compressed, zstd::stream::encode_all(current.as_slice(), /*level*/ 3)?, @@ -41,6 +114,16 @@ fn compressed_offsets_preserve_exact_bytes_and_open_snapshots() -> io::Result<() let mut retained = Vec::new(); plain_snapshot.read_to_end(&mut retained)?; assert_eq!(retained, current); + reader.rewind()?; + plain_snapshot.rewind()?; + cancelled.set(true); + for mut reader in [reader, plain_snapshot] { + assert_eq!( + reader.read_to_end(&mut Vec::new()).unwrap_err().kind(), + io::ErrorKind::Other + ); + assert_eq!(reader.rewind().unwrap_err().kind(), io::ErrorKind::Other); + } Ok(()) } @@ -67,11 +150,12 @@ fn prefix_bounds_support_sized_unsized_and_concatenated_frames() -> io::Result<( (concatenated, [first, second].concat()), ] { fs::write(&path, compressed)?; - assert!(rollout_contains_prefix(&path, first.len() as u64)?); - assert!(rollout_contains_prefix(&path, expected.len() as u64)?); - assert!(!rollout_contains_prefix(&path, expected.len() as u64 + 1)?); + let contains = |end| rollout_contains_prefix(&path, end, || false); + assert!(contains(first.len() as u64)?); + assert!(contains(expected.len() as u64)?); + assert!(!contains(expected.len() as u64 + 1)?); let mut actual = Vec::new(); - open_rollout_seekable_reader(&path)?.read_to_end(&mut actual)?; + open_rollout_seekable_reader(&path, || false)?.read_to_end(&mut actual)?; assert_eq!(actual, expected); assert!(!path.with_extension("").exists()); } diff --git a/codex-rs/thread-store/src/local/model_context.rs b/codex-rs/thread-store/src/local/model_context.rs index 13a81b71ad..bcfd460b20 100644 --- a/codex-rs/thread-store/src/local/model_context.rs +++ b/codex-rs/thread-store/src/local/model_context.rs @@ -115,10 +115,13 @@ pub(super) async fn load_for_fork( // Recover only the runtime version before applying the fork cutoff. Stop at the // newest version-bearing context instead of retaining the source's full replay. let source_lineage = lineage.clone(); + let (stop, keep_running) = tokio::sync::oneshot::channel::<()>(); session_meta.meta.multi_agent_version = tokio::task::spawn_blocking(move || { for segment in source_lineage.segments().iter().rev() { - let file = - codex_rollout::open_rollout_seekable_reader(segment.rollout_path.as_path())?; + let file = codex_rollout::open_rollout_seekable_reader( + segment.rollout_path.as_path(), + || stop.is_closed(), + )?; let mut scanner = match segment.end.map(|end| end.end_byte_offset) { Some(end_byte_offset) => ReverseJsonlScanner::new_at(file, end_byte_offset)?, None => ReverseJsonlScanner::new(file)?, @@ -145,6 +148,7 @@ pub(super) async fn load_for_fork( .map_err(|err| ThreadStoreError::Internal { message: format!("failed to read fork runtime version: {err}"), })?; + drop(keep_running); } match history_base { Some(history_base) => { @@ -159,13 +163,15 @@ async fn scan_model_context_from_lineage( lineage: RolloutLineage, session_meta: SessionMetaLine, ) -> ThreadStoreResult> { + let (stop, keep_running) = tokio::sync::oneshot::channel::<()>(); let scan = tokio::task::spawn_blocking(move || { - scan_model_context_from_lineage_blocking(&lineage, session_meta) + scan_model_context_from_lineage_blocking(&lineage, session_meta, || stop.is_closed()) }) .await .map_err(|err| ThreadStoreError::Internal { message: format!("failed to join model context scan: {err}"), })?; + drop(keep_running); match scan { Ok(items) => Ok(items), Err(err) => Err(ThreadStoreError::Internal { @@ -177,10 +183,14 @@ async fn scan_model_context_from_lineage( fn scan_model_context_from_lineage_blocking( lineage: &RolloutLineage, session_meta: SessionMetaLine, + is_cancelled: impl Fn() -> bool, ) -> io::Result> { let mut scan = ModelContextScan::default(); 'segments: for segment in lineage.segments().iter().rev() { - let file = codex_rollout::open_rollout_seekable_reader(segment.rollout_path.as_path())?; + let file = codex_rollout::open_rollout_seekable_reader( + segment.rollout_path.as_path(), + &is_cancelled, + )?; let mut scanner = match segment.end.map(|end| end.end_byte_offset) { Some(end_byte_offset) => ReverseJsonlScanner::new_at(file, end_byte_offset)?, None => ReverseJsonlScanner::new(file)?, diff --git a/codex-rs/thread-store/src/local/rollout_lineage.rs b/codex-rs/thread-store/src/local/rollout_lineage.rs index 556ca03767..c38440a2dc 100644 --- a/codex-rs/thread-store/src/local/rollout_lineage.rs +++ b/codex-rs/thread-store/src/local/rollout_lineage.rs @@ -284,8 +284,9 @@ async fn validate_cutoff_bounds( } let path = rollout_path.to_path_buf(); let end_byte_offset = end.end_byte_offset; + let (stop, keep_running) = tokio::sync::oneshot::channel::<()>(); let contains_prefix = tokio::task::spawn_blocking(move || { - codex_rollout::rollout_contains_prefix(&path, end_byte_offset) + codex_rollout::rollout_contains_prefix(&path, end_byte_offset, || stop.is_closed()) }) .await .map_err(|err| ThreadStoreError::Internal { @@ -297,6 +298,7 @@ async fn validate_cutoff_bounds( rollout_path.display() ), })?; + drop(keep_running); if !contains_prefix { return Err(malformed_lineage( requested_thread_id, diff --git a/codex-rs/thread-store/src/local/thread_history_materialization.rs b/codex-rs/thread-store/src/local/thread_history_materialization.rs index a543302c38..c6cfa9de8e 100644 --- a/codex-rs/thread-store/src/local/thread_history_materialization.rs +++ b/codex-rs/thread-store/src/local/thread_history_materialization.rs @@ -91,7 +91,8 @@ async fn read_projection_steps( ) -> ThreadStoreResult<(Vec, u64)> { let path = rollout_path.to_path_buf(); let span = tracing::Span::current(); - tokio::task::spawn_blocking(move || { + let (stop, keep_running) = tokio::sync::oneshot::channel::<()>(); + let result = tokio::task::spawn_blocking(move || { let _entered = span.enter(); read_projection_steps_sync( &path, @@ -99,12 +100,15 @@ async fn read_projection_steps( expected_ordinal, thread_id, subagent_history_start_ordinal, + || stop.is_closed(), ) }) .await .map_err(|err| ThreadStoreError::Internal { message: format!("failed to join rollout projection read: {err}"), - })? + })?; + drop(keep_running); + result } /// Read and project the complete durable prefix in one blocking task. @@ -114,15 +118,16 @@ fn read_projection_steps_sync( expected_ordinal: u64, thread_id: ThreadId, subagent_history_start_ordinal: Option, + is_cancelled: impl Fn() -> bool, ) -> ThreadStoreResult<(Vec, u64)> { - let mut file = match codex_rollout::open_rollout_seekable_reader(rollout_path) { + let mut file = match codex_rollout::open_rollout_seekable_reader(rollout_path, &is_cancelled) { Ok(file) => file, Err(err) if err.kind() == std::io::ErrorKind::NotFound && start_offset == 0 => { return Ok((Vec::new(), 0)); } Err(err) => return Err(thread_store_io_error(err)), }; - let file_end_offset = file.metadata().map_err(thread_store_io_error)?.len(); + let file_end_offset = file.seek(SeekFrom::End(0)).map_err(thread_store_io_error)?; let byte_count = file_end_offset .checked_sub(start_offset) @@ -139,15 +144,19 @@ fn read_projection_steps_sync( .map_err(thread_store_io_error)?; // Only project the newline-terminated prefix; leave a trailing partial record for the next // pass. - let complete_byte_count = bytes - .iter() - .rposition(|byte| *byte == b'\n') - .map_or(0, |index| index + 1); let mut projections = Vec::new(); let mut next_ordinal = expected_ordinal; let mut next_offset = start_offset; let mut line_start_offset = start_offset; - for line_bytes in bytes[..complete_byte_count].split_inclusive(|byte| *byte == b'\n') { + for line_bytes in bytes.split_inclusive(|byte| *byte == b'\n') { + if is_cancelled() { + return Err(thread_store_io_error(std::io::Error::other( + "rollout read cancelled", + ))); + } + if !line_bytes.ends_with(b"\n") { + break; + } let line_end_offset = line_start_offset .checked_add(u64::try_from(line_bytes.len()).map_err(|_| { ThreadStoreError::Internal { diff --git a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs index 0d260df188..45bb268678 100644 --- a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs +++ b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs @@ -1,3 +1,4 @@ +use std::cell::Cell; use std::fs; use std::io::Write; use std::path::Path; @@ -2470,6 +2471,20 @@ async fn blank_and_rejected_rollout_lines_advance_projection() { .expect("queue valid retry"); recorder.flush().await.expect("flush valid retry"); + // Replay this small file, cancelling between records after its single read completes. + let checks = Cell::new(0); + assert!( + super::read_projection_steps_sync( + rollout_path.as_path(), + /*start_offset*/ 0, + /*expected_ordinal*/ 0, + thread_id, + /*subagent_history_start_ordinal*/ None, + || checks.replace(checks.get() + 1) >= 6, + ) + .is_err() + ); + super::materialize_to_sqlite(&store, thread_id, rollout_path.as_path()) .await .expect("project valid retry after rejected line"); diff --git a/codex-rs/tui/src/worktree_startup.rs b/codex-rs/tui/src/worktree_startup.rs index 3d194468ac..c1d1ca31ca 100644 --- a/codex-rs/tui/src/worktree_startup.rs +++ b/codex-rs/tui/src/worktree_startup.rs @@ -95,8 +95,10 @@ impl ManagedTuiWorktree { async fn latest_thread_cwd(path: Option, fallback: PathBuf) -> PathBuf { let Some(path) = path else { return fallback }; - tokio::task::spawn_blocking(move || { - let reader = codex_rollout::open_rollout_seekable_reader(&path).ok()?; + let (stop, keep_running) = tokio::sync::oneshot::channel::<()>(); + let result = tokio::task::spawn_blocking(move || { + let reader = + codex_rollout::open_rollout_seekable_reader(&path, || stop.is_closed()).ok()?; let mut scanner = codex_rollout::ReverseJsonlScanner::new(reader).ok()?; while let Some(outcome) = scanner.scan_next_rollout_line().ok()? { if let codex_rollout::ScanOutcome::Parsed(codex_rollout::RolloutLine { @@ -112,7 +114,9 @@ async fn latest_thread_cwd(path: Option, fallback: PathBuf) -> PathBuf .await .ok() .flatten() - .unwrap_or(fallback) + .unwrap_or(fallback); + drop(keep_running); + result } #[allow(clippy::too_many_arguments)]