diff --git a/codex-rs/core/src/rollout/recorder.rs b/codex-rs/core/src/rollout/recorder.rs index 002269d59e..2dd3c7de66 100644 --- a/codex-rs/core/src/rollout/recorder.rs +++ b/codex-rs/core/src/rollout/recorder.rs @@ -5,6 +5,12 @@ use std::fs::{self}; use std::io::Error as IoError; use std::path::Path; use std::path::PathBuf; +#[cfg(test)] +use std::sync::Arc; +#[cfg(test)] +use std::sync::atomic::AtomicUsize; +#[cfg(test)] +use std::sync::atomic::Ordering; use chrono::SecondsFormat; use chrono::Utc; @@ -94,11 +100,11 @@ pub enum RolloutRecorderParams { enum RolloutCmd { AddItems(Vec), Persist { - ack: oneshot::Sender<()>, + ack: oneshot::Sender>, }, /// Ensure all prior writes are processed; respond when flushed. Flush { - ack: oneshot::Sender<()>, + ack: oneshot::Sender>, }, Shutdown { ack: oneshot::Sender<()>, @@ -453,7 +459,7 @@ impl RolloutRecorder { // writes. Using `tokio::fs::File` keeps everything on the async I/O // driver instead of blocking the runtime. tokio::task::spawn(rollout_writer( - file, + file.map(JsonlWriter::new), deferred_log_file_info, rx, meta, @@ -514,6 +520,7 @@ impl RolloutRecorder { .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}"))) + .and_then(|result| result) } /// Flush all queued writes and wait until they are committed by the writer task. @@ -525,6 +532,7 @@ impl RolloutRecorder { .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}"))) + .and_then(|result| result) } pub(crate) async fn load_rollout_items( @@ -647,6 +655,7 @@ fn truncate_fs_page( page } +#[derive(Clone)] struct LogFileInfo { /// Full path to the rollout file. path: PathBuf, @@ -706,7 +715,7 @@ fn open_log_file(path: &Path) -> std::io::Result { #[allow(clippy::too_many_arguments)] async fn rollout_writer( - file: Option, + file: Option, mut deferred_log_file_info: Option, mut rx: mpsc::Receiver, mut meta: Option, @@ -717,8 +726,8 @@ async fn rollout_writer( default_provider: String, generate_memories: bool, ) -> std::io::Result<()> { - let mut writer = file.map(|file| JsonlWriter { file }); - let mut buffered_items = Vec::::new(); + let mut writer = file; + let mut pending_items = Vec::::new(); if let Some(builder) = state_builder.as_mut() { builder.rollout_path = rollout_path.clone(); } @@ -749,35 +758,38 @@ async fn rollout_writer( continue; } + pending_items.extend(items); + if writer.is_none() { - buffered_items.extend(items); continue; } write_and_reconcile_items( writer.as_mut(), - items.as_slice(), + pending_items.as_slice(), &rollout_path, state_db_ctx.as_deref(), state_builder.as_ref(), default_provider.as_str(), ) - .await?; + .await + .map(|()| pending_items.clear()) + .unwrap_or_else(|err| warn!("rollout write failed; keeping writer alive: {err}")); } RolloutCmd::Persist { ack } => { - if writer.is_none() { + if writer.is_none() || meta.is_some() || !pending_items.is_empty() { 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 writer.is_none() { + let Some(log_file_info) = deferred_log_file_info.as_ref() 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::new(tokio::fs::File::from_std(file))); + } - if let Some(session_meta) = meta.take() { + if let Some(session_meta) = meta.clone() { write_session_meta( writer.as_mut(), session_meta, @@ -789,41 +801,45 @@ async fn rollout_writer( generate_memories, ) .await?; + meta = None; } - if !buffered_items.is_empty() { + if !pending_items.is_empty() { write_and_reconcile_items( writer.as_mut(), - buffered_items.as_slice(), + pending_items.as_slice(), &rollout_path, state_db_ctx.as_deref(), state_builder.as_ref(), default_provider.as_str(), ) .await?; - buffered_items.clear(); + pending_items.clear(); } + deferred_log_file_info = None; Ok(()) } .await; if let Err(err) = result { - let _ = ack.send(()); - return Err(err); + warn!("rollout persist failed; keeping writer alive: {err}"); + let _ = ack.send(Err(err)); + continue; } } - let _ = ack.send(()); + let _ = ack.send(Ok(())); } 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); + warn!("rollout flush failed; keeping writer alive: {e}"); + let _ = ack.send(Err(e)); + continue; } - let _ = ack.send(()); + let _ = ack.send(Ok(())); } RolloutCmd::Shutdown { ack } => { let _ = ack.send(()); @@ -879,9 +895,7 @@ async fn write_and_reconcile_items( default_provider: &str, ) -> std::io::Result<()> { if let Some(writer) = writer.as_mut() { - for item in items { - writer.write_rollout_item(item).await?; - } + writer.write_rollout_items(items).await?; } sync_thread_state_after_write( state_db_ctx, @@ -946,6 +960,10 @@ async fn sync_thread_state_after_write( struct JsonlWriter { file: tokio::fs::File, + #[cfg(test)] + fail_next_writes: Option>, + #[cfg(test)] + fail_next_flushes: Option>, } #[derive(serde::Serialize)] @@ -956,7 +974,87 @@ struct RolloutLineRef<'a> { } impl JsonlWriter { + fn new(file: tokio::fs::File) -> Self { + Self { + file, + #[cfg(test)] + fail_next_writes: None, + #[cfg(test)] + fail_next_flushes: None, + } + } + + #[cfg(test)] + fn with_fail_next_writes(file: tokio::fs::File, fail_next_writes: Arc) -> Self { + Self { + file, + fail_next_writes: Some(fail_next_writes), + fail_next_flushes: None, + } + } + + #[cfg(test)] + fn with_fail_next_flushes(file: tokio::fs::File, fail_next_flushes: Arc) -> Self { + Self { + file, + fail_next_writes: None, + fail_next_flushes: Some(fail_next_flushes), + } + } + async fn write_rollout_item(&mut self, rollout_item: &RolloutItem) -> std::io::Result<()> { + self.write_rollout_items(std::slice::from_ref(rollout_item)) + .await + } + + async fn write_rollout_items(&mut self, rollout_items: &[RolloutItem]) -> std::io::Result<()> { + #[cfg(test)] + if let Some(fail_next_writes) = self.fail_next_writes.as_ref() + && fail_next_writes + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| { + (remaining > 0).then(|| remaining - 1) + }) + .is_ok() + { + return Err(IoError::other("injected rollout write failure")); + } + + let file_len_before_write = self.file.metadata().await?.len(); + let mut json = String::new(); + for rollout_item in rollout_items { + json.push_str(&Self::rollout_line_json(rollout_item)?); + json.push('\n'); + } + + let result = async { + self.file.write_all(json.as_bytes()).await?; + #[cfg(test)] + if let Some(fail_next_flushes) = self.fail_next_flushes.as_ref() + && fail_next_flushes + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| { + (remaining > 0).then(|| remaining - 1) + }) + .is_ok() + { + return Err(IoError::other("injected rollout flush failure")); + } + self.file.flush().await + } + .await; + + if let Err(err) = result { + if let Err(truncate_err) = self.file.set_len(file_len_before_write).await { + return Err(IoError::other(format!( + "failed to roll back partial rollout write after {err}: {truncate_err}" + ))); + } + return Err(err); + } + + Ok(()) + } + + fn rollout_line_json(rollout_item: &RolloutItem) -> std::io::Result { let timestamp_format: &[FormatItem] = format_description!( "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z" ); @@ -968,14 +1066,7 @@ impl JsonlWriter { timestamp, item: rollout_item, }; - self.write_line(&line).await - } - async fn write_line(&mut self, item: &impl serde::Serialize) -> std::io::Result<()> { - let mut json = serde_json::to_string(item)?; - json.push('\n'); - self.file.write_all(json.as_bytes()).await?; - self.file.flush().await?; - Ok(()) + serde_json::to_string(&line).map_err(IoError::other) } } diff --git a/codex-rs/core/src/rollout/recorder_tests.rs b/codex-rs/core/src/rollout/recorder_tests.rs index dbe11ac9f7..8df9eaa033 100644 --- a/codex-rs/core/src/rollout/recorder_tests.rs +++ b/codex-rs/core/src/rollout/recorder_tests.rs @@ -3,7 +3,8 @@ use crate::config::ConfigBuilder; use crate::features::Feature; use chrono::TimeZone; use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; -use codex_protocol::protocol::AgentMessageEvent; +use codex_protocol::items::AgentMessageContent; +use codex_protocol::items::AgentMessageItem; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::SandboxPolicy; @@ -15,6 +16,9 @@ use std::fs::{self}; use std::io::Write; use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use tempfile::TempDir; use uuid::Uuid; @@ -51,6 +55,52 @@ fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result RolloutItem { + let mut events = AgentMessageItem::new(&[AgentMessageContent::Text { + text: message.to_string(), + }]) + .as_legacy_events(); + RolloutItem::EventMsg(events.pop().expect("single agent message event")) +} + +async fn queue_test_rollout_agent_message( + tx: &mpsc::Sender, + message: &str, +) -> Result<(), mpsc::error::SendError> { + tx.send(RolloutCmd::AddItems(vec![test_rollout_agent_message( + message, + )])) + .await +} + +async fn flush_test_rollout_writer(tx: &mpsc::Sender) -> std::io::Result<()> { + let (flush_tx, flush_rx) = oneshot::channel(); + tx.send(RolloutCmd::Flush { ack: flush_tx }) + .await + .map_err(|e| IoError::other(format!("flush should queue: {e}")))?; + flush_rx + .await + .map_err(|e| IoError::other(format!("flush ack should be sent: {e}")))? +} + +async fn wait_for_test_fault_to_fire(fault: &AtomicUsize) { + tokio::time::timeout(Duration::from_secs(5), async { + while fault.load(Ordering::Acquire) != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("injected rollout fault should be consumed"); +} + +fn assert_rollout_message_count(text: &str, message: &str, expected: usize) { + assert_eq!( + text.matches(message).count(), + expected, + "unexpected rollout count for {message:?}" + ); +} + #[tokio::test] async fn recorder_materializes_only_after_explicit_persist() -> std::io::Result<()> { let home = TempDir::new().expect("temp dir"); @@ -81,13 +131,7 @@ async fn recorder_materializes_only_after_explicit_persist() -> std::io::Result< ); recorder - .record_items(&[RolloutItem::EventMsg(EventMsg::AgentMessage( - AgentMessageEvent { - message: "buffered-event".to_string(), - phase: None, - memory_citation: None, - }, - ))]) + .record_items(&[test_rollout_agent_message("buffered-event")]) .await?; recorder.flush().await?; assert!( @@ -138,6 +182,118 @@ async fn recorder_materializes_only_after_explicit_persist() -> std::io::Result< Ok(()) } +#[tokio::test] +async fn rollout_writer_recovers_after_transient_write_failure() -> std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let rollout_path = home.path().join("rollout.jsonl"); + let fail_next_writes = Arc::new(AtomicUsize::new(1)); + + let file = tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&rollout_path) + .await?; + let writer = JsonlWriter::with_fail_next_writes(file, fail_next_writes.clone()); + let (tx, rx) = mpsc::channel::(256); + let writer_task = tokio::spawn(rollout_writer( + Some(writer), + None, + rx, + None, + home.path().to_path_buf(), + rollout_path.clone(), + None, + None, + "test-provider".to_string(), + false, + )); + + queue_test_rollout_agent_message(&tx, "first-write-fails") + .await + .expect("first write should queue"); + + wait_for_test_fault_to_fire(&fail_next_writes).await; + + let _ = queue_test_rollout_agent_message(&tx, "second-write-succeeds").await; + let _ = flush_test_rollout_writer(&tx).await; + + drop(tx); + + let text = std::fs::read_to_string(&rollout_path)?; + assert!( + text.contains("second-write-succeeds"), + "expected the message sent after a failure to show up in the rollout" + ); + assert!( + text.contains("first-write-fails"), + "expected the message from the failed batch to be retried into the rollout" + ); + + writer_task + .await + .expect("writer task should join cleanly") + .expect("writer task should exit cleanly after channel closes"); + + Ok(()) +} + +#[tokio::test] +async fn rollout_writer_rolls_back_partial_batch_and_retries_without_duplicates() +-> std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let rollout_path = home.path().join("rollout.jsonl"); + let fail_next_flushes = Arc::new(AtomicUsize::new(1)); + + let file = tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&rollout_path) + .await?; + let writer = JsonlWriter::with_fail_next_flushes(file, fail_next_flushes.clone()); + let (tx, rx) = mpsc::channel::(256); + let writer_task = tokio::spawn(rollout_writer( + Some(writer), + None, + rx, + None, + home.path().to_path_buf(), + rollout_path.clone(), + None, + None, + "test-provider".to_string(), + false, + )); + + tx.send(RolloutCmd::AddItems(vec![ + test_rollout_agent_message("batch-item-1"), + test_rollout_agent_message("batch-item-2"), + ])) + .await + .expect("failed batch should queue"); + wait_for_test_fault_to_fire(&fail_next_flushes).await; + + queue_test_rollout_agent_message(&tx, "batch-item-3") + .await + .expect("retry trigger should queue"); + flush_test_rollout_writer(&tx).await?; + drop(tx); + + let text = std::fs::read_to_string(&rollout_path)?; + assert_rollout_message_count(&text, "batch-item-1", 1); + assert_rollout_message_count(&text, "batch-item-2", 1); + assert_rollout_message_count(&text, "batch-item-3", 1); + for line in text.lines() { + serde_json::from_str::(line).expect("rollout line should be valid JSON"); + } + + writer_task + .await + .expect("writer task should join cleanly") + .expect("writer task should exit cleanly after channel closes"); + + Ok(()) +} + #[tokio::test] async fn metadata_irrelevant_events_touch_state_db_updated_at() -> std::io::Result<()> { let home = TempDir::new().expect("temp dir"); @@ -198,13 +354,7 @@ async fn metadata_irrelevant_events_touch_state_db_updated_at() -> std::io::Resu tokio::time::sleep(Duration::from_secs(1)).await; recorder - .record_items(&[RolloutItem::EventMsg(EventMsg::AgentMessage( - AgentMessageEvent { - message: "assistant text".to_string(), - phase: None, - memory_citation: None, - }, - ))]) + .record_items(&[test_rollout_agent_message("assistant text")]) .await?; recorder.flush().await?; @@ -249,13 +399,7 @@ async fn metadata_irrelevant_events_fall_back_to_upsert_when_thread_missing() -> Utc::now(), SessionSource::Cli, ); - let items = vec![RolloutItem::EventMsg(EventMsg::AgentMessage( - AgentMessageEvent { - message: "assistant text".to_string(), - phase: None, - memory_citation: None, - }, - ))]; + let items = vec![test_rollout_agent_message("assistant text")]; sync_thread_state_after_write( Some(state_db.as_ref()),