mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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:
committed by
GitHub
Unverified
parent
085ffb4456
commit
8035cb03f1
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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);
|
||||
};
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user