From 8035cb03f1a5061d0342cb8fa3a10a18068ca683 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 10 Apr 2026 14:12:33 +0100 Subject: [PATCH] feat: make rollout recorder reliable against errors (#17214) The rollout writer now keeps an owned/monitored task handle, returns real Result acks for flush/persist/shutdown, retries failed flushes by reopening the rollout file, and keeps buffered items until they are successfully written. Session flushes are now real durability barriers for fork/rollback/read-after-write paths, while turn completion surfaces a warning if the rollout still cannot be saved after recovery. --- codex-rs/core/src/agent/control.rs | 4 +- codex-rs/core/src/agent/control_tests.rs | 21 +- codex-rs/core/src/codex.rs | 33 +- codex-rs/core/src/codex_tests.rs | 15 +- codex-rs/core/src/codex_thread.rs | 4 +- codex-rs/core/src/guardian/review_session.rs | 2 +- codex-rs/core/src/memories/tests.rs | 7 +- codex-rs/core/src/tasks/mod.rs | 18 +- codex-rs/rollout/src/recorder.rs | 521 +++++++++++++------ codex-rs/rollout/src/recorder_tests.rs | 102 +++- 10 files changed, 536 insertions(+), 191 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index d60d1a9ff..509a578f5 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -362,7 +362,7 @@ impl AgentControl { .session .ensure_rollout_materialized() .await; - parent_thread.codex.session.flush_rollout().await; + parent_thread.codex.session.flush_rollout().await?; } let rollout_path = parent_thread @@ -663,7 +663,7 @@ impl AgentControl { let state = self.upgrade()?; let result = if let Ok(thread) = state.get_thread(agent_id).await { thread.codex.session.ensure_rollout_materialized().await; - thread.codex.session.flush_rollout().await; + thread.codex.session.flush_rollout().await?; if matches!(thread.agent_status().await, AgentStatus::Shutdown) { Ok(String::new()) } else { diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index a2c2b2279..d33285316 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -194,7 +194,12 @@ async fn persist_thread_for_tree_resume(thread: &Arc, message: &str .inject_user_message_without_turn(message.to_string()) .await; thread.codex.session.ensure_rollout_materialized().await; - thread.codex.session.flush_rollout().await; + thread + .codex + .session + .flush_rollout() + .await + .expect("test thread rollout should flush"); } async fn wait_for_live_thread_spawn_children( @@ -624,7 +629,12 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { .session .ensure_rollout_materialized() .await; - parent_thread.codex.session.flush_rollout().await; + parent_thread + .codex + .session + .flush_rollout() + .await + .expect("parent rollout should flush"); let child_thread_id = harness .control @@ -821,7 +831,12 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { .session .ensure_rollout_materialized() .await; - parent_thread.codex.session.flush_rollout().await; + parent_thread + .codex + .session + .flush_rollout() + .await + .expect("parent rollout should flush"); let child_thread_id = harness .control diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 2dea5e3a3..68de25b6f 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2216,16 +2216,16 @@ impl Session { self.services.state_db.clone() } - /// Ensure rollout file writes are durably flushed. - pub(crate) async fn flush_rollout(&self) { + /// Flush rollout writes and return the final durability-barrier result. + pub(crate) async fn flush_rollout(&self) -> std::io::Result<()> { let recorder = { let guard = self.services.rollout.lock().await; guard.clone() }; - if let Some(rec) = recorder - && let Err(e) = rec.flush().await - { - warn!("failed to flush rollout recorder: {e}"); + if let Some(recorder) = recorder { + recorder.flush().await + } else { + Ok(()) } } @@ -2372,7 +2372,7 @@ impl Session { // Defer seeding the session's initial context until the first turn starts so // turn/start overrides can be merged before we write to the rollout. if !is_subagent { - self.flush_rollout().await; + let _ = self.flush_rollout().await; } } InitialHistory::Forked(rollout_items) => { @@ -2396,7 +2396,7 @@ impl Session { // Flush after seeding history and any persisted rollout copy. if !is_subagent { - self.flush_rollout().await; + let _ = self.flush_rollout().await; } } } @@ -5600,13 +5600,24 @@ mod handlers { .into_iter() .chain(std::iter::once(RolloutItem::EventMsg(rollback_msg.clone()))) .collect::>(); - sess.persist_rollout_items(&[RolloutItem::EventMsg(rollback_msg.clone())]) - .await; - sess.flush_rollout().await; sess.apply_rollout_reconstruction(turn_context.as_ref(), replay_items.as_slice()) .await; sess.recompute_token_usage(turn_context.as_ref()).await; + sess.persist_rollout_items(&[RolloutItem::EventMsg(rollback_msg.clone())]) + .await; + if let Err(err) = sess.flush_rollout().await { + sess.send_event( + turn_context.as_ref(), + EventMsg::Warning(WarningEvent { + message: format!( + "Rolled the thread back, but failed to save the rollback marker. Codex will continue retrying. Error: {err}" + ), + }), + ) + .await; + } + sess.deliver_event_raw(Event { id: turn_context.sub_id.clone(), msg: rollback_msg, diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index fa6e608f4..cd0895a8b 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -1313,7 +1313,11 @@ async fn fork_startup_context_then_first_turn_diff_snapshot() -> anyhow::Result< // Forking reads the persisted rollout JSONL, so force the completed source turn to disk // before snapshotting from it. initial.codex.ensure_rollout_materialized().await; - initial.codex.flush_rollout().await; + initial + .codex + .flush_rollout() + .await + .expect("source rollout should flush before fork"); let mut fork_config = initial.config.clone(); fork_config.permissions.approval_policy = @@ -2359,7 +2363,10 @@ async fn attach_rollout_recorder(session: &Arc) -> PathBuf { *rollout = Some(recorder); } session.ensure_rollout_materialized().await; - session.flush_rollout().await; + session + .flush_rollout() + .await + .expect("attached rollout should flush"); rollout_path } @@ -4422,7 +4429,7 @@ async fn record_context_updates_and_set_reference_context_item_persists_baseline .expect("serialize expected context item") ); session.ensure_rollout_materialized().await; - session.flush_rollout().await; + session.flush_rollout().await.expect("rollout should flush"); let InitialHistory::Resumed(resumed) = RolloutRecorder::get_rollout_history(&rollout_path) .await @@ -4524,7 +4531,7 @@ async fn record_context_updates_and_set_reference_context_item_persists_full_rei .record_context_updates_and_set_reference_context_item(&turn_context) .await; session.ensure_rollout_materialized().await; - session.flush_rollout().await; + session.flush_rollout().await.expect("rollout should flush"); let InitialHistory::Resumed(resumed) = RolloutRecorder::get_rollout_history(&rollout_path) .await diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index f95c615a6..e90dd0f80 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -82,8 +82,8 @@ impl CodexThread { } #[doc(hidden)] - pub async fn flush_rollout(&self) { - self.codex.session.flush_rollout().await; + pub async fn flush_rollout(&self) -> std::io::Result<()> { + self.codex.session.flush_rollout().await } pub async fn submit_with_trace( diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 4f812b8ef..d84bc0da6 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -564,7 +564,7 @@ async fn append_guardian_followup_reminder(review_session: &GuardianReviewSessio async fn load_rollout_items_for_fork( session: &Session, ) -> anyhow::Result>> { - session.flush_rollout().await; + session.flush_rollout().await?; let Some(rollout_path) = session.current_rollout_path().await else { return Ok(None); }; diff --git a/codex-rs/core/src/memories/tests.rs b/codex-rs/core/src/memories/tests.rs index 6255eacde..fd9202b99 100644 --- a/codex-rs/core/src/memories/tests.rs +++ b/codex-rs/core/src/memories/tests.rs @@ -689,7 +689,12 @@ mod phase2 { other => panic!("unexpected sandbox policy: {other:?}"), } subagent.codex.session.ensure_rollout_materialized().await; - subagent.codex.session.flush_rollout().await; + subagent + .codex + .session + .flush_rollout() + .await + .expect("subagent rollout should flush"); let rollout_path = subagent .rollout_path() .expect("consolidation thread should have a rollout path"); diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 5da3650bf..b46db8e98 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -46,6 +46,7 @@ use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; +use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; use codex_features::Feature; @@ -302,7 +303,18 @@ impl Session { ) .await; let sess = session_ctx.clone_session(); - sess.flush_rollout().await; + if let Err(err) = sess.flush_rollout().await { + warn!("failed to flush rollout before completing turn: {err}"); + sess.send_event( + ctx_for_finish.as_ref(), + EventMsg::Warning(WarningEvent { + message: format!( + "Failed to save the conversation transcript; Codex will continue retrying. Error: {err}" + ), + }), + ) + .await; + } if !task_cancellation_token.is_cancelled() { // Emit completion uniformly from spawn site so all tasks share the same lifecycle. sess.on_task_finished(Arc::clone(&ctx_for_finish), last_agent_message) @@ -591,7 +603,9 @@ impl Session { .await; // Ensure the marker is durably visible before emitting TurnAborted: some clients // synchronously re-read the rollout on receipt of the abort event. - self.flush_rollout().await; + if let Err(err) = self.flush_rollout().await { + warn!("failed to flush interrupted-turn marker before emitting TurnAborted: {err}"); + } } let (completed_at, duration_ms) = task diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 6098b805e..a14214408 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -5,6 +5,8 @@ use std::fs::File; use std::io::Error as IoError; use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; use chrono::SecondsFormat; use chrono::Utc; @@ -21,6 +23,8 @@ use tokio::io::AsyncWriteExt; use tokio::sync::mpsc; use tokio::sync::mpsc::Sender; use tokio::sync::oneshot; +use tokio::task::JoinHandle; +use tracing::error; use tracing::info; use tracing::trace; use tracing::warn; @@ -70,6 +74,7 @@ use codex_utils_path as path_utils; #[derive(Clone)] pub struct RolloutRecorder { tx: Sender, + writer_task: Arc, pub(crate) rollout_path: PathBuf, state_db: Option, event_persistence_mode: EventPersistenceMode, @@ -98,13 +103,60 @@ enum RolloutCmd { }, /// Ensure all prior writes are processed; respond when flushed. Flush { - ack: oneshot::Sender<()>, + ack: oneshot::Sender>, }, Shutdown { - ack: oneshot::Sender<()>, + ack: oneshot::Sender>, }, } +/// Observable state for the background rollout writer task. +struct RolloutWriterTask { + handle: Mutex>>, + terminal_failure: Mutex>>, +} + +impl RolloutWriterTask { + /// Create task observability state before spawning the writer. + fn new() -> Self { + Self { + handle: Mutex::new(None), + terminal_failure: Mutex::new(None), + } + } + + /// Store the spawned task handle so it remains owned for the lifetime of recorder clones. + fn set_handle(&self, handle: JoinHandle<()>) { + let mut guard = self + .handle + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *guard = Some(handle); + } + + /// Remember a terminal task failure for future recorder API calls. + fn mark_failed(&self, err: &IoError) { + let mut guard = self + .terminal_failure + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *guard = Some(Arc::new(clone_io_error(err))); + } + + /// Return the terminal writer-task failure, if the task exited with an error. + fn terminal_failure(&self) -> Option { + let guard = self + .terminal_failure + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + guard.as_ref().map(|err| clone_io_error(err.as_ref())) + } +} + +fn clone_io_error(err: &IoError) -> IoError { + IoError::new(err.kind(), err.to_string()) +} + impl RolloutRecorderParams { pub fn new( conversation_id: ThreadId, @@ -472,21 +524,41 @@ 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. - tokio::task::spawn(rollout_writer( - file, - deferred_log_file_info, - rx, - meta, - cwd, - rollout_path.clone(), - state_db_ctx.clone(), - state_builder, - config.model_provider_id().to_string(), - config.generate_memories(), - )); + let writer_task = Arc::new(RolloutWriterTask::new()); + let writer_task_for_spawn = Arc::clone(&writer_task); + let rollout_path_for_spawn = rollout_path.clone(); + let default_provider = config.model_provider_id().to_string(); + let generate_memories = config.generate_memories(); + let state_db_ctx_for_spawn = state_db_ctx.clone(); + let handle = tokio::task::spawn(async move { + let result = rollout_writer( + file, + deferred_log_file_info, + rx, + meta, + cwd, + rollout_path_for_spawn.clone(), + state_db_ctx_for_spawn, + state_builder, + default_provider, + generate_memories, + ) + .await; + if let Err(err) = result { + // This is the terminal background-task failure path. Normal I/O failures stay inside + // `rollout_writer`, are reported through command acks, and leave items buffered for retry. + error!( + "rollout writer task failed for {}: {err}", + rollout_path_for_spawn.display() + ); + writer_task_for_spawn.mark_failed(&err); + } + }); + writer_task.set_handle(handle); Ok(Self { tx, + writer_task, rollout_path, state_db: state_db_ctx, event_persistence_mode, @@ -520,31 +592,53 @@ impl RolloutRecorder { self.tx .send(RolloutCmd::AddItems(filtered)) .await - .map_err(|e| IoError::other(format!("failed to queue rollout items: {e}"))) + .map_err(|e| { + self.writer_task.terminal_failure().unwrap_or_else(|| { + IoError::other(format!("failed to queue rollout items: {e}")) + }) + }) } /// Materialize the rollout file and persist all buffered items. /// - /// This is idempotent; after first materialization, repeated calls are no-ops. + /// This is idempotent. If materialization fails, the recorder keeps all pending items in memory + /// and a later `persist()` or `flush()` can retry opening and writing the rollout file. pub async fn persist(&self) -> std::io::Result<()> { let (tx, rx) = oneshot::channel(); self.tx .send(RolloutCmd::Persist { ack: tx }) .await - .map_err(|e| IoError::other(format!("failed to queue rollout persist: {e}")))?; - rx.await - .map_err(|e| IoError::other(format!("failed waiting for rollout persist: {e}")))? + .map_err(|e| { + self.writer_task.terminal_failure().unwrap_or_else(|| { + IoError::other(format!("failed to queue rollout persist: {e}")) + }) + })?; + rx.await.map_err(|e| { + self.writer_task.terminal_failure().unwrap_or_else(|| { + IoError::other(format!("failed waiting for rollout persist: {e}")) + }) + })? } /// Flush all queued writes and wait until they are committed by the writer task. + /// + /// If the first writer attempt fails, the writer drops and reopens the file handle before + /// retrying. This returns an error only when that retry also fails or the writer task is gone. pub async fn flush(&self) -> std::io::Result<()> { let (tx, rx) = oneshot::channel(); self.tx .send(RolloutCmd::Flush { ack: tx }) .await - .map_err(|e| IoError::other(format!("failed to queue rollout flush: {e}")))?; - rx.await - .map_err(|e| IoError::other(format!("failed waiting for rollout flush: {e}"))) + .map_err(|e| { + self.writer_task.terminal_failure().unwrap_or_else(|| { + IoError::other(format!("failed to queue rollout flush: {e}")) + }) + })?; + rx.await.map_err(|e| { + self.writer_task + .terminal_failure() + .unwrap_or_else(|| IoError::other(format!("failed waiting for rollout flush: {e}"))) + })? } pub async fn load_rollout_items( @@ -629,13 +723,24 @@ impl RolloutRecorder { })) } + /// Drain pending items before stopping the writer task. + /// + /// If draining fails, the writer stays alive so callers can continue retrying flush/shutdown. pub async fn shutdown(&self) -> std::io::Result<()> { let (tx_done, rx_done) = oneshot::channel(); match self.tx.send(RolloutCmd::Shutdown { ack: tx_done }).await { - Ok(_) => rx_done - .await - .map_err(|e| IoError::other(format!("failed waiting for rollout shutdown: {e}")))?, + Ok(_) => rx_done.await.map_err(|e| { + self.writer_task.terminal_failure().unwrap_or_else(|| { + IoError::other(format!("failed waiting for rollout shutdown: {e}")) + }) + })??, Err(e) => { + if let Some(err) = self.writer_task.terminal_failure() { + warn!( + "failed to send rollout shutdown command because writer task failed: {err}" + ); + return Err(err); + } warn!("failed to send rollout shutdown command: {e}"); return Err(IoError::other(format!( "failed to send rollout shutdown command: {e}" @@ -724,132 +829,259 @@ fn open_log_file(path: &Path) -> std::io::Result { .open(path) } +/// Mutable state owned by the background rollout writer. +/// +/// Items are first appended to `pending_items`; persist/flush/shutdown remove each item from that +/// queue only after it is written successfully. I/O failures drop the file handle but keep the +/// unwritten suffix so the next barrier can reopen the file and retry. +struct RolloutWriterState { + writer: Option, + deferred_log_file_info: Option, + pending_items: Vec, + meta: Option, + cwd: PathBuf, + rollout_path: PathBuf, + state_db_ctx: Option, + state_builder: Option, + default_provider: String, + generate_memories: bool, + last_logged_error: Option, +} + +impl RolloutWriterState { + #[allow(clippy::too_many_arguments)] + fn new( + file: Option, + deferred_log_file_info: Option, + meta: Option, + cwd: PathBuf, + rollout_path: PathBuf, + state_db_ctx: Option, + mut state_builder: Option, + default_provider: String, + generate_memories: bool, + ) -> Self { + if let Some(builder) = state_builder.as_mut() { + builder.rollout_path = rollout_path.clone(); + } + Self { + writer: file.map(|file| JsonlWriter { file }), + deferred_log_file_info, + pending_items: Vec::new(), + meta, + cwd, + rollout_path, + state_db_ctx, + state_builder, + default_provider, + generate_memories, + last_logged_error: None, + } + } + + fn add_items(&mut self, items: Vec) { + self.pending_items.extend(items); + } + + async fn flush_if_materialized(&mut self) { + if self.is_deferred() { + return; + } + if let Err(err) = self.flush().await { + self.enter_recovery_mode(&err); + } + } + + async fn persist(&mut self) -> std::io::Result<()> { + self.write_pending_with_recovery("persist").await + } + + async fn flush(&mut self) -> std::io::Result<()> { + if self.is_deferred() && self.pending_items.is_empty() { + return Ok(()); + } + self.write_pending_with_recovery("flush").await + } + + async fn shutdown(&mut self) -> std::io::Result<()> { + if self.is_deferred() && self.pending_items.is_empty() { + return Ok(()); + } + self.write_pending_with_recovery("shutdown").await + } + + async fn write_pending_with_recovery(&mut self, operation: &str) -> std::io::Result<()> { + match self.write_pending_once().await { + Ok(()) => { + self.last_logged_error = None; + Ok(()) + } + Err(first_err) => { + self.enter_recovery_mode(&first_err); + warn!("failed to {operation} rollout writer; reopening and retrying: {first_err}"); + match self.write_pending_once().await { + Ok(()) => { + self.last_logged_error = None; + Ok(()) + } + Err(second_err) => { + self.enter_recovery_mode(&second_err); + warn!( + "retrying rollout writer {operation} failed; first error: \ + {first_err}; final error: {second_err}" + ); + Err(second_err) + } + } + } + } + } + + fn is_deferred(&self) -> bool { + self.writer.is_none() && self.deferred_log_file_info.is_some() + } + + fn enter_recovery_mode(&mut self, err: &IoError) { + let message = err.to_string(); + if self.last_logged_error.as_ref() != Some(&message) { + error!( + "rollout writer failed for {}; buffered rollout items will be retried: {err}", + self.rollout_path.display() + ); + } + self.last_logged_error = Some(message); + self.writer = None; + } + + async fn ensure_writer_open(&mut self) -> std::io::Result<()> { + if self.writer.is_some() { + return Ok(()); + } + + let path = self + .deferred_log_file_info + .as_ref() + .map(|info| info.path.as_path()) + .unwrap_or(self.rollout_path.as_path()); + let file = open_log_file(path)?; + self.writer = Some(JsonlWriter { + file: tokio::fs::File::from_std(file), + }); + self.deferred_log_file_info = None; + Ok(()) + } + + async fn write_session_meta_if_needed(&mut self) -> std::io::Result<()> { + let Some(session_meta) = self.meta.as_ref().cloned() else { + return Ok(()); + }; + write_session_meta( + self.writer.as_mut(), + session_meta, + &self.cwd, + &self.rollout_path, + self.state_db_ctx.as_deref(), + &mut self.state_builder, + self.default_provider.as_str(), + self.generate_memories, + ) + .await?; + self.meta = None; + Ok(()) + } + + async fn write_pending_once(&mut self) -> std::io::Result<()> { + self.ensure_writer_open().await?; + self.write_session_meta_if_needed().await?; + + self.write_pending_items_once().await?; + + if let Some(writer) = self.writer.as_mut() { + writer.file.flush().await?; + } + Ok(()) + } + + async fn write_pending_items_once(&mut self) -> std::io::Result<()> { + let Some(writer) = self.writer.as_mut() else { + return Err(IoError::other("rollout writer is not open")); + }; + + let mut written_count = 0usize; + let mut write_result = Ok(()); + for item in &self.pending_items { + if let Err(err) = writer.write_rollout_item(item).await { + write_result = Err(err); + break; + } + written_count += 1; + } + + if written_count > 0 { + let written_items: Vec = + self.pending_items.drain(..written_count).collect(); + sync_thread_state_after_write( + self.state_db_ctx.as_deref(), + &self.rollout_path, + self.state_builder.as_ref(), + written_items.as_slice(), + self.default_provider.as_str(), + /*new_thread_memory_mode*/ None, + ) + .await; + } + + write_result + } +} + #[allow(clippy::too_many_arguments)] async fn rollout_writer( file: Option, - mut deferred_log_file_info: Option, + deferred_log_file_info: Option, mut rx: mpsc::Receiver, - mut meta: Option, - cwd: std::path::PathBuf, + meta: Option, + cwd: PathBuf, rollout_path: PathBuf, state_db_ctx: Option, - mut state_builder: Option, + state_builder: Option, default_provider: String, generate_memories: bool, ) -> std::io::Result<()> { - let mut writer = file.map(|file| JsonlWriter { file }); - let mut buffered_items = Vec::::new(); - if let Some(builder) = state_builder.as_mut() { - builder.rollout_path = rollout_path.clone(); - } - - // Resumed sessions already have a file handle open, so session metadata can - // be written immediately if present. - if writer.is_some() - && let Some(session_meta) = meta.take() - { - write_session_meta( - writer.as_mut(), - session_meta, - &cwd, - &rollout_path, - state_db_ctx.as_deref(), - &mut state_builder, - default_provider.as_str(), - generate_memories, - ) - .await?; - } + let mut state = RolloutWriterState::new( + file, + deferred_log_file_info, + meta, + cwd, + rollout_path, + state_db_ctx, + state_builder, + default_provider, + generate_memories, + ); // Process rollout commands while let Some(cmd) = rx.recv().await { match cmd { RolloutCmd::AddItems(items) => { - if items.is_empty() { - continue; - } - - if writer.is_none() { - buffered_items.extend(items); - continue; - } - - write_and_reconcile_items( - writer.as_mut(), - items.as_slice(), - &rollout_path, - state_db_ctx.as_deref(), - state_builder.as_ref(), - default_provider.as_str(), - ) - .await?; + state.add_items(items); + state.flush_if_materialized().await; } RolloutCmd::Persist { ack } => { - if writer.is_none() { - let result = async { - let Some(log_file_info) = deferred_log_file_info.take() else { - return Err(IoError::other( - "deferred rollout recorder missing log file metadata", - )); - }; - let file = open_log_file(log_file_info.path.as_path())?; - writer = Some(JsonlWriter { - file: tokio::fs::File::from_std(file), - }); - - if let Some(session_meta) = meta.take() { - write_session_meta( - writer.as_mut(), - session_meta, - &cwd, - &rollout_path, - state_db_ctx.as_deref(), - &mut state_builder, - default_provider.as_str(), - generate_memories, - ) - .await?; - } - - if !buffered_items.is_empty() { - write_and_reconcile_items( - writer.as_mut(), - buffered_items.as_slice(), - &rollout_path, - state_db_ctx.as_deref(), - state_builder.as_ref(), - default_provider.as_str(), - ) - .await?; - buffered_items.clear(); - } - - Ok(()) - } - .await; - - if let Err(err) = result { - let kind = err.kind(); - let message = err.to_string(); - let _ = ack.send(Err(IoError::new(kind, message.clone()))); - return Err(IoError::new(kind, message)); - } - } - let _ = ack.send(Ok(())); + let _ = ack.send(state.persist().await); } RolloutCmd::Flush { ack } => { - // Deferred fresh threads may not have an initialized file yet. - if let Some(writer) = writer.as_mut() - && let Err(e) = writer.file.flush().await - { - let _ = ack.send(()); - return Err(e); + let _ = ack.send(state.flush().await); + } + RolloutCmd::Shutdown { ack } => match state.shutdown().await { + Ok(()) => { + let _ = ack.send(Ok(())); + break; } - let _ = ack.send(()); - } - RolloutCmd::Shutdown { ack } => { - let _ = ack.send(()); - } + Err(err) => { + let _ = ack.send(Err(err)); + } + }, } } @@ -896,31 +1128,6 @@ async fn write_session_meta( Ok(()) } -async fn write_and_reconcile_items( - mut writer: Option<&mut JsonlWriter>, - items: &[RolloutItem], - rollout_path: &Path, - state_db_ctx: Option<&StateRuntime>, - state_builder: Option<&ThreadMetadataBuilder>, - default_provider: &str, -) -> std::io::Result<()> { - if let Some(writer) = writer.as_mut() { - for item in items { - writer.write_rollout_item(item).await?; - } - } - sync_thread_state_after_write( - state_db_ctx, - rollout_path, - state_builder, - items, - default_provider, - /*new_thread_memory_mode*/ None, - ) - .await; - Ok(()) -} - async fn sync_thread_state_after_write( state_db_ctx: Option<&StateRuntime>, rollout_path: &Path, diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 163c8a1ee..3e2ac5219 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -63,7 +63,7 @@ fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result std::io::Result<()> { +async fn recorder_materializes_on_flush_with_pending_items() -> std::io::Result<()> { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); let thread_id = ThreadId::new(); @@ -85,7 +85,7 @@ async fn recorder_materializes_only_after_explicit_persist() -> std::io::Result< let rollout_path = recorder.rollout_path().to_path_buf(); assert!( !rollout_path.exists(), - "rollout file should not exist before first user message" + "rollout file should not exist before the first recordable item" ); recorder @@ -99,8 +99,8 @@ async fn recorder_materializes_only_after_explicit_persist() -> std::io::Result< .await?; recorder.flush().await?; assert!( - !rollout_path.exists(), - "rollout file should remain deferred before first user message" + rollout_path.exists(), + "flush with pending items should materialize the rollout" ); recorder @@ -114,10 +114,6 @@ async fn recorder_materializes_only_after_explicit_persist() -> std::io::Result< ))]) .await?; recorder.flush().await?; - assert!( - !rollout_path.exists(), - "user-message-like items should not materialize without explicit persist" - ); recorder.persist().await?; // Second call verifies `persist()` is idempotent after materialization. @@ -146,6 +142,96 @@ async fn recorder_materializes_only_after_explicit_persist() -> std::io::Result< Ok(()) } +#[tokio::test] +async fn persist_reports_filesystem_error_and_retries_buffered_items() -> std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let thread_id = ThreadId::new(); + let recorder = RolloutRecorder::new( + &config, + RolloutRecorderParams::new( + thread_id, + /*forked_from_id*/ None, + SessionSource::Exec, + BaseInstructions::default(), + Vec::new(), + EventPersistenceMode::Limited, + ), + /*state_db_ctx*/ None, + /*state_builder*/ None, + ) + .await?; + let rollout_path = recorder.rollout_path().to_path_buf(); + + recorder + .record_items(&[RolloutItem::EventMsg(EventMsg::AgentMessage( + AgentMessageEvent { + message: "buffered-before-persist".to_string(), + phase: None, + memory_citation: None, + }, + ))]) + .await?; + let sessions_blocker_path = home.path().join("sessions"); + File::create(&sessions_blocker_path)?; + + let err = recorder + .persist() + .await + .expect_err("blocked sessions directory should fail persist"); + assert_ne!(err.kind(), std::io::ErrorKind::Interrupted); + assert!( + !rollout_path.exists(), + "failed persist should keep the rollout deferred" + ); + + fs::remove_file(sessions_blocker_path)?; + recorder.flush().await?; + let text = std::fs::read_to_string(&rollout_path)?; + assert!( + text.contains("buffered-before-persist"), + "retry should preserve items buffered before the failed persist" + ); + + recorder.shutdown().await?; + Ok(()) +} + +#[tokio::test] +async fn writer_state_retries_write_error_before_reporting_flush_success() -> std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let rollout_path = home.path().join("rollout.jsonl"); + File::create(&rollout_path)?; + let read_only_file = std::fs::OpenOptions::new().read(true).open(&rollout_path)?; + let mut state = RolloutWriterState::new( + Some(tokio::fs::File::from_std(read_only_file)), + /*deferred_log_file_info*/ None, + /*meta*/ None, + home.path().to_path_buf(), + rollout_path.clone(), + /*state_db_ctx*/ None, + /*state_builder*/ None, + config.model_provider_id.clone(), + config.generate_memories, + ); + state.add_items(vec![RolloutItem::EventMsg(EventMsg::AgentMessage( + AgentMessageEvent { + message: "queued-after-writer-error".to_string(), + phase: None, + memory_citation: None, + }, + ))]); + + state.flush().await?; + let text_after_retry = std::fs::read_to_string(&rollout_path)?; + assert!( + text_after_retry.contains("queued-after-writer-error"), + "flush should retry after reopening and write buffered items" + ); + Ok(()) +} + #[tokio::test] async fn metadata_irrelevant_events_touch_state_db_updated_at() -> std::io::Result<()> { let home = TempDir::new().expect("temp dir");