diff --git a/codex-rs/rollout/src/compression.rs b/codex-rs/rollout/src/compression.rs index 2e9ed39120..d71d56c804 100644 --- a/codex-rs/rollout/src/compression.rs +++ b/codex-rs/rollout/src/compression.rs @@ -62,11 +62,17 @@ pub(crate) fn compressed_rollout_path(path: &Path) -> PathBuf { } /// Materializes a compressed rollout back to plain `.jsonl` for async append paths. -pub(crate) async fn materialize_rollout_for_append(path: &Path) -> io::Result { +pub(crate) async fn materialize_rollout_for_append( + path: &Path, + writer_lock: Option>, +) -> io::Result { let path = path.to_path_buf(); - tokio::task::spawn_blocking(move || materialize_rollout_for_append_blocking(path.as_path())) - .await - .map_err(io::Error::other)? + tokio::task::spawn_blocking(move || { + let _writer_lock = writer_lock; + materialize_rollout_for_append_blocking(path.as_path()) + }) + .await + .map_err(io::Error::other)? } /// Materializes a compressed rollout back to plain `.jsonl` for blocking append paths. @@ -233,6 +239,7 @@ mod worker { use std::io::Write; use std::path::Path; use std::path::PathBuf; + use std::sync::Arc; use std::time::Duration; use std::time::Instant; use std::time::SystemTime; @@ -371,6 +378,7 @@ mod worker { metrics::run("started"); let started_at = Instant::now(); + let writer_locks = Arc::new(crate::WriterLockCoordinator::new(&codex_home)); let result = async { cleanup_stale_temps(codex_home.as_path()).await?; let mut stats = CompressionStats::default(); @@ -381,7 +389,8 @@ mod worker { if started_at.elapsed() >= WORKER_MAX_RUNTIME { break; } - compress_rollouts_in_root(root.as_path(), started_at, &mut stats).await?; + compress_rollouts_in_root(root.as_path(), started_at, &mut stats, &writer_locks) + .await?; } Ok::<_, io::Error>(stats) } @@ -422,6 +431,7 @@ mod worker { root: &Path, started_at: Instant, stats: &mut CompressionStats, + writer_locks: &Arc, ) -> io::Result<()> { if !tokio::fs::try_exists(root).await.unwrap_or(false) { return Ok(()); @@ -484,19 +494,24 @@ mod worker { metrics::file("skipped_unreadable_meta"); continue; } - if crate::read_session_meta_line(path.as_path()).await.is_err() { - stats.skipped = stats.skipped.saturating_add(1); - metrics::file("skipped_unreadable_meta"); - continue; - } + let thread_id = match crate::read_session_meta_line(path.as_path()).await { + Ok(metadata) => metadata.meta.id, + Err(_) => { + stats.skipped = stats.skipped.saturating_add(1); + metrics::file("skipped_unreadable_meta"); + continue; + } + }; stats.scanned = stats.scanned.saturating_add(1); metrics::file("scanned"); while jobs.len() >= MAX_CONCURRENT_COMPRESSION_JOBS { collect_next_compression_job(&mut jobs, stats).await; } + let writer_locks = Arc::clone(writer_locks); jobs.spawn_blocking(move || { let started_at = Instant::now(); - let result = compress_rollout_if_cold_blocking(path.as_path()); + let result = + compress_rollout_if_cold_blocking(path.as_path(), &writer_locks, thread_id); let duration = started_at.elapsed(); (path, duration, result) }); @@ -512,6 +527,7 @@ mod worker { enum CompressionOutcome { Compressed, SkippedNotCold, + SkippedBusy, SkippedChanged, SkippedAlreadyCompressed, } @@ -521,6 +537,7 @@ mod worker { match self { CompressionOutcome::Compressed => "compressed", CompressionOutcome::SkippedNotCold => "skipped_not_cold", + CompressionOutcome::SkippedBusy => "skipped_busy", CompressionOutcome::SkippedChanged => "skipped_changed", CompressionOutcome::SkippedAlreadyCompressed => "skipped_already_compressed", } @@ -576,6 +593,7 @@ mod worker { stats.compressed = stats.compressed.saturating_add(1); } CompressionOutcome::SkippedNotCold + | CompressionOutcome::SkippedBusy | CompressionOutcome::SkippedChanged | CompressionOutcome::SkippedAlreadyCompressed => { stats.skipped = stats.skipped.saturating_add(1); @@ -607,7 +625,11 @@ mod worker { } } - fn compress_rollout_if_cold_blocking(path: &Path) -> io::Result { + fn compress_rollout_if_cold_blocking( + path: &Path, + writer_locks: &Arc, + thread_id: codex_protocol::ThreadId, + ) -> io::Result { let before = match cold_file_state(path)? { ColdFileState::Cold(state) => state, ColdFileState::NotCold(state) => { @@ -651,6 +673,23 @@ mod worker { temp_file.as_file().sync_all()?; let compressed_bytes = temp_file.as_file().metadata()?.len(); + // Encoding and verification do not block writers. Coordination prevents writer + // acquisition while we recheck, publish, and remove the original file. + let Some(_publication_guard) = writer_locks.try_acquire_for_publication(thread_id)? else { + return Ok(CompressionMeasurement::new( + CompressionOutcome::SkippedBusy, + source_bytes, + /*compressed_bytes*/ None, + )); + }; + if !same_file_state(path, &before)? { + return Ok(CompressionMeasurement::new( + CompressionOutcome::SkippedChanged, + source_bytes, + /*compressed_bytes*/ None, + )); + } + match temp_file.persist_noclobber(compressed_path.as_path()) { Ok(_) => {} Err(err) if err.error.kind() == io::ErrorKind::AlreadyExists => { diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index a46b2adb8e..4639100ad4 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -764,3 +764,28 @@ fn set_old_mtime(path: &std::path::Path) -> anyhow::Result<()> { .set_times(times)?; Ok(()) } + +#[tokio::test] +async fn worker_compresses_parallel_cold_candidates_without_skipping_other_publications() +-> anyhow::Result<()> { + let home = TempDir::new()?; + let mut paths = Vec::new(); + for index in 0..32 { + let id = Uuid::from_u128(2000 + index); + let path = rollout_path(home.path(), "2025-01-03T12-00-00", id); + write_rollout( + &path, + ThreadId::from_string(&id.to_string()).unwrap(), + "cold candidate", + )?; + set_old_mtime(&path)?; + paths.push(path); + } + worker::run(home.path().to_path_buf()).await?; + let representations = paths + .iter() + .map(|path| (path.exists(), compressed_rollout_path(path).exists())) + .collect::>(); + assert_eq!(representations, vec![(false, true); paths.len()]); + Ok(()) +} diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 462d0f0260..28e6e3c79c 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -24,6 +24,7 @@ mod seekable_reader; pub(crate) mod session_index; mod sqlite_metrics; pub mod state_db; +mod writer_lock; pub use codex_history::CompactedItem; pub use codex_history::InitialHistory; @@ -104,7 +105,7 @@ pub use seekable_reader::rollout_contains_prefix; pub async fn materialize_rollout_for_reference( path: &std::path::Path, ) -> std::io::Result { - compression::materialize_rollout_for_append(path).await + compression::materialize_rollout_for_append(path, /*writer_lock*/ None).await } pub use config::Config; pub use config::RolloutConfig; @@ -158,6 +159,8 @@ pub use session_index::find_thread_names_by_ids; pub use session_index::remove_thread_name_entries; pub use state_db::StateDbHandle; pub use state_db::sqlite_telemetry_recorder; +pub use writer_lock::WriterLockCoordinator; +pub use writer_lock::WriterLockGuard; #[cfg(test)] mod tests; diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 8388541c64..3da60cbd4d 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -135,18 +135,24 @@ enum RolloutCmd { Shutdown { ack: oneshot::Sender>, }, + Discard { + ack: oneshot::Sender<()>, + }, } /// Observable state for the background rollout writer task. struct RolloutWriterTask { + // The task, not just its caller, owns the lock until queued file writes finish. + _writer_lock: Option>, handle: Mutex>>, terminal_failure: Mutex>>, } impl RolloutWriterTask { /// Create task observability state before spawning the writer. - fn new() -> Self { + fn new(writer_lock: Option>) -> Self { Self { + _writer_lock: writer_lock, handle: Mutex::new(None), terminal_failure: Mutex::new(None), } @@ -851,6 +857,24 @@ impl RolloutRecorder { pub async fn new( config: &impl RolloutConfigView, params: RolloutRecorderParams, + ) -> std::io::Result { + Self::new_inner(config, params, /*writer_lock*/ None).await + } + + /// Opens a recorder under existing thread-store ownership, retaining it through background IO. + /// The caller must supply the guard for this rollout's stable thread ID and Codex home. + pub async fn new_with_writer_lock( + config: &impl RolloutConfigView, + params: RolloutRecorderParams, + writer_lock: Arc, + ) -> std::io::Result { + Self::new_inner(config, params, Some(writer_lock)).await + } + + async fn new_inner( + config: &impl RolloutConfigView, + params: RolloutRecorderParams, + writer_lock: Option>, ) -> std::io::Result { // Clone the cwd for the spawned task to collect git info asynchronously. let cwd = config.cwd().to_path_buf(); @@ -933,7 +957,8 @@ impl RolloutRecorder { } } RolloutRecorderParams::Resume { path } => { - let (path, file, ordinal_state) = open_rollout_for_append(path.as_path()).await?; + let (path, file, ordinal_state) = + open_rollout_for_append(path.as_path(), writer_lock.clone()).await?; RolloutWriterState { writer: Some(JsonlWriter { file }), deferred_creation: false, @@ -955,7 +980,7 @@ impl RolloutRecorder { // Spawn a Tokio task that owns the file handle and performs async // writes. Using `tokio::fs::File` keeps everything on the async I/O // driver instead of blocking the runtime. - let writer_task = Arc::new(RolloutWriterTask::new()); + let writer_task = Arc::new(RolloutWriterTask::new(writer_lock)); let writer_task_for_spawn = Arc::clone(&writer_task); let rollout_path_for_spawn = rollout_path.clone(); let handle = tokio::task::spawn(async move { @@ -1147,6 +1172,28 @@ impl RolloutRecorder { ))); } }; + self.wait_for_exit().await + } + + /// Stops the writer without materializing deferred items, after already-running IO completes. + pub async fn discard(&self) -> std::io::Result<()> { + let (ack, done) = oneshot::channel(); + if self.tx.send(RolloutCmd::Discard { ack }).await.is_ok() { + let _ = done.await; + } + self.wait_for_exit().await + } + + async fn wait_for_exit(&self) -> std::io::Result<()> { + let handle = self + .writer_task + .handle + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + if let Some(handle) = handle { + handle.await.map_err(IoError::other)?; + } Ok(()) } } @@ -1880,6 +1927,10 @@ async fn rollout_writer( RolloutCmd::Flush { ack } => { let _ = ack.send(state.flush().await); } + RolloutCmd::Discard { ack } => { + let _ = ack.send(()); + break; + } RolloutCmd::Shutdown { ack } => match state.shutdown().await { Ok(()) => { let _ = ack.send(Ok(())); @@ -1933,7 +1984,8 @@ pub async fn append_rollout_item_to_path( rollout_path: &Path, item: &RolloutItem, ) -> std::io::Result<()> { - let (_rollout_path, file, ordinal_state) = open_rollout_for_append(rollout_path).await?; + let (_rollout_path, file, ordinal_state) = + open_rollout_for_append(rollout_path, /*writer_lock*/ None).await?; let ordinal = ordinal_state.current()?; let mut writer = JsonlWriter { file }; writer.write_rollout_item(item, ordinal).await @@ -1941,12 +1993,14 @@ pub async fn append_rollout_item_to_path( async fn open_rollout_for_append( path: &Path, + writer_lock: Option>, ) -> std::io::Result<(PathBuf, tokio::fs::File, RolloutOrdinalState)> { let refresh_modified_time = !tokio::fs::try_exists(compression::plain_rollout_path(path)).await?; - let path = compression::materialize_rollout_for_append(path).await?; + let path = compression::materialize_rollout_for_append(path, writer_lock.clone()).await?; let path_for_open = path.clone(); let (file, ordinal_state) = tokio::task::spawn_blocking(move || { + let _writer_lock = writer_lock; let mut file = File::options() .read(true) .append(true) diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index c9d4b1f0f4..3db2318926 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -169,7 +169,7 @@ async fn opening_existing_rollout_preserves_modified_time() -> std::io::Result<( drop(open_log_file(&rollout_path)?); assert_eq!(fs::metadata(&rollout_path)?.modified()?, modified); - drop(open_rollout_for_append(&rollout_path).await?); + drop(open_rollout_for_append(&rollout_path, /*writer_lock*/ None).await?); assert_eq!(fs::metadata(&rollout_path)?.modified()?, modified); Ok(()) } diff --git a/codex-rs/thread-store/src/local/writer_lock.rs b/codex-rs/rollout/src/writer_lock.rs similarity index 66% rename from codex-rs/thread-store/src/local/writer_lock.rs rename to codex-rs/rollout/src/writer_lock.rs index f6b2c8564b..ceae18e743 100644 --- a/codex-rs/thread-store/src/local/writer_lock.rs +++ b/codex-rs/rollout/src/writer_lock.rs @@ -1,3 +1,6 @@ +//! Shared cross-process ownership for local thread writers and rollout publication. +//! Thread lock-file creation and removal uses the existing home coordination lock. + use std::fs; use std::fs::File; use std::fs::OpenOptions; @@ -11,36 +14,34 @@ use std::sync::atomic::Ordering; use codex_protocol::ThreadId; use tracing::warn; -use crate::ThreadStoreError; -use crate::ThreadStoreResult; - const WRITER_LOCK_DIR: &str = "thread-writer-locks"; const COORDINATION_LOCK_FILE: &str = ".coordination.lock"; -pub(super) struct WriterLockCoordinator { +/// Coordinates writer ownership within one Codex home. +pub struct WriterLockCoordinator { directory: PathBuf, cleanup_attempted: AtomicBool, } -pub(super) struct WriterLockGuard { +/// Keeps a thread owned until all file work has finished. +pub struct WriterLockGuard { coordinator: Arc, path: PathBuf, file: Option, } impl WriterLockCoordinator { - pub(super) fn new(codex_home: &Path) -> Self { + /// Uses the same lock namespace as existing local thread-store writers. + pub fn new(codex_home: &Path) -> Self { Self { directory: codex_home.join(WRITER_LOCK_DIR), cleanup_attempted: AtomicBool::new(false), } } - pub(super) fn acquire( - self: &Arc, - thread_id: ThreadId, - ) -> ThreadStoreResult { - let coordination_lock = self.lock_coordination()?; + /// Acquires exclusive writer ownership, returning `WouldBlock` for an active writer. + pub fn acquire(self: &Arc, thread_id: ThreadId) -> io::Result { + let _coordination_lock = self.lock_coordination()?; if !self.cleanup_attempted.swap(true, Ordering::Relaxed) && let Err(err) = self.remove_stale_thread_locks() { @@ -54,31 +55,29 @@ impl WriterLockCoordinator { .create(true) .truncate(false) .open(&path) - .map_err(|err| ThreadStoreError::Internal { - message: format!( + .map_err(|err| { + io::Error::other(format!( "failed to open thread writer lock {}: {err}", path.display() - ), + )) })?; match file.try_lock() { Ok(()) => {} Err(std::fs::TryLockError::WouldBlock) => { - return Err(ThreadStoreError::Conflict { - message: format!("thread {thread_id} already has an active writer"), - }); + return Err(io::Error::new( + io::ErrorKind::WouldBlock, + format!("thread {thread_id} already has an active writer"), + )); } Err(std::fs::TryLockError::Error(err)) => { - return Err(ThreadStoreError::Internal { - message: format!( - "failed to acquire thread writer lock {}: {err}", - path.display() - ), - }); + return Err(io::Error::other(format!( + "failed to acquire thread writer lock {}: {err}", + path.display() + ))); } } - drop(coordination_lock); Ok(WriterLockGuard { coordinator: Arc::clone(self), path, @@ -86,31 +85,38 @@ impl WriterLockCoordinator { }) } - fn lock_coordination(&self) -> ThreadStoreResult { - fs::create_dir_all(&self.directory).map_err(|err| ThreadStoreError::Internal { - message: format!( - "failed to create thread writer lock directory {}: {err}", - self.directory.display() - ), - })?; - let path = self.directory.join(COORDINATION_LOCK_FILE); + /// Holds coordination through publication after probing that the thread is idle. + /// Every writer takes coordination before opening its thread lock, so the probe + /// itself can be released. Encoding and verification must finish before this call. + pub(crate) fn try_acquire_for_publication( + &self, + thread_id: ThreadId, + ) -> io::Result> { + let coordination_lock = self.lock_coordination()?; + let path = self.directory.join(format!("{thread_id}.lock")); + match OpenOptions::new().read(true).write(true).open(path) { + Ok(file) => match file.try_lock() { + Ok(()) => Ok(Some(coordination_lock)), + Err(std::fs::TryLockError::WouldBlock) => Ok(None), + Err(std::fs::TryLockError::Error(err)) => Err(err), + }, + Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(Some(coordination_lock)), + Err(err) => Err(err), + } + } + + fn lock_coordination(&self) -> io::Result { + fs::create_dir_all(&self.directory)?; let file = OpenOptions::new() .read(true) .write(true) .create(true) .truncate(false) - .open(&path) - .map_err(|err| ThreadStoreError::Internal { - message: format!( - "failed to open thread writer coordination lock {}: {err}", - path.display() - ), - })?; - file.lock().map_err(|err| ThreadStoreError::Internal { - message: format!( - "failed to acquire thread writer coordination lock {}: {err}", - path.display() - ), + .open(self.directory.join(COORDINATION_LOCK_FILE))?; + file.lock().map_err(|err| { + io::Error::other(format!( + "failed to acquire thread writer coordination lock: {err}" + )) })?; Ok(file) } diff --git a/codex-rs/thread-store/src/local/writer_lock_tests.rs b/codex-rs/rollout/src/writer_lock_tests.rs similarity index 61% rename from codex-rs/thread-store/src/local/writer_lock_tests.rs rename to codex-rs/rollout/src/writer_lock_tests.rs index 6406314225..9fad0aea36 100644 --- a/codex-rs/thread-store/src/local/writer_lock_tests.rs +++ b/codex-rs/rollout/src/writer_lock_tests.rs @@ -7,7 +7,8 @@ use tempfile::TempDir; use super::COORDINATION_LOCK_FILE; use super::WRITER_LOCK_DIR; use super::WriterLockCoordinator; -use crate::ThreadStoreError; +use pretty_assertions::assert_eq; +use std::io::ErrorKind; #[test] fn writer_locks_reject_competing_owners_and_release_their_files() { @@ -28,7 +29,7 @@ fn writer_locks_reject_competing_owners_and_release_their_files() { Ok(_) => panic!("competing owner should fail"), Err(err) => err, }; - assert!(matches!(err, ThreadStoreError::Conflict { .. })); + assert_eq!(err.kind(), ErrorKind::WouldBlock); let other_owner = secondary .acquire(other_thread_id) .expect("other thread should acquire its own lock"); @@ -74,8 +75,52 @@ fn first_acquisition_removes_stale_locks_without_removing_active_locks() { Ok(_) => panic!("active writer should survive cleanup"), Err(err) => err, }; - assert!(matches!(err, ThreadStoreError::Conflict { .. })); + assert_eq!(err.kind(), ErrorKind::WouldBlock); drop(secondary_owner); drop(active_owner); } + +#[test] +fn publication_skips_live_writers_and_keeps_coordination_locked() { + let home = TempDir::new().expect("temp dir"); + let publisher = Arc::new(WriterLockCoordinator::new(home.path())); + let writer = Arc::new(WriterLockCoordinator::new(home.path())); + let thread_id = ThreadId::default(); + let owner = writer.acquire(thread_id).expect("writer owns thread"); + assert!( + publisher + .try_acquire_for_publication(thread_id) + .unwrap() + .is_none() + ); + drop(owner); + + let publication = publisher + .try_acquire_for_publication(thread_id) + .unwrap() + .unwrap(); + let coordination = fs::OpenOptions::new() + .read(true) + .write(true) + .open(writer.directory.join(COORDINATION_LOCK_FILE)) + .unwrap(); + assert!(matches!( + coordination.try_lock(), + Err(fs::TryLockError::WouldBlock) + )); + drop(publication); + coordination + .try_lock() + .expect("publication releases coordination"); + drop(coordination); + // An existing but unlocked file is also idle; file existence is not ownership. + fs::File::create(writer.directory.join(format!("{thread_id}.lock"))).unwrap(); + drop( + publisher + .try_acquire_for_publication(thread_id) + .unwrap() + .expect("stale lock is idle"), + ); + assert!(writer.acquire(thread_id).is_ok()); +} diff --git a/codex-rs/thread-store/src/local/archive_thread.rs b/codex-rs/thread-store/src/local/archive_thread.rs index 075a0a873f..397ab0f9e2 100644 --- a/codex-rs/thread-store/src/local/archive_thread.rs +++ b/codex-rs/thread-store/src/local/archive_thread.rs @@ -243,8 +243,7 @@ mod tests { ) .expect("child session file"); let _owner_guard = owner - .writer_lock_coordinator - .acquire(child_thread_id) + .acquire_writer_lock(child_thread_id) .expect("acquire child writer lock"); let error = store diff --git a/codex-rs/thread-store/src/local/compression_writer_tests.rs b/codex-rs/thread-store/src/local/compression_writer_tests.rs new file mode 100644 index 0000000000..56ed82867a --- /dev/null +++ b/codex-rs/thread-store/src/local/compression_writer_tests.rs @@ -0,0 +1,197 @@ +//! Exercises compression through local-store ownership, including detached recorder work. + +use std::fs; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; +use std::time::SystemTime; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ThreadMemoryMode; +use codex_protocol::protocol::UserMessageEvent; +use codex_rollout::RolloutItem; +use codex_rollout::RolloutRecorder; +use codex_rollout::WriterLockCoordinator; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; +use uuid::Uuid; + +use super::LocalThreadStore; +use super::test_support::test_config; +use super::test_support::write_session_file; +use crate::ResumeThreadParams; +use crate::ThreadMetadataPatch; +use crate::ThreadPersistenceMetadata; +use crate::ThreadStore; +use crate::UpdateThreadMetadataParams; + +type TestResult = Result>; + +fn fixture() -> TestResult<(TempDir, ThreadId, PathBuf)> { + let home = TempDir::new()?; + let id = Uuid::new_v4(); + let path = write_session_file(home.path(), "2025-01-03T12-00-00", id)?; + Ok((home, ThreadId::from_string(&id.to_string())?, path)) +} + +fn age(path: &Path) -> TestResult<()> { + fs::OpenOptions::new().write(true).open(path)?.set_times( + fs::FileTimes::new().set_modified(SystemTime::now() - Duration::from_secs(8 * 86400)), + )?; + Ok(()) +} + +fn resume(thread_id: ThreadId, path: &Path, home: &Path) -> ResumeThreadParams { + ResumeThreadParams { + thread_id, + rollout_path: Some(path.to_path_buf()), + history: None, + include_archived: true, + metadata: ThreadPersistenceMetadata { + cwd: Some(home.to_path_buf()), + model_provider: "test-provider".into(), + memory_mode: ThreadMemoryMode::Enabled, + }, + } +} + +fn message() -> RolloutItem { + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: "acknowledged write".into(), + ..Default::default() + })) +} + +async fn wait_until(ready: impl Fn() -> bool) { + tokio::time::timeout(Duration::from_secs(10), async { + while !ready() { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .expect("background file work completes"); +} + +async fn compress(home: &Path) -> TestResult<()> { + let marker = home.join(".tmp/rollout-compression.lock"); + if marker.exists() { + fs::remove_file(&marker)?; + } + codex_rollout::spawn_rollout_compression_worker(home.to_path_buf()); + // The marker proves startup; the maintenance lock proves every blocking job finished. + wait_until(|| { + marker.exists() + && codex_rollout::try_acquire_rollout_maintenance_lock(home) + .unwrap() + .is_some() + }) + .await; + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn idle_recorder_retains_stable_thread_ownership_after_store_drop() -> TestResult<()> { + let (home, thread_id, old_path) = fixture()?; + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let path = old_path.with_file_name(format!( + "rollout-2025-01-03T12-00-00-{thread_id}_{}.jsonl", + Uuid::new_v4() + )); + fs::rename(old_path, &path)?; + let cold = write_session_file(home.path(), "2025-01-03T12-00-01", Uuid::new_v4())?; + age(&cold)?; + store + .resume_thread(resume(thread_id, &path, home.path())) + .await?; + let recorder = store + .live_recorders + .lock() + .await + .get(&thread_id) + .unwrap() + .recorder + .clone(); + drop(store); + age(&path)?; + let (mut expected, _, _) = RolloutRecorder::load_rollout_items(&path).await?; + compress(home.path()).await?; + assert!(path.exists()); + assert!(!path.with_extension("jsonl.zst").exists()); + assert!(cold.with_extension("jsonl.zst").exists()); + + recorder.record_canonical_items(&[message()]).await?; + drop(recorder); + let locks = Arc::new(WriterLockCoordinator::new(home.path())); + // No yield since enqueueing: only the background task can own the pending write now. + assert!( + matches!(locks.acquire(thread_id), Err(err) if err.kind() == std::io::ErrorKind::WouldBlock) + ); + wait_until(|| locks.acquire(thread_id).is_ok()).await; + age(&path)?; + compress(home.path()).await?; + assert!(!path.exists()); + expected.push(message()); + let (actual, _, _) = RolloutRecorder::load_rollout_items(&path).await?; + assert_eq!(json!(actual), json!(expected)); + Ok(()) +} + +#[tokio::test] +async fn metadata_updates_share_ownership_and_resume_compressed_rollouts() -> TestResult<()> { + let (home, thread_id, path) = fixture()?; + let config = test_config(home.path()); + let db = codex_state::StateRuntime::init( + config.sqlite.clone(), + config.default_model_provider_id.clone(), + ) + .await?; + let owner = LocalThreadStore::new(config.clone(), Some(db.clone())); + owner + .resume_thread(resume(thread_id, &path, home.path())) + .await?; + let competitor = LocalThreadStore::new(config, Some(db.clone())); + let mut patch = UpdateThreadMetadataParams { + thread_id, + include_archived: true, + patch: ThreadMetadataPatch { + memory_mode: Some(ThreadMemoryMode::Enabled), + ..Default::default() + }, + }; + owner.update_thread_metadata(patch.clone()).await?; + let original = fs::read(&path)?; + let indexed_mode = db.get_thread_memory_mode(thread_id).await?; + patch.patch.memory_mode = Some(ThreadMemoryMode::Disabled); + let error = competitor + .update_thread_metadata(patch.clone()) + .await + .unwrap_err(); + assert!(matches!(error, crate::ThreadStoreError::Conflict { .. })); + assert_eq!(fs::read(&path)?, original); + assert_eq!(db.get_thread_memory_mode(thread_id).await?, indexed_mode); + owner.update_thread_metadata(patch.clone()).await?; + owner.shutdown_thread(thread_id).await?; + let (items, _, _) = RolloutRecorder::load_rollout_items(&path).await?; + let mut expected = codex_rollout::read_session_meta_line(&path).await?; + expected.meta.memory_mode = Some("disabled".into()); + expected.git = None; + assert_eq!( + json!(items.last()), + json!(Some(RolloutItem::SessionMeta(expected.clone()))) + ); + + age(&path)?; + compress(home.path()).await?; + patch.patch.memory_mode = Some(ThreadMemoryMode::Enabled); + competitor.update_thread_metadata(patch).await?; + expected.meta.memory_mode = Some("enabled".into()); + let (items, _, _) = RolloutRecorder::load_rollout_items(&path).await?; + assert_eq!( + json!(items.last()), + json!(Some(RolloutItem::SessionMeta(expected))) + ); + Ok(()) +} diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index a8cd5b636e..70b27e7efa 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -11,6 +11,7 @@ use codex_utils_absolute_path::AbsolutePathBuf; pub(super) async fn create_thread( store: &LocalThreadStore, params: CreateThreadParams, + writer_lock: super::WriterLockGuard, ) -> ThreadStoreResult { let cwd = params .metadata @@ -26,7 +27,7 @@ pub(super) async fn create_thread( model_provider_id: params.metadata.model_provider.clone(), generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled), }; - RolloutRecorder::new( + RolloutRecorder::new_with_writer_lock( &config, RolloutRecorderParams::new( params.thread_id, @@ -57,6 +58,7 @@ pub(super) async fn create_thread( ) .with_subagent_history_start_ordinal(params.subagent_history_start_ordinal) .with_initial_window_id(params.initial_window_id), + writer_lock, ) .await .map_err(|err| ThreadStoreError::Internal { diff --git a/codex-rs/thread-store/src/local/delete_thread.rs b/codex-rs/thread-store/src/local/delete_thread.rs index 810153b800..f089ba6434 100644 --- a/codex-rs/thread-store/src/local/delete_thread.rs +++ b/codex-rs/thread-store/src/local/delete_thread.rs @@ -166,7 +166,7 @@ fn referenced_thread_error(thread_id: codex_protocol::ThreadId) -> ThreadStoreEr async fn delete_thread_after_reference_check( store: &LocalThreadStore, mut thread_rollouts: ThreadRollouts, - writer_guards: &mut Vec, + writer_guards: &mut Vec, ) -> ThreadStoreResult<()> { let thread_id = thread_rollouts.thread_id; let thread_id_str = thread_id.to_string(); @@ -206,9 +206,17 @@ async fn delete_thread_after_reference_check( super::thread_history::delete_thread(store, rollout_id).await?; } - // Drop the recorder before removing files, but retain its writer lock until cleanup finishes. - if let Some(entry) = store.live_recorders.lock().await.remove(&thread_id) { + // Stop queued file work before removing files, retaining ownership until cleanup finishes. + let live_entry = store.live_recorders.lock().await.remove(&thread_id); + if let Some(entry) = live_entry { writer_guards.push(entry.writer_lock); + entry + .recorder + .discard() + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to stop thread writer before deletion: {err}"), + })?; } let found_rollout_path = !thread_rollouts.paths.is_empty(); for rollout_path in thread_rollouts.paths { @@ -546,8 +554,7 @@ mod tests { ) .expect("child session file"); let _owner_guard = owner - .writer_lock_coordinator - .acquire(child_thread_id) + .acquire_writer_lock(child_thread_id) .expect("acquire child writer lock"); let error = store @@ -570,8 +577,7 @@ mod tests { let owner = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); let thread_id = ThreadId::default(); let _owner_guard = owner - .writer_lock_coordinator - .acquire(thread_id) + .acquire_writer_lock(thread_id) .expect("acquire writer lock before rollout exists"); let error = store diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index 3c48136272..6a2d95f2b0 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -30,8 +30,8 @@ pub(super) async fn create_thread( let _live_writer_guard = store.live_writer_locks.lock(thread_id).await; let history_mode = params.history_mode; store.ensure_live_recorder_absent(thread_id).await?; - let writer_lock = store.writer_lock_coordinator.acquire(thread_id)?; - let recorder = create_thread::create_thread(store, params).await?; + let writer_lock = store.acquire_writer_lock(thread_id)?; + let recorder = create_thread::create_thread(store, params, writer_lock.clone()).await?; store .insert_live_recorder(thread_id, recorder, thread_id, history_mode, writer_lock) .await @@ -43,7 +43,7 @@ pub(super) async fn resume_thread( ) -> ThreadStoreResult<()> { let _live_writer_guard = store.live_writer_locks.lock(params.thread_id).await; store.ensure_live_recorder_absent(params.thread_id).await?; - let writer_lock = store.writer_lock_coordinator.acquire(params.thread_id)?; + let writer_lock = store.acquire_writer_lock(params.thread_id)?; let history_mode = if let Some(history) = params.history.as_deref() { canonical_history_mode_from_rollout_items(history) } else if let Some(rollout_path) = params.rollout_path.as_ref() { @@ -105,11 +105,15 @@ pub(super) async fn resume_thread( params.thread_id, history_mode, )?; - let recorder = RolloutRecorder::new(&config, RolloutRecorderParams::resume(rollout_path)) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!("failed to resume local thread recorder: {err}"), - })?; + let recorder = RolloutRecorder::new_with_writer_lock( + &config, + RolloutRecorderParams::resume(rollout_path), + writer_lock.clone(), + ) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to resume local thread recorder: {err}"), + })?; store .insert_live_recorder( params.thread_id, @@ -207,13 +211,17 @@ pub(super) async fn discard_thread( if pending_metadata.take().is_some() { store.pending_thread_metadata.remove(thread_id).await; } - store + let entry = store .live_recorders .lock() .await .remove(&thread_id) - .map(|_| ()) - .ok_or(ThreadStoreError::ThreadNotFound { thread_id }) + .ok_or(ThreadStoreError::ThreadNotFound { thread_id })?; + entry + .recorder + .discard() + .await + .map_err(thread_store_io_error) } pub(super) async fn rollout_path( diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index a83134e443..18ee7a609c 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -22,8 +22,10 @@ mod thread_rollout_resolver; mod thread_sections; mod unarchive_thread; mod update_thread_metadata; -mod writer_lock; +#[cfg(test)] +#[path = "compression_writer_tests.rs"] +mod compression_writer_tests; #[cfg(test)] #[path = "daybreak_metadata_tests.rs"] mod daybreak_metadata_tests; @@ -37,6 +39,7 @@ use codex_protocol::ThreadId; use codex_protocol::protocol::ThreadHistoryMode; use codex_rollout::RolloutRecorder; use codex_rollout::StateDbHandle; +use codex_rollout::WriterLockCoordinator; use codex_state::SqliteConfig; use std::collections::HashMap; use std::collections::hash_map::Entry; @@ -101,8 +104,6 @@ use crate::TurnPage; use crate::UpdateProjectParams; use crate::UpdateThreadMetadataParams; use crate::UpdatedProject; -use crate::local::writer_lock::WriterLockCoordinator; -use crate::local::writer_lock::WriterLockGuard; pub use rollout_migration::RolloutMigrationFailureReason; pub use rollout_migration::RolloutMigrationMode; @@ -136,6 +137,8 @@ pub struct LocalThreadStore { thread_history_db: Arc>, } +type WriterLockGuard = Arc; + struct LiveRecorderEntry { recorder: RolloutRecorder, // Rollout projection rows are keyed by immutable rollout ID, not the stable thread ID used @@ -304,6 +307,23 @@ impl LocalThreadStore { Ok(()) } + fn acquire_writer_lock(&self, thread_id: ThreadId) -> ThreadStoreResult { + self.writer_lock_coordinator + .acquire(thread_id) + .map(Arc::new) + .map_err(|err| { + if err.kind() == std::io::ErrorKind::WouldBlock { + ThreadStoreError::Conflict { + message: err.to_string(), + } + } else { + ThreadStoreError::Internal { + message: err.to_string(), + } + } + }) + } + async fn acquire_writer_locks( &self, thread_ids: &[ThreadId], @@ -313,7 +333,7 @@ impl LocalThreadStore { if self.live_recorders.lock().await.contains_key(&thread_id) { continue; } - writer_locks.push(self.writer_lock_coordinator.acquire(thread_id)?); + writer_locks.push(self.acquire_writer_lock(thread_id)?); } Ok(writer_locks) } @@ -1366,6 +1386,20 @@ mod tests { .live_rollout_path(thread_id) .await .expect("load rollout path"); + { + let recorder = store + .live_recorders + .lock() + .await + .get(&thread_id) + .expect("live recorder") + .recorder + .clone(); + recorder + .record_canonical_items(&[user_message_item("deferred item to discard")]) + .await + .expect("queue deferred item without materializing it"); + } assert!(!rollout_path.exists()); let lock_path = home diff --git a/codex-rs/thread-store/src/local/revert_thread.rs b/codex-rs/thread-store/src/local/revert_thread.rs index 6012462c44..9f6e8f139d 100644 --- a/codex-rs/thread-store/src/local/revert_thread.rs +++ b/codex-rs/thread-store/src/local/revert_thread.rs @@ -34,7 +34,7 @@ pub(super) async fn revert( let _lifecycle_guard = store.live_writer_locks.lock_lifecycle(thread_id).await; let _live_writer_guard = store.live_writer_locks.lock(thread_id).await; store.ensure_live_recorder_absent(thread_id).await?; - let _writer_lock = store.writer_lock_coordinator.acquire(thread_id)?; + let writer_lock = store.acquire_writer_lock(thread_id)?; // Resolution may return a compressed sibling. Keep SQLite's exact stored path for the CAS. let expected_sqlite_path = state_db @@ -115,6 +115,7 @@ pub(super) async fn revert( rollout_id, history_base, forked_from_ordinal_exclusive, + writer_lock, ) .await?; let replacement_path = recorder.rollout_path().to_path_buf(); @@ -146,6 +147,7 @@ async fn create_replacement_recorder( rollout_id: ThreadId, history_base: Option, forked_from_ordinal_exclusive: Option, + writer_lock: super::WriterLockGuard, ) -> ThreadStoreResult { let config = RolloutConfig { codex_home: store.config.codex_home.clone(), @@ -179,7 +181,7 @@ async fn create_replacement_recorder( if let Some(context_window) = source_meta.context_window { params = params.with_initial_window_id(context_window.window_id); } - RolloutRecorder::new(&config, params) + RolloutRecorder::new_with_writer_lock(&config, params, writer_lock) .await .map_err(|err| ThreadStoreError::Internal { message: format!("failed to create reverted rollout: {err}"), diff --git a/codex-rs/thread-store/src/local/rollout_migration.rs b/codex-rs/thread-store/src/local/rollout_migration.rs index 79f7dd1bf1..59eaede016 100644 --- a/codex-rs/thread-store/src/local/rollout_migration.rs +++ b/codex-rs/thread-store/src/local/rollout_migration.rs @@ -411,7 +411,7 @@ impl LocalThreadStore { && options.mode == RolloutMigrationMode::Apply && let Some(thread_id) = thread_id { - let _writer_guard = match self.writer_lock_coordinator.acquire(thread_id) { + let _writer_guard = match self.acquire_writer_lock(thread_id) { Ok(guard) => guard, Err(ThreadStoreError::Conflict { message }) => { return Ok(Some(skipped_busy_outcome( @@ -535,7 +535,7 @@ impl LocalThreadStore { } let _live_writer_guard = self.live_writer_locks.lock(thread_id).await; - let _writer_guard = match self.writer_lock_coordinator.acquire(thread_id) { + let _writer_guard = match self.acquire_writer_lock(thread_id) { Ok(guard) => guard, Err(ThreadStoreError::Conflict { message }) => { return Ok(Some(skipped_busy_outcome( @@ -1003,7 +1003,7 @@ impl LocalThreadStore { legacy_names: &HashMap, limiter: &mut RolloutMigrationRateLimiter, ) -> ThreadStoreResult { - let _writer_guard = self.writer_lock_coordinator.acquire(thread_id)?; + let _writer_guard = self.acquire_writer_lock(thread_id)?; let decompressed_path = rollout_path_is_compressed(rollout_path) .then(|| decompressed_staged_rollout_path(rollout_path)) .transpose()?; diff --git a/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs b/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs index 7a4512bfcd..d2e144ca33 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs @@ -213,8 +213,7 @@ async fn recovers_pending_migrations_after_retrying_busy_rollouts() { let store = indexed_store(home.path()).await; let state_db = store.state_db().await.expect("state db"); let writer_guard = store - .writer_lock_coordinator - .acquire(busy_thread_id) + .acquire_writer_lock(busy_thread_id) .expect("hold cross-process writer lock"); store @@ -419,8 +418,7 @@ async fn retries_busy_rollouts_after_archive_and_compression_move() { let store = indexed_store(home.path()).await; let state_db = store.state_db().await.expect("state db"); let writer_guard = store - .writer_lock_coordinator - .acquire(thread_id) + .acquire_writer_lock(thread_id) .expect("hold cross-process writer lock"); store @@ -514,8 +512,7 @@ async fn treats_writer_owned_empty_rollouts_as_busy() { .await .expect("seed thread metadata"); let writer_guard = store - .writer_lock_coordinator - .acquire(thread_id) + .acquire_writer_lock(thread_id) .expect("hold cross-process writer lock"); store diff --git a/codex-rs/thread-store/src/local/rollout_migration_tests.rs b/codex-rs/thread-store/src/local/rollout_migration_tests.rs index 12ceaf44f7..2190bc0f53 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration_tests.rs @@ -2290,8 +2290,7 @@ async fn migration_skips_threads_with_an_active_writer() { ); let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); let _writer = store - .writer_lock_coordinator - .acquire(thread_id) + .acquire_writer_lock(thread_id) .expect("acquire live writer lock"); let original = fs::read(&path).expect("read active rollout"); @@ -2359,8 +2358,7 @@ async fn migration_recovers_a_published_rollout_with_missing_projection() { .expect("simulate pending migration journal"); let writer = store - .writer_lock_coordinator - .acquire(thread_id) + .acquire_writer_lock(thread_id) .expect("acquire live writer lock"); let busy = store .migrate_rollouts(apply_options()) diff --git a/codex-rs/thread-store/src/local/unarchive_thread.rs b/codex-rs/thread-store/src/local/unarchive_thread.rs index 4ef1bdb0af..017d18ad00 100644 --- a/codex-rs/thread-store/src/local/unarchive_thread.rs +++ b/codex-rs/thread-store/src/local/unarchive_thread.rs @@ -23,7 +23,7 @@ pub(super) async fn unarchive_thread( let _lifecycle_guard = store.live_writer_locks.lock_lifecycle(thread_id).await; // Archive, delete, and revert use the same cross-process lock while moving or selecting // rollout files. Unarchive must participate before it moves those files back. - let _writer_lock = store.writer_lock_coordinator.acquire(thread_id)?; + let _writer_lock = store.acquire_writer_lock(thread_id)?; let state_db_ctx = store.state_db().await; let selected_archived_path = thread_rollout_resolver::resolve_current_including_archived(store, thread_id) @@ -198,8 +198,7 @@ mod tests { let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid) .expect("archived session file"); let writer_guard = owner - .writer_lock_coordinator - .acquire(thread_id) + .acquire_writer_lock(thread_id) .expect("acquire writer lock"); let error = store diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index f1ac8a1a58..6caf5b294d 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -5,11 +5,14 @@ use chrono::Utc; use codex_protocol::SanitizedGitUrl; use codex_protocol::ThreadId; use codex_protocol::protocol::GitInfo; +use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; +use codex_rollout::RolloutConfig; use codex_rollout::RolloutItem; -use codex_rollout::append_rollout_item_to_path; +use codex_rollout::RolloutRecorder; +use codex_rollout::RolloutRecorderParams; use codex_rollout::append_thread_name; use codex_rollout::read_session_meta_line; use codex_state::ThreadMetadataBuilder; @@ -91,6 +94,25 @@ pub(super) async fn update_thread_metadata( None }; let paginated = matches!(history_mode, Some(ThreadHistoryMode::Paginated)); + let needs_rollout_compat = requires_rollout_compat || patch.name.is_some(); + // Reject competing writers before committing any part of a legacy rollout patch to SQLite. + let writer_lock = if !paginated + && needs_rollout_compat + && (patch.memory_mode.is_some() || patch.git_info.is_some()) + { + let live = store + .live_recorders + .lock() + .await + .get(&thread_id) + .map(|entry| entry.writer_lock.clone()); + Some(match live { + Some(guard) => guard, + None => store.acquire_writer_lock(thread_id)?, + }) + } else { + None + }; let require_sqlite_write = pending_patch.is_some() || sqlite_write_failure_should_block(&patch) || paginated; let mut updated = apply_metadata_update( @@ -142,7 +164,6 @@ pub(super) async fn update_thread_metadata( } return Ok(updated); } - let needs_rollout_compat = requires_rollout_compat || patch.name.is_some(); if !needs_rollout_compat { if pending_patch.is_some() { remove_pending_thread_metadata(store, thread_id, &mut pending_metadata).await; @@ -163,8 +184,21 @@ pub(super) async fn update_thread_metadata( })?; let name = patch.name; let git_info = patch.git_info; - if let Some(memory_mode) = patch.memory_mode { - apply_thread_memory_mode(resolved_rollout.path.as_path(), thread_id, memory_mode).await?; + if let Some(memory_mode) = patch.memory_mode + && let Some(writer_lock) = writer_lock.as_ref() + { + update_rollout_metadata( + store, + thread_id, + &resolved_rollout.path, + writer_lock, + |meta| { + // Replay retains prior git metadata when this marker omits it. + meta.git = None; + meta.meta.memory_mode = Some(memory_mode_as_str(memory_mode).to_string()); + }, + ) + .await?; refresh_resolved_rollout_path(&mut resolved_rollout).await; } @@ -232,14 +266,22 @@ pub(super) async fn update_thread_metadata( } None => None, }; - if let Some(((sha, branch, origin_url), memory_mode)) = resolved_git_info.as_ref() { - apply_thread_git_info_to_rollout( - resolved_rollout.path.as_path(), + if let Some(((sha, branch, origin_url), memory_mode)) = resolved_git_info.as_ref() + && let Some(writer_lock) = writer_lock.as_ref() + { + update_rollout_metadata( + store, thread_id, - sha, - branch, - origin_url, - memory_mode.as_deref(), + &resolved_rollout.path, + writer_lock, + |meta| { + meta.git = Some(GitInfo { + commit_hash: sha.as_deref().map(codex_git_utils::GitSha::new), + branch: branch.clone(), + repository_url: origin_url.clone(), + }); + meta.meta.memory_mode = memory_mode.clone(); + }, ) .await?; refresh_resolved_rollout_path(&mut resolved_rollout).await; @@ -797,71 +839,63 @@ fn resolve_git_info_patch( (sha, branch, origin_url) } -async fn apply_thread_git_info_to_rollout( - rollout_path: &Path, +// Both metadata patches use the same owned recorder; temporary writers also survive cancellation. +async fn update_rollout_metadata( + store: &LocalThreadStore, thread_id: ThreadId, - sha: &Option, - branch: &Option, - origin_url: &Option, - memory_mode: Option<&str>, + path: &Path, + writer_lock: &super::WriterLockGuard, + patch: impl FnOnce(&mut SessionMetaLine), ) -> ThreadStoreResult<()> { - let mut session_meta = - read_session_meta_line(rollout_path) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!("failed to set thread git metadata: {err}"), - })?; - if session_meta.meta.id != thread_id { + // Keep shutdown from closing the selected recorder before its metadata append is flushed. + let _live_writer_guard = store.live_writer_locks.lock(thread_id).await; + let io_error = |err| ThreadStoreError::Internal { + message: format!("failed to update rollout metadata: {err}"), + }; + let recorder = store + .live_recorders + .lock() + .await + .get(&thread_id) + .map(|entry| entry.recorder.clone()); + let mut metadata = codex_rollout::read_session_meta_line(path) + .await + .map_err(io_error)?; + if metadata.meta.id != thread_id { return Err(ThreadStoreError::Internal { message: format!( - "failed to set thread git metadata: rollout session metadata id mismatch: expected {thread_id}, found {}", - session_meta.meta.id + "rollout session metadata id mismatch: expected {thread_id}, found {}", + metadata.meta.id ), }); } - - session_meta.git = Some(GitInfo { - commit_hash: sha.as_deref().map(codex_git_utils::GitSha::new), - branch: branch.clone(), - repository_url: origin_url.clone(), - }); - session_meta.meta.memory_mode = memory_mode.map(str::to_string); - append_rollout_item_to_path(rollout_path, &RolloutItem::SessionMeta(session_meta)) + let temporary = recorder.is_none(); + let recorder = match recorder { + Some(recorder) => recorder, + None => RolloutRecorder::new_with_writer_lock( + &RolloutConfig { + codex_home: store.config.codex_home.clone(), + sqlite: store.config.sqlite.clone(), + cwd: metadata.meta.cwd.clone(), + model_provider_id: store.config.default_model_provider_id.clone(), + generate_memories: metadata.meta.memory_mode.as_deref() != Some("disabled"), + }, + RolloutRecorderParams::resume(path.to_path_buf()), + writer_lock.clone(), + ) .await - .map_err(|err| ThreadStoreError::Internal { - message: format!("failed to set thread git metadata: {err}"), - }) -} - -async fn apply_thread_memory_mode( - rollout_path: &Path, - thread_id: ThreadId, - memory_mode: ThreadMemoryMode, -) -> ThreadStoreResult<()> { - let mut session_meta = - read_session_meta_line(rollout_path) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!("failed to set thread memory mode: {err}"), - })?; - if session_meta.meta.id != thread_id { - return Err(ThreadStoreError::Internal { - message: format!( - "failed to set thread memory mode: rollout session metadata id mismatch: expected {thread_id}, found {}", - session_meta.meta.id - ), - }); + .map_err(io_error)?, + }; + patch(&mut metadata); + recorder + .record_canonical_items(&[RolloutItem::SessionMeta(metadata)]) + .await + .map_err(io_error)?; + recorder.flush().await.map_err(io_error)?; + if temporary { + recorder.shutdown().await.map_err(io_error)?; } - - // Memory-mode updates should not modify git metadata. The rollout replay - // code will preserve the latest prior git marker when this field is absent. - session_meta.git = None; - session_meta.meta.memory_mode = Some(memory_mode_as_str(memory_mode).to_string()); - append_rollout_item_to_path(rollout_path, &RolloutItem::SessionMeta(session_meta)) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!("failed to set thread memory mode: {err}"), - }) + Ok(()) } fn memory_mode_as_str(mode: ThreadMemoryMode) -> &'static str { @@ -881,6 +915,7 @@ mod tests { use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::ThreadHistoryMode; use codex_utils_absolute_path::test_support::PathExt; + use futures::FutureExt; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -1222,6 +1257,49 @@ mod tests { assert_eq!(memory_mode.as_deref(), Some("disabled")); } + #[tokio::test] + async fn metadata_patch_completes_before_concurrent_shutdown() + -> Result<(), Box> { + let home = TempDir::new()?; + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let uuid = Uuid::new_v4(); + let thread_id = ThreadId::from_string(&uuid.to_string())?; + let path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid)?; + store + .resume_thread(ResumeThreadParams { + thread_id, + rollout_path: Some(path.clone()), + history: None, + include_archived: true, + metadata: test_thread_metadata(), + }) + .await?; + let writer_lock = store.live_recorders.lock().await[&thread_id] + .writer_lock + .clone(); + let (mut expected, _, _) = RolloutRecorder::load_rollout_items(&path).await?; + let mut metadata = read_session_meta_line(&path).await?; + metadata.meta.memory_mode = Some("disabled".into()); + expected.push(RolloutItem::SessionMeta(metadata)); + + let mut shutdown = Box::pin(store.shutdown_thread(thread_id)); + update_rollout_metadata(&store, thread_id, &path, &writer_lock, |meta| { + meta.meta.memory_mode = Some("disabled".into()); + // Poll shutdown in the read/append gap; the writer mutex must keep it pending. + assert!( + tokio::task::unconstrained(shutdown.as_mut()) + .now_or_never() + .is_none() + ); + }) + .await?; + shutdown.await?; + + let (actual, _, _) = RolloutRecorder::load_rollout_items(&path).await?; + assert_eq!(json!(actual), json!(expected)); + Ok(()) + } + #[tokio::test] async fn update_thread_metadata_updates_paginated_git_info_in_sqlite_only() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs b/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs index 8116703d11..906fe7724d 100644 --- a/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs +++ b/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs @@ -3411,8 +3411,10 @@ terminal_visualization_instructions = true _ => None, }) .expect("browser should resolve a saved owner"); + // Archive stops the writer; unarchive leaves a stored thread eligible for compression. + server.thread_archive(saved_owner.id).await?; let saved_path = server - .thread_read(saved_owner.id, /*include_turns*/ false) + .thread_unarchive(saved_owner.id) .await? .path .expect("saved rollout path");