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

@@ -1826,8 +1826,10 @@ async fn latest_thread_cwd(thread: &AppServerThread) -> PathBuf {
 
 async fn parse_latest_turn_context_cwd(path: &Path) -> Option<PathBuf> {
     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<PathBuf> {
     })
     .await
     .ok()
-    .flatten()
+    .flatten();
+    drop(keep_running);
+    result
 }
 
 fn cwds_match(current_cwd: &Path, session_cwd: &Path) -> bool {

codex-rs/memories/write/src/rollout_input.rs

@@ -36,7 +36,7 @@ pub(crate) async fn read_recent_rollout_items(path: &Path) -> anyhow::Result<Vec
     let path = path.to_owned();
     let (stop, keep_running) = tokio::sync::oneshot::channel::<()>();
     let task = tokio::task::spawn_blocking(move || -> anyhow::Result<Vec<RolloutItem>> {
-        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();

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(())

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<Option<V
 ///
 /// Compressed files are decoded into an anonymous temporary file, keeping memory use bounded and
 /// preserving read-only access to the Codex home. Plain files retain their existing seek fast path.
-pub fn open_rollout_seekable_reader(path: &Path) -> io::Result<File> {
-    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<impl Read + Seek> {
+    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<R, F> {
+    inner: R,
+    is_cancelled: F,
+}
+
+impl<R, F: Fn() -> bool> CancellableReader<R, F> {
+    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<R: Read, F: Fn() -> bool> Read for CancellableReader<R, F> {
+    fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
+        self.check_cancelled()?;
+        let limit = buffer.len().min(64 * 1024);
+        let read = self.inner.read(&mut buffer[..limit])?;
+        self.check_cancelled()?;
+        Ok(read)
+    }
+}
+
+impl<R: Seek, F: Fn() -> bool> Seek for CancellableReader<R, F> {
+    fn seek(&mut self, position: SeekFrom) -> io::Result<u64> {
+        self.check_cancelled()?;
+        self.inner.seek(position)
     }
 }
 
@@ -84,10 +139,21 @@ pub fn open_rollout_seekable_reader(path: &Path) -> io::Result<File> {
 /// 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<bool> {
+pub fn rollout_contains_prefix(
+    path: &Path,
+    end_byte_offset: u64,
+    is_cancelled: impl Fn() -> bool,
+) -> io::Result<bool> {
+    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)
         }
     }

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());
     }

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<Vec<RolloutItem>> {
+    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<Vec<RolloutItem>> {
     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)?,

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,

codex-rs/thread-store/src/local/thread_history_materialization.rs

@@ -91,7 +91,8 @@ async fn read_projection_steps(
 ) -> ThreadStoreResult<(Vec<RolloutProjectionStep>, 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<u64>,
+    is_cancelled: impl Fn() -> bool,
 ) -> ThreadStoreResult<(Vec<RolloutProjectionStep>, 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 {

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");

codex-rs/tui/src/worktree_startup.rs

@@ -95,8 +95,10 @@ impl ManagedTuiWorktree {
 
 async fn latest_thread_cwd(path: Option<PathBuf>, 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<PathBuf>, fallback: PathBuf) -> PathBuf
     .await
     .ok()
     .flatten()
-    .unwrap_or(fallback)
+    .unwrap_or(fallback);
+    drop(keep_running);
+    result
 }
 
 #[allow(clippy::too_many_arguments)]