mirror of
https://github.com/openai/codex.git
synced 2026-09-13 11:47:17 +00:00
[WIP] fix: rollout recorder not adding turns to rollout after an error until next restart
I have been chasing this bug for quite some time. I noticed that sometimes when I closed the Codex app, or I tried to resume a CLI thread, that a bunch of messages were missing. Turns out that this can happen if the rollout recorder hits an error (for example if you’re out of disk space). After that happens, it will stop writing to disk. I’m not confident about the Codex-written fix, but I’m sharing it as-is at least to highlight the problem and start the conversation!
This commit is contained in:
@@ -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<RolloutItem>),
|
||||
Persist {
|
||||
ack: oneshot::Sender<()>,
|
||||
ack: oneshot::Sender<std::io::Result<()>>,
|
||||
},
|
||||
/// Ensure all prior writes are processed; respond when flushed.
|
||||
Flush {
|
||||
ack: oneshot::Sender<()>,
|
||||
ack: oneshot::Sender<std::io::Result<()>>,
|
||||
},
|
||||
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<File> {
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn rollout_writer(
|
||||
file: Option<tokio::fs::File>,
|
||||
file: Option<JsonlWriter>,
|
||||
mut deferred_log_file_info: Option<LogFileInfo>,
|
||||
mut rx: mpsc::Receiver<RolloutCmd>,
|
||||
mut meta: Option<SessionMeta>,
|
||||
@@ -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::<RolloutItem>::new();
|
||||
let mut writer = file;
|
||||
let mut pending_items = Vec::<RolloutItem>::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<Arc<AtomicUsize>>,
|
||||
#[cfg(test)]
|
||||
fail_next_flushes: Option<Arc<AtomicUsize>>,
|
||||
}
|
||||
|
||||
#[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<AtomicUsize>) -> 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<AtomicUsize>) -> 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<String> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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<Path
|
||||
Ok(path)
|
||||
}
|
||||
|
||||
fn test_rollout_agent_message(message: &str) -> 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<RolloutCmd>,
|
||||
message: &str,
|
||||
) -> Result<(), mpsc::error::SendError<RolloutCmd>> {
|
||||
tx.send(RolloutCmd::AddItems(vec![test_rollout_agent_message(
|
||||
message,
|
||||
)]))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn flush_test_rollout_writer(tx: &mpsc::Sender<RolloutCmd>) -> 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::<RolloutCmd>(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::<RolloutCmd>(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::<serde_json::Value>(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()),
|
||||
|
||||
Reference in New Issue
Block a user