diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 052093ab86..fe2e5df8ba 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -778,6 +778,9 @@ "local_thread_store_compression": { "type": "boolean" }, + "local_thread_store_shared_compression": { + "type": "boolean" + }, "mcp_2026_07_28": { "type": "boolean" }, @@ -5978,6 +5981,9 @@ "local_thread_store_compression": { "type": "boolean" }, + "local_thread_store_shared_compression": { + "type": "boolean" + }, "mcp_2026_07_28": { "type": "boolean" }, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 333e59d2d0..90fe9e06f7 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -67,6 +67,7 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; +use codex_rollout::RolloutCompressionMode; use codex_rollout::state_db::StateDbHandle; use codex_skills_extension::HostSkillsService; use codex_thread_store::InMemoryThreadStore; @@ -384,6 +385,14 @@ pub fn thread_store_from_config( let compression_enabled = config .features .enabled(Feature::LocalThreadStoreCompression); + let compression_mode = if config + .features + .enabled(Feature::LocalThreadStoreSharedCompression) + { + RolloutCompressionMode::IncludeShared + } else { + RolloutCompressionMode::Standalone + }; let background_migration_enabled = config .features .enabled(Feature::BackgroundPaginatedRolloutMigration); @@ -400,11 +409,17 @@ pub fn thread_store_from_config( warn!("failed to migrate legacy rollouts on startup: {err}"); } if compression_enabled { - codex_rollout::spawn_rollout_compression_worker(codex_home); + codex_rollout::spawn_rollout_compression_worker( + codex_home, + compression_mode, + ); } }); } else if compression_enabled { - codex_rollout::spawn_rollout_compression_worker(config.codex_home.to_path_buf()); + codex_rollout::spawn_rollout_compression_worker( + config.codex_home.to_path_buf(), + compression_mode, + ); } store } diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index da71498d32..3724a4f2d8 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -173,6 +173,8 @@ pub enum Feature { ExternalAgentMemoryImport, /// Compress cold local thread-store rollout files. LocalThreadStoreCompression, + /// Allow rollout compression on homes used exclusively by compressed-lineage-aware readers. + LocalThreadStoreSharedCompression, /// Migrate legacy local rollout files to paginated history in the background. BackgroundPaginatedRolloutMigration, /// Enable the Chronicle sidecar for passive screen-context memories. @@ -1068,6 +1070,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::LocalThreadStoreSharedCompression, + key: "local_thread_store_shared_compression", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::BackgroundPaginatedRolloutMigration, key: "background_paginated_rollout_migration", diff --git a/codex-rs/rollout/src/compression.rs b/codex-rs/rollout/src/compression.rs index 64d1745495..36336c91db 100644 --- a/codex-rs/rollout/src/compression.rs +++ b/codex-rs/rollout/src/compression.rs @@ -21,13 +21,22 @@ const OPEN_ROLLOUT_LINE_READER_RETRY_DELAY: Duration = Duration::from_millis(50) const TEMP_SUFFIX: &str = ".tmp"; static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0); +/// Whether cold files belonging to shared rollout lineages may be compressed. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RolloutCompressionMode { + /// Preserve compatibility with readers that require shared lineages to remain plain JSONL. + Standalone, + /// Requires every reader of this Codex home to support compressed shared lineages. + IncludeShared, +} + /// Starts a best-effort background job that compresses cold local rollout files. /// /// The worker is fire-and-forget: failures are logged, startup is not blocked, /// and a run marker under `codex_home` prevents overlapping or too-frequent /// compression runs from the same local store. -pub fn spawn_rollout_compression_worker(codex_home: PathBuf) { - worker::spawn(codex_home) +pub fn spawn_rollout_compression_worker(codex_home: PathBuf, mode: RolloutCompressionMode) { + worker::spawn(codex_home, mode) } /// Returns the modified time for the existing plain or compressed rollout file. @@ -247,6 +256,7 @@ mod worker { use crate::RolloutReferenceIndex; use crate::SESSIONS_SUBDIR; + use super::RolloutCompressionMode; use super::RolloutFile; use super::metrics; use super::path; @@ -324,7 +334,7 @@ mod worker { } } - pub(super) fn spawn(codex_home: PathBuf) { + pub(super) fn spawn(codex_home: PathBuf, mode: RolloutCompressionMode) { let Ok(handle) = tokio::runtime::Handle::try_current() else { metrics::run("skipped_no_runtime"); warn!( @@ -334,7 +344,7 @@ mod worker { return; }; handle.spawn(async move { - if let Err(err) = run(codex_home.clone()).await { + if let Err(err) = run(codex_home.clone(), mode).await { warn!( "rollout compression worker failed for {}: {err}", codex_home.display() @@ -343,7 +353,7 @@ mod worker { }); } - pub(super) async fn run(codex_home: PathBuf) -> io::Result<()> { + pub(super) async fn run(codex_home: PathBuf, mode: RolloutCompressionMode) -> io::Result<()> { let Some(_maintenance_guard) = crate::try_acquire_rollout_maintenance_lock(codex_home.as_path())? else { @@ -391,8 +401,14 @@ mod worker { if started_at.elapsed() >= WORKER_MAX_RUNTIME { break; } - compress_rollouts_in_root(root.as_path(), started_at, &reference_index, &mut stats) - .await?; + compress_rollouts_in_root( + root.as_path(), + started_at, + &reference_index, + mode, + &mut stats, + ) + .await?; } Ok::<_, io::Error>(stats) } @@ -433,6 +449,7 @@ mod worker { root: &Path, started_at: Instant, reference_index: &RolloutReferenceIndex, + mode: RolloutCompressionMode, stats: &mut CompressionStats, ) -> io::Result<()> { if !tokio::fs::try_exists(root).await.unwrap_or(false) { @@ -501,12 +518,14 @@ mod worker { metrics::file("skipped_unreadable_meta"); continue; }; - if reference_index.reference_count(rollout_id) > 0 { + if mode == RolloutCompressionMode::Standalone + && reference_index.reference_count(rollout_id) > 0 + { stats.skipped = stats.skipped.saturating_add(1); metrics::file("skipped_referenced"); continue; } - if meta.meta.history_base.is_some() { + if mode == RolloutCompressionMode::Standalone && meta.meta.history_base.is_some() { stats.skipped = stats.skipped.saturating_add(1); metrics::file("skipped_fork_pointer"); continue; @@ -745,6 +764,8 @@ mod worker { fn encode_zstd_to_writer(source: &Path, output: impl Write) -> io::Result<()> { let mut input = File::open(source)?; let mut encoder = zstd::stream::write::Encoder::new(output, COMPRESSION_LEVEL)?; + // Preserve fast byte-bound checks for paginated history without decoding the whole file. + encoder.set_pledged_src_size(Some(input.metadata()?.len()))?; io::copy(&mut input, &mut encoder)?; encoder.finish()?; Ok(()) diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index 48a862ae07..f9d7a74921 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -289,7 +289,11 @@ async fn worker_compresses_old_active_and_archived_rollouts() -> anyhow::Result< let fresh_temp = active_path.with_file_name("rollout-fresh.jsonl.zst.tmp"); fs::write(&fresh_temp, "fresh temp")?; - worker::run(home.path().to_path_buf()).await?; + worker::run( + home.path().to_path_buf(), + RolloutCompressionMode::Standalone, + ) + .await?; assert!(!active_path.exists()); assert!(compressed_rollout_path(&active_path).exists()); @@ -319,44 +323,72 @@ async fn worker_waits_for_rollout_maintenance_before_compressing() -> anyhow::Re let guard = crate::try_acquire_rollout_maintenance_lock(home.path())? .expect("claim rollout maintenance lock"); - worker::run(home.path().to_path_buf()).await?; + worker::run( + home.path().to_path_buf(), + RolloutCompressionMode::Standalone, + ) + .await?; assert!(path.exists()); assert!(!compressed_rollout_path(&path).exists()); drop(guard); - worker::run(home.path().to_path_buf()).await?; + worker::run( + home.path().to_path_buf(), + RolloutCompressionMode::Standalone, + ) + .await?; assert!(!path.exists()); assert!(compressed_rollout_path(&path).exists()); Ok(()) } #[tokio::test] -async fn worker_skips_archived_paginated_fork_pointer_chain() -> anyhow::Result<()> { - let home = TempDir::new()?; - let thread_id = ThreadId::from_string(&Uuid::from_u128(15).to_string())?; - let source_uuid = Uuid::from_u128(16); - let source_rollout_id = ThreadId::from_string(&source_uuid.to_string())?; - let source_path = rollout_path(home.path(), "2025-01-03T12-00-00", source_uuid); - write_rollout(&source_path, thread_id, "referenced source")?; - set_old_mtime(&source_path)?; +async fn worker_compresses_archived_fork_chain_only_with_shared_mode() -> anyhow::Result<()> { + for mode in [ + RolloutCompressionMode::Standalone, + RolloutCompressionMode::IncludeShared, + ] { + // Each mode gets a fresh home, without bypassing the worker's maintenance cooldown. + let home = TempDir::new()?; + let thread_id = ThreadId::from_string(&Uuid::from_u128(15).to_string())?; + let source_uuid = Uuid::from_u128(16); + let source_rollout_id = ThreadId::from_string(&source_uuid.to_string())?; + let source_path = rollout_path(home.path(), "2025-01-03T12-00-00", source_uuid); + write_rollout(&source_path, thread_id, "referenced source")?; + set_old_mtime(&source_path)?; - let child_uuid = Uuid::from_u128(17); - let child_path = archived_rollout_path(home.path(), "2025-01-03T12-00-01", child_uuid); - write_rollout(&child_path, thread_id, "fork child")?; - set_history_base( - child_path.as_path(), - HistoryPosition { - thread_id: source_rollout_id, - end_ordinal_exclusive: 2, - end_byte_offset: std::fs::metadata(source_path.as_path())?.len(), - }, - )?; - set_old_mtime(&child_path)?; + let child_uuid = Uuid::from_u128(17); + let child_path = archived_rollout_path(home.path(), "2025-01-03T12-00-01", child_uuid); + write_rollout(&child_path, thread_id, "fork child")?; + set_history_base( + child_path.as_path(), + HistoryPosition { + thread_id: source_rollout_id, + end_ordinal_exclusive: 2, + end_byte_offset: std::fs::metadata(source_path.as_path())?.len(), + }, + )?; + set_old_mtime(&child_path)?; + let original_source = fs::read(&source_path)?; + let original_child = fs::read(&child_path)?; - worker::run(home.path().to_path_buf()).await?; + worker::run(home.path().to_path_buf(), mode).await?; - assert!(source_path.exists()); - assert!(child_path.exists()); + for (path, original) in [ + (&source_path, original_source), + (&child_path, original_child), + ] { + match mode { + RolloutCompressionMode::Standalone => assert_eq!(fs::read(path)?, original), + RolloutCompressionMode::IncludeShared => { + assert!(!path.exists(), "shared compression should replace {path:?}"); + let mut restored = Vec::new(); + crate::open_rollout_seekable_reader(path)?.read_to_end(&mut restored)?; + assert_eq!(restored, original); + } + } + } + } Ok(()) } @@ -383,7 +415,11 @@ async fn worker_skips_source_referenced_by_archived_compressed_rollout() -> anyh )?; compress_now(child_path.as_path())?; - worker::run(home.path().to_path_buf()).await?; + worker::run( + home.path().to_path_buf(), + RolloutCompressionMode::Standalone, + ) + .await?; assert!(source_path.exists()); Ok(()) @@ -402,7 +438,11 @@ async fn worker_skips_unreadable_metadata_without_blocking_other_compression() - let unreadable_path = rollout_path(home.path(), "2025-01-03T12-00-01", Uuid::from_u128(21)); fs::write(unreadable_path.as_path(), "{not json}\n")?; - worker::run(home.path().to_path_buf()).await?; + worker::run( + home.path().to_path_buf(), + RolloutCompressionMode::Standalone, + ) + .await?; assert!(!source_path.exists()); assert!(compressed_rollout_path(&source_path).exists()); @@ -476,7 +516,11 @@ async fn compression_preserves_rollout_permissions() -> anyhow::Result<()> { fs::set_permissions(&rollout_path, fs::Permissions::from_mode(0o600))?; set_old_mtime(&rollout_path)?; - worker::run(home.path().to_path_buf()).await?; + worker::run( + home.path().to_path_buf(), + RolloutCompressionMode::Standalone, + ) + .await?; let compressed_path = compressed_rollout_path(&rollout_path); assert!(!rollout_path.exists()); @@ -558,7 +602,11 @@ async fn compression_preserves_read_only_rollout_permissions() -> anyhow::Result fs::set_permissions(&rollout_path, fs::Permissions::from_mode(0o400))?; let source_modified = fs::metadata(&rollout_path)?.modified()?; - worker::run(home.path().to_path_buf()).await?; + worker::run( + home.path().to_path_buf(), + RolloutCompressionMode::Standalone, + ) + .await?; let compressed_path = compressed_rollout_path(&rollout_path); let compressed_metadata = fs::metadata(&compressed_path)?; @@ -579,7 +627,11 @@ async fn worker_skips_existing_compressed_archived_rollouts() -> anyhow::Result< let compressed_path = compressed_rollout_path(&rollout_path); set_old_mtime(&compressed_path)?; - worker::run(home.path().to_path_buf()).await?; + worker::run( + home.path().to_path_buf(), + RolloutCompressionMode::Standalone, + ) + .await?; assert!(!rollout_path.exists()); assert!(compressed_path.exists()); @@ -603,7 +655,11 @@ async fn worker_skips_when_fresh_run_marker_exists() -> anyhow::Result<()> { fs::create_dir_all(marker_dir.as_path())?; fs::write(marker_dir.join("rollout-compression.lock"), "recent run")?; - worker::run(home.path().to_path_buf()).await?; + worker::run( + home.path().to_path_buf(), + RolloutCompressionMode::Standalone, + ) + .await?; assert!(rollout_path.exists()); assert!(!compressed_rollout_path(&rollout_path).exists()); diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 04733d33bf..e319c933c1 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -20,6 +20,7 @@ mod reverse_jsonl_scanner; mod rollout_file_name; mod rollout_reference_index; pub(crate) mod search; +mod seekable_reader; pub(crate) mod session_index; mod sqlite_metrics; pub mod state_db; @@ -76,11 +77,14 @@ pub static INTERACTIVE_SESSION_SOURCES: LazyLock> = LazyLock: }); pub use codex_protocol::protocol::SessionMeta; +pub use compression::RolloutCompressionMode; pub use compression::RolloutLineReader; pub use compression::existing_rollout_path; pub use compression::open_rollout_line_reader; pub use compression::plain_rollout_path; pub use compression::spawn_rollout_compression_worker; +pub use seekable_reader::open_rollout_seekable_reader; +pub use seekable_reader::rollout_contains_prefix; /// Materializes a compressed rollout as plain JSONL before another rollout references it. pub async fn materialize_rollout_for_reference( diff --git a/codex-rs/rollout/src/seekable_reader.rs b/codex-rs/rollout/src/seekable_reader.rs new file mode 100644 index 0000000000..29b4f0173b --- /dev/null +++ b/codex-rs/rollout/src/seekable_reader.rs @@ -0,0 +1,86 @@ +//! Reads rollout byte positions independently of their plain or compressed representation. +//! +//! 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. + +use std::fs::File; +use std::io; +use std::io::Read; +use std::io::Seek; +use std::path::Path; +use std::time::Duration; + +use crate::plain_rollout_path; + +enum RolloutReader { + Plain(File), + Compressed(File), +} + +impl RolloutReader { + fn open(path: &Path) -> io::Result { + let plain_path = plain_rollout_path(path); + let compressed_path = plain_path.with_extension("jsonl.zst"); + for attempt in 0..4 { + match File::open(&plain_path) { + Ok(file) => return Ok(Self::Plain(file)), + Err(err) if err.kind() == io::ErrorKind::NotFound => {} + Err(err) => return Err(err), + } + match File::open(&compressed_path) { + Ok(file) => return Ok(Self::Compressed(file)), + Err(err) if err.kind() == io::ErrorKind::NotFound && attempt < 3 => { + std::thread::sleep(Duration::from_millis(50)); + } + Err(err) => return Err(err), + } + } + unreachable!("the final open attempt returns its error") + } +} + +/// Opens the original JSONL bytes for blocking offset reads without changing the rollout on disk. +/// +/// 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 { + match RolloutReader::open(path)? { + RolloutReader::Plain(file) => Ok(file), + RolloutReader::Compressed(file) => { + let mut decoded = tempfile::tempfile()?; + io::copy(&mut zstd::stream::read::Decoder::new(file)?, &mut decoded)?; + decoded.rewind()?; + Ok(decoded) + } + } +} + +/// Checks a frozen prefix's byte bound using a blocking read of the logical JSONL representation. +/// +/// 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 { + match RolloutReader::open(path)? { + RolloutReader::Plain(file) => Ok(end_byte_offset <= file.metadata()?.len()), + RolloutReader::Compressed(mut file) => { + // A zstd frame header occupies at most 18 bytes. + let mut header = [0; 18]; + let read = file.read(&mut header)?; + if zstd::zstd_safe::get_frame_content_size(&header[..read]) + .ok() + .flatten() + .is_some_and(|size| end_byte_offset <= size) + { + return Ok(true); + } + file.rewind()?; + let mut prefix = zstd::stream::read::Decoder::new(file)?.take(end_byte_offset); + Ok(io::copy(&mut prefix, &mut io::sink())? == end_byte_offset) + } + } +} + +#[cfg(test)] +#[path = "seekable_reader_tests.rs"] +mod tests; diff --git a/codex-rs/rollout/src/seekable_reader_tests.rs b/codex-rs/rollout/src/seekable_reader_tests.rs new file mode 100644 index 0000000000..abc2866490 --- /dev/null +++ b/codex-rs/rollout/src/seekable_reader_tests.rs @@ -0,0 +1,79 @@ +use std::fs; +use std::io::Read; +use std::io::Seek; +use std::io::SeekFrom; +use std::io::Write; + +use pretty_assertions::assert_eq; + +use super::*; + +#[test] +fn compressed_offsets_preserve_exact_bytes_and_open_snapshots() -> io::Result<()> { + let home = tempfile::tempdir()?; + let path = home.path().join("rollout.jsonl"); + let compressed = path.with_extension("jsonl.zst"); + let original = "first\r\nsecond 🦀\nthird\n".as_bytes(); + fs::write( + &compressed, + zstd::stream::encode_all(original, /*level*/ 3)?, + )?; + + let mut reader = open_rollout_seekable_reader(&path)?; + // The retained snapshot must remain valid after another caller materializes the logical path. + fs::write(&path, b"replacement\n")?; + fs::remove_file(&compressed)?; + reader.seek(SeekFrom::Start(7))?; + let mut suffix = Vec::new(); + reader.read_to_end(&mut suffix)?; + assert_eq!(suffix, original[7..]); + + let mut current = Vec::new(); + open_rollout_seekable_reader(&compressed)?.read_to_end(&mut current)?; + assert_eq!(current, b"replacement\n"); + + let mut plain_snapshot = open_rollout_seekable_reader(&path)?; + fs::write( + &compressed, + zstd::stream::encode_all(current.as_slice(), /*level*/ 3)?, + )?; + fs::remove_file(&path)?; + let mut retained = Vec::new(); + plain_snapshot.read_to_end(&mut retained)?; + assert_eq!(retained, current); + Ok(()) +} + +#[test] +fn prefix_bounds_support_sized_unsized_and_concatenated_frames() -> io::Result<()> { + let home = tempfile::tempdir()?; + let path = home.path().join("rollout.jsonl.zst"); + let first = "first 🦀\n".as_bytes(); + let second = b"second\n"; + let mut encoder = zstd::stream::write::Encoder::new(Vec::new(), /*level*/ 3)?; + encoder.set_pledged_src_size(Some(first.len() as u64))?; + encoder.write_all(first)?; + let sized = encoder.finish()?; + let without_size = zstd::stream::encode_all(first, /*level*/ 3)?; + let mut concatenated = sized.clone(); + concatenated.extend(zstd::stream::encode_all( + second.as_slice(), + /*level*/ 3, + )?); + + for (compressed, expected) in [ + (sized, first.to_vec()), + (without_size, first.to_vec()), + (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 mut actual = Vec::new(); + open_rollout_seekable_reader(&path)?.read_to_end(&mut actual)?; + assert_eq!(actual, expected); + assert!(!path.with_extension("").exists()); + } + Ok(()) +} diff --git a/codex-rs/thread-store/src/local/model_context.rs b/codex-rs/thread-store/src/local/model_context.rs index ce70a00bf4..d27718e12f 100644 --- a/codex-rs/thread-store/src/local/model_context.rs +++ b/codex-rs/thread-store/src/local/model_context.rs @@ -1,4 +1,3 @@ -use std::fs::File; use std::io; use codex_protocol::protocol::HistoryPosition; @@ -26,13 +25,14 @@ mod tests; /// Loads rollout items needed to reconstruct the latest model-visible context. /// -/// Plain paginated JSONL rollouts use a reverse scan. When it finds both a usable replacement- +/// Paginated JSONL rollouts use a reverse scan. When it finds both a usable replacement- /// history checkpoint and the completed user-turn context needed for resume metadata, the returned /// replay starts with the canonical `SessionMeta` followed by that newest suffix. When no /// bounded cutoff is available, the scan continues to the beginning and returns the complete /// replay it already accumulated. /// -/// Legacy and compressed rollout shapes keep the existing full-history path. +/// Compressed segments are decoded before applying their original JSONL offsets. Legacy rollouts +/// keep the existing full-history path. pub(super) async fn load_latest_model_context( store: &LocalThreadStore, params: LoadThreadHistoryParams, @@ -65,12 +65,7 @@ pub(super) async fn load_latest_model_context( }); } - let items = if matches!(session_meta.meta.history_mode, ThreadHistoryMode::Paginated) - && !path - .file_name() - .and_then(|file_name| file_name.to_str()) - .is_some_and(|file_name| file_name.ends_with(".jsonl.zst")) - { + let items = if matches!(session_meta.meta.history_mode, ThreadHistoryMode::Paginated) { let lineage = store.resolve_rollout_lineage(params.thread_id).await?; scan_model_context_from_lineage(lineage, session_meta).await? } else { @@ -137,7 +132,7 @@ fn scan_model_context_from_lineage_blocking( ) -> io::Result> { let mut scan = ModelContextScan::default(); 'segments: for segment in lineage.segments().iter().rev() { - let file = File::open(segment.rollout_path.as_path())?; + let file = codex_rollout::open_rollout_seekable_reader(segment.rollout_path.as_path())?; 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/model_context_tests.rs b/codex-rs/thread-store/src/local/model_context_tests.rs index e86f040704..af8e75ae53 100644 --- a/codex-rs/thread-store/src/local/model_context_tests.rs +++ b/codex-rs/thread-store/src/local/model_context_tests.rs @@ -386,9 +386,34 @@ async fn replays_nested_archived_lineage_from_frozen_prefix() { turn_complete("child-turn"), ]; assert_eq!( - serde_json::to_value(context.items).expect("serialize context"), + serde_json::to_value(&context.items).expect("serialize context"), + serde_json::to_value(&expected).expect("serialize expected context") + ); + // The same frozen lineage must replay from compressed files, without materializing or + // accidentally including the archived root's records after the inherited cutoff. + for path in [&archived_root, &middle_path, &child_path] { + let input = std::fs::File::open(path).expect("open rollout"); + let output = std::fs::File::create(path.with_extension("jsonl.zst")) + .expect("create compressed rollout"); + zstd::stream::copy_encode(input, output, /*level*/ 3).expect("compress rollout"); + std::fs::remove_file(path).expect("remove plain rollout"); + } + let compressed_context = store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id: child_id, + include_archived: false, + }) + .await + .expect("load compressed lineage model context"); + assert_eq!( + serde_json::to_value(compressed_context.items).expect("serialize compressed context"), serde_json::to_value(expected).expect("serialize expected context") ); + assert!( + [archived_root, middle_path, child_path] + .iter() + .all(|path| !path.exists()) + ); } fn write_paginated_rollout( diff --git a/codex-rs/thread-store/src/local/revert_thread.rs b/codex-rs/thread-store/src/local/revert_thread.rs index e2418f7845..66d47c518d 100644 --- a/codex-rs/thread-store/src/local/revert_thread.rs +++ b/codex-rs/thread-store/src/local/revert_thread.rs @@ -68,19 +68,22 @@ pub(super) async fn revert( }); } - // HistoryPosition stores byte offsets in plain JSONL, so materialize any compressed - // immutable lineage before deriving the retained boundary. + // Preserve old-reader compatibility when introducing the first reference to a standalone + // source. Already-shared ancestors stay read-only; their offsets address decoded JSONL bytes. let mut lineage = store.resolve_rollout_lineage(thread_id).await?; for segment in &mut lineage.segments { - segment.rollout_path = - codex_rollout::materialize_rollout_for_reference(segment.rollout_path.as_path()) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!( - "failed to materialize rollout {} for revert: {err}", - segment.rollout_path.display() - ), - })?; + if segment.rollout_id() == current_rollout.rollout_id && source_meta.history_base.is_none() + { + segment.rollout_path = + codex_rollout::materialize_rollout_for_reference(segment.rollout_path.as_path()) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to materialize rollout {} for revert: {err}", + segment.rollout_path.display() + ), + })?; + } super::thread_history_materialization::materialize_to_sqlite( store, segment.rollout_id(), diff --git a/codex-rs/thread-store/src/local/rollout_lineage.rs b/codex-rs/thread-store/src/local/rollout_lineage.rs index e5d5a83437..e290756ee4 100644 --- a/codex-rs/thread-store/src/local/rollout_lineage.rs +++ b/codex-rs/thread-store/src/local/rollout_lineage.rs @@ -94,21 +94,11 @@ impl LocalThreadStore { } let rollout_path = match representation { LineageRepresentation::Existing => rollout_path, - LineageRepresentation::PlainForReference => { - let rollout_path = super::helpers::scoped_rollout_path( - self.config.codex_home.clone(), - rollout_path.as_path(), - "Codex home", - )?; - codex_rollout::materialize_rollout_for_reference(rollout_path.as_path()) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!( - "failed to materialize referenced rollout {}: {err}", - rollout_path.display() - ), - })? - } + LineageRepresentation::PlainForReference => super::helpers::scoped_rollout_path( + self.config.codex_home.clone(), + rollout_path.as_path(), + "Codex home", + )?, }; let meta = codex_rollout::read_session_meta_line(rollout_path.as_path()) .await @@ -130,6 +120,26 @@ impl LocalThreadStore { "source rollout is not paginated", )); } + let rollout_path = match representation { + LineageRepresentation::Existing => rollout_path, + LineageRepresentation::PlainForReference + if next_rollout_id.is_none() && meta.meta.history_base.is_none() => + { + // A newly shared standalone source must remain readable by older binaries. + codex_rollout::materialize_rollout_for_reference(rollout_path.as_path()) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to materialize referenced rollout {}: {err}", + rollout_path.display() + ), + })? + } + // Already-shared compressed history requires a compatible reader regardless of + // new forks. Read it without publishing decoded copies into ancestors' folders; + // their owners may concurrently archive or unarchive those immutable files. + LineageRepresentation::PlainForReference => rollout_path, + }; if let Some(end) = end { validate_cutoff_bounds(requested_thread_id, rollout_path.as_path(), &end).await?; } @@ -238,16 +248,22 @@ async fn validate_cutoff_bounds( "cutoff cannot include source session metadata", )); } - let file_len = tokio::fs::metadata(rollout_path) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!( - "failed to read lineage metadata {}: {err}", - rollout_path.display() - ), - })? - .len(); - if end.end_byte_offset > file_len { + let path = rollout_path.to_path_buf(); + let end_byte_offset = end.end_byte_offset; + let contains_prefix = tokio::task::spawn_blocking(move || { + codex_rollout::rollout_contains_prefix(&path, end_byte_offset) + }) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to join rollout prefix validation: {err}"), + })? + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to read lineage metadata {}: {err}", + rollout_path.display() + ), + })?; + if !contains_prefix { return Err(malformed_lineage( requested_thread_id, "cutoff byte offset is past the source rollout", 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 a63ed7fcb1..58ea278803 100644 --- a/codex-rs/thread-store/src/local/thread_history_materialization.rs +++ b/codex-rs/thread-store/src/local/thread_history_materialization.rs @@ -29,9 +29,9 @@ pub(super) async fn materialize_to_sqlite( .as_ref() .map_or(0, |state| state.next_byte_offset); if projection_state.is_none() - && !tokio::fs::try_exists(rollout_path) + && codex_rollout::existing_rollout_path(rollout_path) .await - .map_err(thread_store_io_error)? + .is_none() { return Ok(()); } @@ -76,13 +76,21 @@ async fn read_projection_steps( thread_id: ThreadId, subagent_history_start_ordinal: Option, ) -> ThreadStoreResult<(Vec, u64)> { - let file_end_offset = match tokio::fs::metadata(rollout_path).await { - Ok(metadata) => metadata.len(), + let path = rollout_path.to_path_buf(); + let file = + tokio::task::spawn_blocking(move || codex_rollout::open_rollout_seekable_reader(&path)) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to join rollout projection read: {err}"), + })?; + let mut file = match file { + Ok(file) => tokio::fs::File::from_std(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().await.map_err(thread_store_io_error)?.len(); let byte_count = file_end_offset .checked_sub(start_offset) @@ -93,9 +101,6 @@ async fn read_projection_steps( message: "durable rollout append exceeds addressable memory".to_string(), })?; let mut bytes = vec![0; byte_count]; - let mut file = tokio::fs::File::open(rollout_path) - .await - .map_err(thread_store_io_error)?; file.seek(SeekFrom::Start(start_offset)) .await .map_err(thread_store_io_error)?; 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 4354e39244..295287f065 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 @@ -908,6 +908,19 @@ async fn active_turn_stores_only_its_start_position() { assert!(prepared.model_context.iter().any(|item| { matches!(item, RolloutItem::EventMsg(EventMsg::TurnStarted(event)) if event.turn_id == "turn-1") })); + // Another store may fork the persisted prefix while this store keeps the source writer open. + let other_store = projection_store(home.path()).await; + let other_prepared = + prepare_paginated_fork(&other_store, thread_id, ForkBoundary::Latest).await; + assert_eq!( + serde_json::to_value(( + other_prepared.history_base, + other_prepared.model_context.as_ref() + )) + .expect("serialize other store's fork snapshot"), + serde_json::to_value((prepared.history_base, prepared.model_context.as_ref())) + .expect("serialize live store's fork snapshot"), + ); assert_eq!( prepare_paginated_fork( &store, @@ -943,7 +956,7 @@ async fn paginated_fork_persists_empty_source() { } #[tokio::test] -async fn paginated_fork_materializes_compressed_source_and_ancestor() { +async fn paginated_fork_reads_compressed_shared_lineage_without_materializing() { let home = TempDir::new().expect("temp dir"); let store = projection_store(home.path()).await; let ancestor_thread_id = ThreadId::default(); @@ -976,6 +989,17 @@ async fn paginated_fork_materializes_compressed_source_and_ancestor() { .await .expect("shutdown ancestor"); + // A standalone source still becomes plain before its first shared reference, so the default + // mode does not introduce compressed lineages that older readers cannot follow. + compress_rollout(ancestor_path.as_path()); + assert_eq!( + prepare_paginated_fork(&store, ancestor_thread_id, ForkBoundary::Latest) + .await + .history_base, + Some(ancestor_base) + ); + assert!(ancestor_path.exists()); + let source_thread_id = ThreadId::default(); create_paginated_subagent_thread( &store, @@ -1022,20 +1046,20 @@ async fn paginated_fork_materializes_compressed_source_and_ancestor() { prepare_paginated_fork(&store, source_thread_id, ForkBoundary::Latest), prepare_paginated_fork(&store, source_thread_id, ForkBoundary::Latest), ); - assert!(ancestor_path.exists()); - assert!(source_path.exists()); - assert!(!ancestor_compressed_path.exists()); - assert!(!source_compressed_path.exists()); + assert!(!ancestor_path.exists()); + assert!(!source_path.exists()); + assert!(ancestor_compressed_path.exists()); + assert!(source_compressed_path.exists()); assert_eq!( - fs::metadata(&ancestor_path) + fs::metadata(&ancestor_compressed_path) .and_then(|metadata| metadata.modified()) - .expect("read materialized ancestor timestamp"), + .expect("read compressed ancestor timestamp"), ancestor_modified ); assert_eq!( - fs::metadata(&source_path) + fs::metadata(&source_compressed_path) .and_then(|metadata| metadata.modified()) - .expect("read materialized source timestamp"), + .expect("read compressed source timestamp"), source_modified ); for prepared in [first, second] { @@ -1050,10 +1074,48 @@ async fn paginated_fork_materializes_compressed_source_and_ancestor() { )); } } + + // Skipping materialization must not allow a reference that cannot be resolved inside this home. + let external_home = TempDir::new().expect("external temp dir"); + let external_path = external_home.path().join( + source_compressed_path + .file_name() + .expect("source rollout filename"), + ); + fs::rename(&source_compressed_path, &external_path).expect("move shared source outside home"); + store + .resume_thread(ResumeThreadParams { + thread_id: source_thread_id, + rollout_path: Some(external_path), + history: None, + include_archived: true, + metadata: ThreadPersistenceMetadata { + cwd: Some(home.path().to_path_buf()), + model_provider: "test-provider".to_string(), + memory_mode: ThreadMemoryMode::Enabled, + }, + }) + .await + .expect("resume external shared source"); + let error = store + .prepare_fork(PrepareForkParams { + thread_id: source_thread_id, + boundary: ForkBoundary::Latest, + }) + .await + .expect_err("external shared source cannot be referenced by rollout id"); + assert!(matches!( + error, + crate::ThreadStoreError::InvalidRequest { message } if message.contains("must be in Codex home") + )); + store + .shutdown_thread(source_thread_id) + .await + .expect("shutdown external source"); } #[tokio::test] -async fn cancelled_fork_keeps_source_reserved_until_lineage_materialization_finishes() { +async fn cancelled_fork_keeps_source_reserved_until_lineage_resolution_finishes() { let home = TempDir::new().expect("temp dir"); let store = projection_store(home.path()).await; let ancestor_thread_id = ThreadId::default(); @@ -1096,6 +1158,7 @@ async fn cancelled_fork_keeps_source_reserved_until_lineage_materialization_fini compress_rollout(source_path.as_path()); let ancestor_writer_guard = store.live_writer_locks.lock(ancestor_thread_id).await; + let source_coordination = store.live_writer_locks.coordination(source_thread_id).await; let preparation_store = store.clone(); let preparation = tokio::spawn(async move { preparation_store @@ -1106,12 +1169,12 @@ async fn cancelled_fork_keeps_source_reserved_until_lineage_materialization_fini .await }); tokio::time::timeout(Duration::from_secs(10), async { - while !source_path.exists() { + while source_coordination.lifecycle.try_write().is_ok() { tokio::task::yield_now().await; } }) .await - .expect("detached lineage task should materialize the source"); + .expect("fork preparation should reserve the source"); preparation.abort(); assert!( preparation @@ -1126,14 +1189,14 @@ async fn cancelled_fork_keeps_source_reserved_until_lineage_materialization_fini tokio::select! { biased; result = &mut delete => { - panic!("source deletion completed while lineage materialization was active: {result:?}") + panic!("source deletion completed while lineage resolution was active: {result:?}") } _ = tokio::task::yield_now() => {} } drop(ancestor_writer_guard); delete .await - .expect("delete source after lineage materialization finishes"); + .expect("delete source after lineage resolution finishes"); } #[tokio::test]