Support compression for shared rollout lineages (#41357)

## Why

Rollout compression skips files that participate in paginated fork lineages because lineage readers rely on byte offsets into the original JSONL. This prevents shared rollout history from being compressed.

## What changed

- Add a seekable rollout reader that preserves logical JSONL offsets for plain and zstd-compressed files without materializing compressed data beside the rollout.
- Use the reader for lineage validation, model-context replay, and SQLite history projection, including when compression races with a read.
- Add the opt-in `local_thread_store_shared_compression` feature. The default mode continues to leave shared lineages uncompressed for compatibility, while the new mode includes referenced sources and fork pointers.
- Record the uncompressed size in new zstd frames so common prefix-bound checks can avoid decoding the full file.

## Testing

Add coverage for seekable compressed reads, prefix validation across zstd frame shapes, compressed lineage replay and projection, concurrent forks, and compression-mode behavior.

GitOrigin-RevId: 83e519f7eabf7448977eff37c55899858cc17a07
This commit is contained in:
jif authored and copyberry committed 2026-08-28 16:04:26 +00:00
1 parent 39507eea53
commit 1cc81ca89a
14 files changed
+493 -111

No files matched your search

+6
View File
@@ -778,6 +778,9 @@
"local_thread_store_compression": { "local_thread_store_compression": {
"type": "boolean" "type": "boolean"
}, },
"local_thread_store_shared_compression": {
"type": "boolean"
},
"mcp_2026_07_28": { "mcp_2026_07_28": {
"type": "boolean" "type": "boolean"
}, },
@@ -5978,6 +5981,9 @@
"local_thread_store_compression": { "local_thread_store_compression": {
"type": "boolean" "type": "boolean"
}, },
"local_thread_store_shared_compression": {
"type": "boolean"
},
"mcp_2026_07_28": { "mcp_2026_07_28": {
"type": "boolean" "type": "boolean"
}, },
+17 -2
View File
@@ -67,6 +67,7 @@ use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_protocol::protocol::W3cTraceContext; use codex_protocol::protocol::W3cTraceContext;
use codex_rollout::RolloutCompressionMode;
use codex_rollout::state_db::StateDbHandle; use codex_rollout::state_db::StateDbHandle;
use codex_skills_extension::HostSkillsService; use codex_skills_extension::HostSkillsService;
use codex_thread_store::InMemoryThreadStore; use codex_thread_store::InMemoryThreadStore;
@@ -384,6 +385,14 @@ pub fn thread_store_from_config(
let compression_enabled = config let compression_enabled = config
.features .features
.enabled(Feature::LocalThreadStoreCompression); .enabled(Feature::LocalThreadStoreCompression);
let compression_mode = if config
.features
.enabled(Feature::LocalThreadStoreSharedCompression)
{
RolloutCompressionMode::IncludeShared
} else {
RolloutCompressionMode::Standalone
};
let background_migration_enabled = config let background_migration_enabled = config
.features .features
.enabled(Feature::BackgroundPaginatedRolloutMigration); .enabled(Feature::BackgroundPaginatedRolloutMigration);
@@ -400,11 +409,17 @@ pub fn thread_store_from_config(
warn!("failed to migrate legacy rollouts on startup: {err}"); warn!("failed to migrate legacy rollouts on startup: {err}");
} }
if compression_enabled { 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 { } 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 store
} }
+8
View File
@@ -173,6 +173,8 @@ pub enum Feature {
ExternalAgentMemoryImport, ExternalAgentMemoryImport,
/// Compress cold local thread-store rollout files. /// Compress cold local thread-store rollout files.
LocalThreadStoreCompression, 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. /// Migrate legacy local rollout files to paginated history in the background.
BackgroundPaginatedRolloutMigration, BackgroundPaginatedRolloutMigration,
/// Enable the Chronicle sidecar for passive screen-context memories. /// Enable the Chronicle sidecar for passive screen-context memories.
@@ -1068,6 +1070,12 @@ pub const FEATURES: &[FeatureSpec] = &[
stage: Stage::UnderDevelopment, stage: Stage::UnderDevelopment,
default_enabled: false, default_enabled: false,
}, },
FeatureSpec {
id: Feature::LocalThreadStoreSharedCompression,
key: "local_thread_store_shared_compression",
stage: Stage::UnderDevelopment,
default_enabled: false,
},
FeatureSpec { FeatureSpec {
id: Feature::BackgroundPaginatedRolloutMigration, id: Feature::BackgroundPaginatedRolloutMigration,
key: "background_paginated_rollout_migration", key: "background_paginated_rollout_migration",
+30 -9
View File
@@ -21,13 +21,22 @@ const OPEN_ROLLOUT_LINE_READER_RETRY_DELAY: Duration = Duration::from_millis(50)
const TEMP_SUFFIX: &str = ".tmp"; const TEMP_SUFFIX: &str = ".tmp";
static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0); 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. /// 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, /// 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 /// and a run marker under `codex_home` prevents overlapping or too-frequent
/// compression runs from the same local store. /// compression runs from the same local store.
pub fn spawn_rollout_compression_worker(codex_home: PathBuf) { pub fn spawn_rollout_compression_worker(codex_home: PathBuf, mode: RolloutCompressionMode) {
worker::spawn(codex_home) worker::spawn(codex_home, mode)
} }
/// Returns the modified time for the existing plain or compressed rollout file. /// Returns the modified time for the existing plain or compressed rollout file.
@@ -247,6 +256,7 @@ mod worker {
use crate::RolloutReferenceIndex; use crate::RolloutReferenceIndex;
use crate::SESSIONS_SUBDIR; use crate::SESSIONS_SUBDIR;
use super::RolloutCompressionMode;
use super::RolloutFile; use super::RolloutFile;
use super::metrics; use super::metrics;
use super::path; 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 { let Ok(handle) = tokio::runtime::Handle::try_current() else {
metrics::run("skipped_no_runtime"); metrics::run("skipped_no_runtime");
warn!( warn!(
@@ -334,7 +344,7 @@ mod worker {
return; return;
}; };
handle.spawn(async move { handle.spawn(async move {
if let Err(err) = run(codex_home.clone()).await { if let Err(err) = run(codex_home.clone(), mode).await {
warn!( warn!(
"rollout compression worker failed for {}: {err}", "rollout compression worker failed for {}: {err}",
codex_home.display() 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) = let Some(_maintenance_guard) =
crate::try_acquire_rollout_maintenance_lock(codex_home.as_path())? crate::try_acquire_rollout_maintenance_lock(codex_home.as_path())?
else { else {
@@ -391,8 +401,14 @@ mod worker {
if started_at.elapsed() >= WORKER_MAX_RUNTIME { if started_at.elapsed() >= WORKER_MAX_RUNTIME {
break; break;
} }
compress_rollouts_in_root(root.as_path(), started_at, &reference_index, &mut stats) compress_rollouts_in_root(
.await?; root.as_path(),
started_at,
&reference_index,
mode,
&mut stats,
)
.await?;
} }
Ok::<_, io::Error>(stats) Ok::<_, io::Error>(stats)
} }
@@ -433,6 +449,7 @@ mod worker {
root: &Path, root: &Path,
started_at: Instant, started_at: Instant,
reference_index: &RolloutReferenceIndex, reference_index: &RolloutReferenceIndex,
mode: RolloutCompressionMode,
stats: &mut CompressionStats, stats: &mut CompressionStats,
) -> io::Result<()> { ) -> io::Result<()> {
if !tokio::fs::try_exists(root).await.unwrap_or(false) { if !tokio::fs::try_exists(root).await.unwrap_or(false) {
@@ -501,12 +518,14 @@ mod worker {
metrics::file("skipped_unreadable_meta"); metrics::file("skipped_unreadable_meta");
continue; 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); stats.skipped = stats.skipped.saturating_add(1);
metrics::file("skipped_referenced"); metrics::file("skipped_referenced");
continue; 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); stats.skipped = stats.skipped.saturating_add(1);
metrics::file("skipped_fork_pointer"); metrics::file("skipped_fork_pointer");
continue; continue;
@@ -745,6 +764,8 @@ mod worker {
fn encode_zstd_to_writer(source: &Path, output: impl Write) -> io::Result<()> { fn encode_zstd_to_writer(source: &Path, output: impl Write) -> io::Result<()> {
let mut input = File::open(source)?; let mut input = File::open(source)?;
let mut encoder = zstd::stream::write::Encoder::new(output, COMPRESSION_LEVEL)?; 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)?; io::copy(&mut input, &mut encoder)?;
encoder.finish()?; encoder.finish()?;
Ok(()) Ok(())
+88 -32
View File
@@ -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"); let fresh_temp = active_path.with_file_name("rollout-fresh.jsonl.zst.tmp");
fs::write(&fresh_temp, "fresh temp")?; 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!(!active_path.exists());
assert!(compressed_rollout_path(&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())? let guard = crate::try_acquire_rollout_maintenance_lock(home.path())?
.expect("claim rollout maintenance lock"); .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!(path.exists());
assert!(!compressed_rollout_path(&path).exists()); assert!(!compressed_rollout_path(&path).exists());
drop(guard); drop(guard);
worker::run(home.path().to_path_buf()).await?; worker::run(
home.path().to_path_buf(),
RolloutCompressionMode::Standalone,
)
.await?;
assert!(!path.exists()); assert!(!path.exists());
assert!(compressed_rollout_path(&path).exists()); assert!(compressed_rollout_path(&path).exists());
Ok(()) Ok(())
} }
#[tokio::test] #[tokio::test]
async fn worker_skips_archived_paginated_fork_pointer_chain() -> anyhow::Result<()> { async fn worker_compresses_archived_fork_chain_only_with_shared_mode() -> anyhow::Result<()> {
let home = TempDir::new()?; for mode in [
let thread_id = ThreadId::from_string(&Uuid::from_u128(15).to_string())?; RolloutCompressionMode::Standalone,
let source_uuid = Uuid::from_u128(16); RolloutCompressionMode::IncludeShared,
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); // Each mode gets a fresh home, without bypassing the worker's maintenance cooldown.
write_rollout(&source_path, thread_id, "referenced source")?; let home = TempDir::new()?;
set_old_mtime(&source_path)?; 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_uuid = Uuid::from_u128(17);
let child_path = archived_rollout_path(home.path(), "2025-01-03T12-00-01", child_uuid); let child_path = archived_rollout_path(home.path(), "2025-01-03T12-00-01", child_uuid);
write_rollout(&child_path, thread_id, "fork child")?; write_rollout(&child_path, thread_id, "fork child")?;
set_history_base( set_history_base(
child_path.as_path(), child_path.as_path(),
HistoryPosition { HistoryPosition {
thread_id: source_rollout_id, thread_id: source_rollout_id,
end_ordinal_exclusive: 2, end_ordinal_exclusive: 2,
end_byte_offset: std::fs::metadata(source_path.as_path())?.len(), end_byte_offset: std::fs::metadata(source_path.as_path())?.len(),
}, },
)?; )?;
set_old_mtime(&child_path)?; 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()); for (path, original) in [
assert!(child_path.exists()); (&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(()) Ok(())
} }
@@ -383,7 +415,11 @@ async fn worker_skips_source_referenced_by_archived_compressed_rollout() -> anyh
)?; )?;
compress_now(child_path.as_path())?; 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()); assert!(source_path.exists());
Ok(()) 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)); 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")?; 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!(!source_path.exists());
assert!(compressed_rollout_path(&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))?; fs::set_permissions(&rollout_path, fs::Permissions::from_mode(0o600))?;
set_old_mtime(&rollout_path)?; 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); let compressed_path = compressed_rollout_path(&rollout_path);
assert!(!rollout_path.exists()); 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))?; fs::set_permissions(&rollout_path, fs::Permissions::from_mode(0o400))?;
let source_modified = fs::metadata(&rollout_path)?.modified()?; 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_path = compressed_rollout_path(&rollout_path);
let compressed_metadata = fs::metadata(&compressed_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); let compressed_path = compressed_rollout_path(&rollout_path);
set_old_mtime(&compressed_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!(!rollout_path.exists());
assert!(compressed_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::create_dir_all(marker_dir.as_path())?;
fs::write(marker_dir.join("rollout-compression.lock"), "recent run")?; 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!(rollout_path.exists());
assert!(!compressed_rollout_path(&rollout_path).exists()); assert!(!compressed_rollout_path(&rollout_path).exists());
+4
View File
@@ -20,6 +20,7 @@ mod reverse_jsonl_scanner;
mod rollout_file_name; mod rollout_file_name;
mod rollout_reference_index; mod rollout_reference_index;
pub(crate) mod search; pub(crate) mod search;
mod seekable_reader;
pub(crate) mod session_index; pub(crate) mod session_index;
mod sqlite_metrics; mod sqlite_metrics;
pub mod state_db; pub mod state_db;
@@ -76,11 +77,14 @@ pub static INTERACTIVE_SESSION_SOURCES: LazyLock<Vec<SessionSource>> = LazyLock:
}); });
pub use codex_protocol::protocol::SessionMeta; pub use codex_protocol::protocol::SessionMeta;
pub use compression::RolloutCompressionMode;
pub use compression::RolloutLineReader; pub use compression::RolloutLineReader;
pub use compression::existing_rollout_path; pub use compression::existing_rollout_path;
pub use compression::open_rollout_line_reader; pub use compression::open_rollout_line_reader;
pub use compression::plain_rollout_path; pub use compression::plain_rollout_path;
pub use compression::spawn_rollout_compression_worker; 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. /// Materializes a compressed rollout as plain JSONL before another rollout references it.
pub async fn materialize_rollout_for_reference( pub async fn materialize_rollout_for_reference(
+86
View File
@@ -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<Self> {
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<File> {
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<bool> {
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;
@@ -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(())
}
@@ -1,4 +1,3 @@
use std::fs::File;
use std::io; use std::io;
use codex_protocol::protocol::HistoryPosition; use codex_protocol::protocol::HistoryPosition;
@@ -26,13 +25,14 @@ mod tests;
/// Loads rollout items needed to reconstruct the latest model-visible context. /// 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 /// 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 /// 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 /// bounded cutoff is available, the scan continues to the beginning and returns the complete
/// replay it already accumulated. /// 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( pub(super) async fn load_latest_model_context(
store: &LocalThreadStore, store: &LocalThreadStore,
params: LoadThreadHistoryParams, 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) 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 lineage = store.resolve_rollout_lineage(params.thread_id).await?; let lineage = store.resolve_rollout_lineage(params.thread_id).await?;
scan_model_context_from_lineage(lineage, session_meta).await? scan_model_context_from_lineage(lineage, session_meta).await?
} else { } else {
@@ -137,7 +132,7 @@ fn scan_model_context_from_lineage_blocking(
) -> io::Result<Vec<RolloutItem>> { ) -> io::Result<Vec<RolloutItem>> {
let mut scan = ModelContextScan::default(); let mut scan = ModelContextScan::default();
'segments: for segment in lineage.segments().iter().rev() { '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) { let mut scanner = match segment.end.map(|end| end.end_byte_offset) {
Some(end_byte_offset) => ReverseJsonlScanner::new_at(file, end_byte_offset)?, Some(end_byte_offset) => ReverseJsonlScanner::new_at(file, end_byte_offset)?,
None => ReverseJsonlScanner::new(file)?, None => ReverseJsonlScanner::new(file)?,
@@ -386,9 +386,34 @@ async fn replays_nested_archived_lineage_from_frozen_prefix() {
turn_complete("child-turn"), turn_complete("child-turn"),
]; ];
assert_eq!( 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") 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<const N: usize>( fn write_paginated_rollout<const N: usize>(
@@ -68,19 +68,22 @@ pub(super) async fn revert(
}); });
} }
// HistoryPosition stores byte offsets in plain JSONL, so materialize any compressed // Preserve old-reader compatibility when introducing the first reference to a standalone
// immutable lineage before deriving the retained boundary. // source. Already-shared ancestors stay read-only; their offsets address decoded JSONL bytes.
let mut lineage = store.resolve_rollout_lineage(thread_id).await?; let mut lineage = store.resolve_rollout_lineage(thread_id).await?;
for segment in &mut lineage.segments { for segment in &mut lineage.segments {
segment.rollout_path = if segment.rollout_id() == current_rollout.rollout_id && source_meta.history_base.is_none()
codex_rollout::materialize_rollout_for_reference(segment.rollout_path.as_path()) {
.await segment.rollout_path =
.map_err(|err| ThreadStoreError::Internal { codex_rollout::materialize_rollout_for_reference(segment.rollout_path.as_path())
message: format!( .await
"failed to materialize rollout {} for revert: {err}", .map_err(|err| ThreadStoreError::Internal {
segment.rollout_path.display() message: format!(
), "failed to materialize rollout {} for revert: {err}",
})?; segment.rollout_path.display()
),
})?;
}
super::thread_history_materialization::materialize_to_sqlite( super::thread_history_materialization::materialize_to_sqlite(
store, store,
segment.rollout_id(), segment.rollout_id(),
@@ -94,21 +94,11 @@ impl LocalThreadStore {
} }
let rollout_path = match representation { let rollout_path = match representation {
LineageRepresentation::Existing => rollout_path, LineageRepresentation::Existing => rollout_path,
LineageRepresentation::PlainForReference => { LineageRepresentation::PlainForReference => super::helpers::scoped_rollout_path(
let rollout_path = super::helpers::scoped_rollout_path( self.config.codex_home.clone(),
self.config.codex_home.clone(), rollout_path.as_path(),
rollout_path.as_path(), "Codex home",
"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()
),
})?
}
}; };
let meta = codex_rollout::read_session_meta_line(rollout_path.as_path()) let meta = codex_rollout::read_session_meta_line(rollout_path.as_path())
.await .await
@@ -130,6 +120,26 @@ impl LocalThreadStore {
"source rollout is not paginated", "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 { if let Some(end) = end {
validate_cutoff_bounds(requested_thread_id, rollout_path.as_path(), &end).await?; 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", "cutoff cannot include source session metadata",
)); ));
} }
let file_len = tokio::fs::metadata(rollout_path) let path = rollout_path.to_path_buf();
.await let end_byte_offset = end.end_byte_offset;
.map_err(|err| ThreadStoreError::Internal { let contains_prefix = tokio::task::spawn_blocking(move || {
message: format!( codex_rollout::rollout_contains_prefix(&path, end_byte_offset)
"failed to read lineage metadata {}: {err}", })
rollout_path.display() .await
), .map_err(|err| ThreadStoreError::Internal {
})? message: format!("failed to join rollout prefix validation: {err}"),
.len(); })?
if end.end_byte_offset > file_len { .map_err(|err| ThreadStoreError::Internal {
message: format!(
"failed to read lineage metadata {}: {err}",
rollout_path.display()
),
})?;
if !contains_prefix {
return Err(malformed_lineage( return Err(malformed_lineage(
requested_thread_id, requested_thread_id,
"cutoff byte offset is past the source rollout", "cutoff byte offset is past the source rollout",
@@ -29,9 +29,9 @@ pub(super) async fn materialize_to_sqlite(
.as_ref() .as_ref()
.map_or(0, |state| state.next_byte_offset); .map_or(0, |state| state.next_byte_offset);
if projection_state.is_none() if projection_state.is_none()
&& !tokio::fs::try_exists(rollout_path) && codex_rollout::existing_rollout_path(rollout_path)
.await .await
.map_err(thread_store_io_error)? .is_none()
{ {
return Ok(()); return Ok(());
} }
@@ -76,13 +76,21 @@ async fn read_projection_steps(
thread_id: ThreadId, thread_id: ThreadId,
subagent_history_start_ordinal: Option<u64>, subagent_history_start_ordinal: Option<u64>,
) -> ThreadStoreResult<(Vec<RolloutProjectionStep>, u64)> { ) -> ThreadStoreResult<(Vec<RolloutProjectionStep>, u64)> {
let file_end_offset = match tokio::fs::metadata(rollout_path).await { let path = rollout_path.to_path_buf();
Ok(metadata) => metadata.len(), 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 => { Err(err) if err.kind() == std::io::ErrorKind::NotFound && start_offset == 0 => {
return Ok((Vec::new(), 0)); return Ok((Vec::new(), 0));
} }
Err(err) => return Err(thread_store_io_error(err)), 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 = let byte_count =
file_end_offset file_end_offset
.checked_sub(start_offset) .checked_sub(start_offset)
@@ -93,9 +101,6 @@ async fn read_projection_steps(
message: "durable rollout append exceeds addressable memory".to_string(), message: "durable rollout append exceeds addressable memory".to_string(),
})?; })?;
let mut bytes = vec![0; byte_count]; 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)) file.seek(SeekFrom::Start(start_offset))
.await .await
.map_err(thread_store_io_error)?; .map_err(thread_store_io_error)?;
@@ -908,6 +908,19 @@ async fn active_turn_stores_only_its_start_position() {
assert!(prepared.model_context.iter().any(|item| { assert!(prepared.model_context.iter().any(|item| {
matches!(item, RolloutItem::EventMsg(EventMsg::TurnStarted(event)) if event.turn_id == "turn-1") 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!( assert_eq!(
prepare_paginated_fork( prepare_paginated_fork(
&store, &store,
@@ -943,7 +956,7 @@ async fn paginated_fork_persists_empty_source() {
} }
#[tokio::test] #[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 home = TempDir::new().expect("temp dir");
let store = projection_store(home.path()).await; let store = projection_store(home.path()).await;
let ancestor_thread_id = ThreadId::default(); let ancestor_thread_id = ThreadId::default();
@@ -976,6 +989,17 @@ async fn paginated_fork_materializes_compressed_source_and_ancestor() {
.await .await
.expect("shutdown ancestor"); .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(); let source_thread_id = ThreadId::default();
create_paginated_subagent_thread( create_paginated_subagent_thread(
&store, &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),
prepare_paginated_fork(&store, source_thread_id, ForkBoundary::Latest), prepare_paginated_fork(&store, source_thread_id, ForkBoundary::Latest),
); );
assert!(ancestor_path.exists()); assert!(!ancestor_path.exists());
assert!(source_path.exists()); assert!(!source_path.exists());
assert!(!ancestor_compressed_path.exists()); assert!(ancestor_compressed_path.exists());
assert!(!source_compressed_path.exists()); assert!(source_compressed_path.exists());
assert_eq!( assert_eq!(
fs::metadata(&ancestor_path) fs::metadata(&ancestor_compressed_path)
.and_then(|metadata| metadata.modified()) .and_then(|metadata| metadata.modified())
.expect("read materialized ancestor timestamp"), .expect("read compressed ancestor timestamp"),
ancestor_modified ancestor_modified
); );
assert_eq!( assert_eq!(
fs::metadata(&source_path) fs::metadata(&source_compressed_path)
.and_then(|metadata| metadata.modified()) .and_then(|metadata| metadata.modified())
.expect("read materialized source timestamp"), .expect("read compressed source timestamp"),
source_modified source_modified
); );
for prepared in [first, second] { 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] #[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 home = TempDir::new().expect("temp dir");
let store = projection_store(home.path()).await; let store = projection_store(home.path()).await;
let ancestor_thread_id = ThreadId::default(); 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()); compress_rollout(source_path.as_path());
let ancestor_writer_guard = store.live_writer_locks.lock(ancestor_thread_id).await; 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_store = store.clone();
let preparation = tokio::spawn(async move { let preparation = tokio::spawn(async move {
preparation_store preparation_store
@@ -1106,12 +1169,12 @@ async fn cancelled_fork_keeps_source_reserved_until_lineage_materialization_fini
.await .await
}); });
tokio::time::timeout(Duration::from_secs(10), async { 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; tokio::task::yield_now().await;
} }
}) })
.await .await
.expect("detached lineage task should materialize the source"); .expect("fork preparation should reserve the source");
preparation.abort(); preparation.abort();
assert!( assert!(
preparation preparation
@@ -1126,14 +1189,14 @@ async fn cancelled_fork_keeps_source_reserved_until_lineage_materialization_fini
tokio::select! { tokio::select! {
biased; biased;
result = &mut delete => { 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() => {} _ = tokio::task::yield_now() => {}
} }
drop(ancestor_writer_guard); drop(ancestor_writer_guard);
delete delete
.await .await
.expect("delete source after lineage materialization finishes"); .expect("delete source after lineage resolution finishes");
} }
#[tokio::test] #[tokio::test]