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.
This commit is contained in:
jif-oai
2026-04-10 14:12:33 +01:00
committed by GitHub
Unverified
parent 085ffb4456
commit 8035cb03f1
10 changed files with 536 additions and 191 deletions
+2 -2
View File
@@ -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 {
+18 -3
View File
@@ -194,7 +194,12 @@ async fn persist_thread_for_tree_resume(thread: &Arc<CodexThread>, 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
+22 -11
View File
@@ -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::<Vec<_>>();
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,
+11 -4
View File
@@ -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<Session>) -> 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
+2 -2
View File
@@ -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(
+1 -1
View File
@@ -564,7 +564,7 @@ async fn append_guardian_followup_reminder(review_session: &GuardianReviewSessio
async fn load_rollout_items_for_fork(
session: &Session,
) -> anyhow::Result<Option<Vec<RolloutItem>>> {
session.flush_rollout().await;
session.flush_rollout().await?;
let Some(rollout_path) = session.current_rollout_path().await else {
return Ok(None);
};
+6 -1
View File
@@ -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");
+16 -2
View File
@@ -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