diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 77ce9b873c..801431fc23 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -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>; +pub(crate) type SessionTaskResult = CodexResult>; #[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, turn_context: Arc, task_result: SessionTaskResult, - ) { - let (last_agent_message, abort_reason) = match task_result { - Ok(last_agent_message) => (last_agent_message, None), + ) -> Option { + 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 { diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 0a05faa8ee..718acfd862 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -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, + 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, + ctx: Arc, + mut next_input: Vec, + mut mcp_startup_requirements: McpStartupRequirements, + mut prewarmed_client_session: Option, + 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(); + } + }) } } diff --git a/codex-rs/core/tests/suite/pending_input_persistence.rs b/codex-rs/core/tests/suite/pending_input_persistence.rs index 0ebb340663..c7dc52e36f 100644 --- a/codex-rs/core/tests/suite/pending_input_persistence.rs +++ b/codex-rs/core/tests/suite/pending_input_persistence.rs @@ -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, checkpoint_context: Option, checkpoints: mpsc::UnboundedSender, + next_flush: Mutex>>, } 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)