openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-06 20:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Handle steers accepted during the final transcript flush (#58877)

Browse Source 
## Why

Input accepted while a regular turn was flushing its transcript could miss the final pending-input check and remain unprocessed when the turn completed.

## What changed

Check pending input atomically with closing steer admission. If input was accepted, resume sampling in the same turn while preserving MCP startup requirements, and complete the turn with the updated answer.

## Testing

Add a regression test that pauses the final flush, submits a steer, and verifies that the model receives it exactly once before the same turn completes with the answer to the steer.

GitOrigin-RevId: 193db3ab6916a86e97c3845590598fd8fbe47081
This commit is contained in:
Charlie Marsh
2026-11-02 12:24:38 +00:00
committed by copyberry
parent 053e5b1ca9
commit ba0d0e1f74
3 changed files with 263 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tasks/mod.rs
+73 -43
View File
@@ -31,6 +31,7 @@ use crate::context::ContextualUserFragment;
use crate::hook_runtime::run_turn_interrupt_hooks;
use crate::session::TurnInput;
use crate::session::session::Session;
use crate::session::turn::McpStartupRequirements;
use crate::session::turn::run_hooks_and_record_inputs;
use crate::session::turn_context::NewTurnContextOptions;
use crate::session::turn_context::TurnContext;
@@ -73,7 +74,7 @@ pub(crate) const GRACEFULL_INTERRUPTION_TIMEOUT_MS: u64 = 100;
const TASK_COMPACT_METRIC: &str = "codex.task.compact";
static ACTIVE_TURNS: Gauge = Gauge::new("core.turns.active");
pub(crate) type SessionTaskResult = CodexResult<Option<String>>;
pub(crate) type SessionTaskResult = CodexResult<Option<regular::RegularTaskOutput>>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum InterruptedTurnHistoryMarker {
@@ -189,11 +190,11 @@ pub(crate) trait SessionTask: Send + Sync + 'static {
/// Executes the task until completion or cancellation.
///
/// Implementations typically stream protocol events using `session` and
/// `ctx`, returning an optional final agent message when finished. The
/// `ctx`, returning the sampling state for a regular turn when finished. The
/// provided `cancellation_token` is cancelled when the session requests an
/// abort; implementers should watch for it and terminate quickly once it
/// fires. Returning [`Some`] yields a final message that
/// [`Session::on_task_finished`] will emit to the client. Returning
/// fires. Returning [`Some`] lets finalization continue sampling
/// if this regular turn accepted more input. Returning
/// [`CodexErr::TurnAborted`] completes the task through the aborted-turn
/// lifecycle instead.
fn run(
@@ -384,7 +385,7 @@ impl Session {
let storage_originator = AuthStorageOriginator::from_client_name(&turn_context.originator);
let task_future = async move {
let ctx_for_finish = Arc::clone(&ctx);
let task_result = task_for_run
let mut task_result = task_for_run
.run(
Arc::clone(&session),
ctx,
@@ -393,30 +394,44 @@ impl Session {
)
.instrument(trace_span!("session_task.run"))
.await;
let sess = Arc::clone(&session);
// Private reviewers save their transcript together with the terminal event.
// Errors and cancellation retain their existing save path.
if (!sess.is_private_guardian_reviewer().await
|| task_cancellation_token.is_cancelled()
|| task_result.is_err())
&& 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}"
),
}),
loop {
// Private reviewers save their transcript together with the terminal event.
// Errors and cancellation retain their existing save path.
if (!session.is_private_guardian_reviewer().await
|| task_cancellation_token.is_cancelled()
|| task_result.is_err())
&& let Err(err) = session.flush_rollout().await
{
warn!("failed to flush rollout before completing turn: {err}");
session.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() {
break;
}
let Some(requirements) = session
.on_task_finished(Arc::clone(&ctx_for_finish), task_result)
.await
else {
break;
};
task_result = RegularTask::run_turns(
Arc::clone(&session),
Arc::clone(&ctx_for_finish),
Vec::new(),
requirements,
/*prewarmed_client_session*/ None,
task_cancellation_token.child_token(),
)
.await;
}
if !task_cancellation_token.is_cancelled() {
// Finish uniformly from the spawn site so all tasks share the same lifecycle.
sess.on_task_finished(Arc::clone(&ctx_for_finish), task_result)
.await;
}
done_clone.notify_waiters();
}
.instrument(task_span);
@@ -652,13 +667,18 @@ impl Session {
}
}
pub async fn on_task_finished(
/// Closes admission, or returns the regular turn's sampling state if a steer won the race.
#[expect(
clippy::await_holding_invalid_type,
reason = "the final input check and task removal must be atomic with steer admission"
)]
pub(crate) async fn on_task_finished(
self: &Arc<Self>,
turn_context: Arc<TurnContext>,
task_result: SessionTaskResult,
) {
let (last_agent_message, abort_reason) = match task_result {
Ok(last_agent_message) => (last_agent_message, None),
) -> Option<McpStartupRequirements> {
let (regular_output, abort_reason) = match task_result {
Ok(output) => (output, None),
Err(err) if matches!(err.details(), CodexErrorDetails::TurnAborted) => {
(None, Some(TurnAbortReason::Interrupted))
}
@@ -679,24 +699,33 @@ impl Session {
(None, None)
}
};
turn_context
.turn_metadata_state
.cancel_git_enrichment_task();
let can_continue =
regular_output.is_some() && turn_context.terminal_error.lock().await.is_none();
let finishing = CancellationToken::new();
let finish_recording = finishing.clone().drop_guard();
let turn_state = {
let mut active = self.active_turn.lock().await;
active.as_mut().and_then(|active_turn| {
let task = active_turn.task.take()?;
task.handle.detach();
active_turn.finishing = Some(finishing);
Some(Arc::clone(&active_turn.turn_state))
})
};
let Some(turn_state) = turn_state else {
return;
let active_turn = active.as_mut()?;
if !Arc::ptr_eq(&active_turn.task.as_ref()?.turn_context, &turn_context) {
return None;
}
if can_continue {
let state = active_turn.turn_state.lock().await;
if state.accepts_mailbox_delivery_for_current_turn()
&& !state.pending_input.is_empty()
{
return regular_output.map(|output| output.mcp_startup_requirements);
}
}
let task = active_turn.task.take()?;
task.handle.detach();
active_turn.finishing = Some(finishing);
Arc::clone(&active_turn.turn_state)
};
turn_context
.turn_metadata_state
.cancel_git_enrichment_task();
let last_agent_message = regular_output.and_then(|output| output.last_agent_message);
let pending_input = self
.input_queue
.take_pending_input_for_turn_state(turn_state.as_ref())
@@ -935,6 +964,7 @@ impl Session {
if cleared_active_turn {
self.maybe_start_turn_for_pending_work().await;
}
None
}
async fn take_active_turn(&self, reason: &TurnAbortReason) -> Option<ActiveTurn> {
codex-rs/core/src/tasks/regular.rs
+60 -29
View File
@@ -2,8 +2,10 @@ use std::sync::Arc;
use codex_async_utils::OrCancelExt;
use codex_extension_api::TurnStartPhase;
use futures::future::BoxFuture;
use tokio_util::sync::CancellationToken;
use crate::client::ModelClientSession;
use crate::session::TurnInput;
use crate::session::session::Session;
use crate::session::startup_prewarm::SessionStartupPrewarmResolution;
@@ -19,6 +21,12 @@ use tracing::trace_span;
use super::SessionTask;
use super::SessionTaskResult;
/// The last answer and the MCP requirements needed if a steer reopens sampling.
pub(crate) struct RegularTaskOutput {
pub(super) last_agent_message: Option<String>,
pub(super) mcp_startup_requirements: McpStartupRequirements,
}
#[derive(Default)]
pub(crate) struct RegularTask;
@@ -52,11 +60,6 @@ impl SessionTask for RegularTask {
{
ctx.turn_timing_state.enable_critical_path();
}
let run_turn_span = trace_span!(
"run_turn",
conversation.id = %sess.thread_id,
turn.id = %ctx.sub_id,
);
// Regular turns emit `TurnStarted` inline so first-turn lifecycle does
// not wait on startup prewarm resolution.
let prewarmed_client_session = async {
@@ -109,29 +112,57 @@ impl SessionTask for RegularTask {
Some(*prewarmed_client_session)
}
};
let mut next_input = input;
let mut prewarmed_client_session = prewarmed_client_session;
let mut mcp_startup_requirements = McpStartupRequirements::default();
loop {
let last_agent_message = run_turn(
Arc::clone(&sess),
Arc::clone(&ctx),
next_input,
&mut mcp_startup_requirements,
prewarmed_client_session.take(),
cancellation_token.child_token(),
)
.instrument(run_turn_span.clone())
.await?;
// Terminal errors are already reported. Let task completion preserve pending
// input instead of restarting the failed turn for that same input.
if ctx.terminal_error.lock().await.is_some() {
return Ok(last_agent_message);
}
if !sess.input_queue.has_pending_input(&sess.active_turn).await {
return Ok(last_agent_message);
}
next_input = Vec::new();
}
Self::run_turns(
sess,
ctx,
input,
McpStartupRequirements::default(),
prewarmed_client_session,
cancellation_token,
)
.await
}
}
impl RegularTask {
/// Runs model requests for a regular turn, retaining explicit MCP requirements across steers.
pub(super) fn run_turns(
sess: Arc<Session>,
ctx: Arc<TurnContext>,
mut next_input: Vec<TurnInput>,
mut mcp_startup_requirements: McpStartupRequirements,
mut prewarmed_client_session: Option<ModelClientSession>,
cancellation_token: CancellationToken,
) -> BoxFuture<'static, SessionTaskResult> {
Box::pin(async move {
let run_turn_span = trace_span!(
"run_turn",
conversation.id = %sess.thread_id,
turn.id = %ctx.sub_id,
);
loop {
let last_agent_message = run_turn(
Arc::clone(&sess),
Arc::clone(&ctx),
next_input,
&mut mcp_startup_requirements,
prewarmed_client_session.take(),
cancellation_token.child_token(),
)
.instrument(run_turn_span.clone())
.await?;
// Terminal errors are already reported. Let task completion preserve pending
// input instead of restarting the failed turn for that same input.
if ctx.terminal_error.lock().await.is_some()
|| !sess.input_queue.has_pending_input(&sess.active_turn).await
{
return Ok(Some(RegularTaskOutput {
last_agent_message,
mcp_startup_requirements,
}));
}
next_input = Vec::new();
}
})
}
}
codex-rs/core/tests/suite/pending_input_persistence.rs
+130 -2
View File
@@ -1,4 +1,4 @@
//! Exercises preparation and execution checkpoints before model requests.
//! Exercises how transcript checkpoints interact with model requests and turn completion.
use std::any::Any;
use std::sync::Arc;
@@ -49,6 +49,8 @@ use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use serde_json::json;
use test_case::test_case;
use tokio::sync::Barrier;
use tokio::sync::Mutex;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::time::timeout;
@@ -93,6 +95,7 @@ struct GatedCheckpointStore {
arm_on: Option<PersistContext>,
checkpoint_context: Option<PersistContext>,
checkpoints: mpsc::UnboundedSender<PendingCheckpoint>,
next_flush: Mutex<Option<Arc<Barrier>>>,
}
macro_rules! delegate_store_methods {
@@ -120,10 +123,21 @@ impl ThreadStore for GatedCheckpointStore {
fn archive_thread(params: ArchiveThreadParams) -> ();
fn unarchive_thread(params: ArchiveThreadParams) -> StoredThread;
fn delete_thread(params: DeleteThreadParams) -> ();
fn flush_thread(thread_id: ThreadId) -> ();
fn shutdown_thread(thread_id: ThreadId) -> ();
}
fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
Box::pin(async move {
let gate = self.next_flush.lock().await.take();
if let Some(gate) = gate {
// Let the test submit input while this flush is still in progress.
gate.wait().await;
gate.wait().await;
}
self.inner.flush_thread(thread_id).await
})
}
fn update_thread_metadata(
&self,
params: UpdateThreadMetadataParams,
@@ -228,6 +242,7 @@ async fn steered_input_checkpoint_controls_next_request(
arm_on: None,
checkpoint_context: None,
checkpoints,
next_flush: Mutex::default(),
});
let base_url = format!("{}/v1", server.uri());
let config_server = responses::start_mock_server().await;
@@ -348,6 +363,117 @@ async fn steered_input_checkpoint_controls_next_request(
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn steer_during_final_flush_reaches_the_model_before_turn_completion() -> anyhow::Result<()> {
let server = responses::start_mock_server().await;
let responses = responses::mount_sse_sequence(
&server,
[
("first", "original answer"),
("second", "answer to the steer"),
]
.into_iter()
.map(|(id, answer)| {
responses::sse(vec![
responses::ev_response_created(id),
responses::ev_assistant_message(id, answer),
responses::ev_completed(id),
])
})
.collect(),
)
.await;
let (checkpoints, _checkpoint_requests) = mpsc::unbounded_channel();
let store = Arc::new(GatedCheckpointStore {
inner: InMemoryThreadStore::default(),
policy: CheckpointPolicy::Synchronous,
armed: AtomicBool::new(false),
arm_on: None,
checkpoint_context: None,
checkpoints,
next_flush: Mutex::default(),
});
let test = test_codex()
.with_thread_store(store.clone())
.with_history_mode(ThreadHistoryMode::Legacy)
.build_with_auto_env(&server)
.await?;
let user_message = |text: &str, client_id| {
TurnInputRequest::new(TurnInput::UserInput {
content: vec![UserInput::Text {
text: text.to_string(),
text_elements: Vec::new(),
}],
client_id,
})
};
let flush = Arc::new(Barrier::new(/*n*/ 2));
*store.next_flush.lock().await = Some(flush.clone());
let TurnInputSubmission::Started {
turn_id,
root_turn_id,
} = test
.codex
.start_or_steer_turn(user_message("first prompt", None))
.await?
else {
panic!("expected a new turn");
};
// The regular task returned its first answer and is flushing before it completes the turn.
timeout(Duration::from_secs(10), flush.wait()).await?;
assert_eq!(responses.requests().len(), 1);
let steer = timeout(
Duration::from_secs(10),
test.codex
.start_or_steer_turn(user_message("late steer", Some("steer-id".to_string()))),
)
.await;
flush.wait().await;
assert_eq!(
steer??,
TurnInputSubmission::Steered {
turn_id: turn_id.clone(),
root_turn_id
}
);
let mut starts = 0;
let mut client_ids = Vec::new();
let EventMsg::TurnComplete(completed) = wait_for_event(&test.codex, |event| {
match event {
EventMsg::TurnStarted(_) => starts += 1,
EventMsg::UserMessage(message) if message.message == "late steer" => {
client_ids.push(message.client_id.clone());
}
_ => {}
}
matches!(event, EventMsg::TurnComplete(_))
})
.await
else {
unreachable!();
};
assert_eq!(completed.turn_id, turn_id);
assert_eq!(
completed.last_agent_message.as_deref(),
Some("answer to the steer")
);
assert_eq!(starts, 1);
assert_eq!(client_ids, [Some("steer-id".to_string())]);
let requests = responses.requests();
assert_eq!(requests.len(), 2);
assert_eq!(
requests[1]
.message_input_texts("user")
.iter()
.filter(|text| *text == "late steer")
.count(),
1
);
Ok(())
}
#[test_case(FirstInputKind::User; "user_input")]
#[test_case(FirstInputKind::InterAgentCommunication; "inter_agent_communication")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -371,6 +497,7 @@ async fn preparation_and_first_sampling_do_not_wait_for_durable_metadata(
arm_on: None,
checkpoint_context: None,
checkpoints,
next_flush: Mutex::default(),
});
let base_url = format!("{}/v1", server.uri());
let config_server = responses::start_mock_server().await;
@@ -614,6 +741,7 @@ async fn cancelled_spawn_discards_provisional_child(caller: SpawnCaller) -> anyh
arm_on: Some(PersistContext::SubagentSpawn),
checkpoint_context: Some(PersistContext::Standard),
checkpoints,
next_flush: Mutex::default(),
});
let test = test_codex()
.with_thread_store(store)