From 13385391c111df0ac7ae01369c527aa61c0ec5b2 Mon Sep 17 00:00:00 2001 From: Winston Howes Date: Mon, 6 Jul 2026 19:43:51 -0700 Subject: [PATCH] core: serialize thread turn transitions --- codex-rs/core/src/session/inject.rs | 14 ++- codex-rs/core/src/session/session.rs | 2 + codex-rs/core/src/session/tests.rs | 123 +++++++++++++++++++++++++++ codex-rs/core/src/tasks/mod.rs | 69 ++++++++++++--- 4 files changed, 191 insertions(+), 17 deletions(-) diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index 2ad72857571b..80f7b491a2a3 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -22,7 +22,7 @@ impl Session { ) -> Result<(), Vec> { let mut active = self.active_turn.lock().await; match active.as_mut() { - Some(active_turn) => { + Some(active_turn) if active_turn.task.is_some() => { self.input_queue .extend_pending_input_for_turn_state( active_turn.turn_state.as_ref(), @@ -31,7 +31,7 @@ impl Session { .await; Ok(()) } - None => Err(input), + Some(_) | None => Err(input), } } @@ -42,6 +42,10 @@ impl Session { /// start a turn when user/client-triggered work is queued, any task is still /// active, or the session is currently in Plan mode. Active Review tasks are /// covered by the active-task check because Review turns are not steerable. + #[expect( + clippy::await_holding_invalid_type, + reason = "turn startup must remain serialized across its async checks" + )] pub(crate) async fn try_start_turn_if_idle( self: &Arc, input: Vec, @@ -49,6 +53,7 @@ impl Session { if input.is_empty() { return Ok(()); } + let turn_start = self.turn_start_lock.lock().await; if self.input_queue.has_trigger_turn_mailbox_items().await { return Err(TryStartTurnIfIdleError::new( TryStartTurnIfIdleRejectionReason::PendingTriggerTurn, @@ -83,6 +88,7 @@ impl Session { if self.input_queue.has_trigger_turn_mailbox_items().await { self.clear_reserved_idle_turn(&turn_state).await; + drop(turn_start); self.maybe_start_turn_for_pending_work().await; return Err(TryStartTurnIfIdleError::new( TryStartTurnIfIdleRejectionReason::PendingTriggerTurn, @@ -95,6 +101,7 @@ impl Session { .await; if turn_context.collaboration_mode.mode == ModeKind::Plan { self.clear_reserved_idle_turn(&turn_state).await; + drop(turn_start); self.maybe_start_turn_for_pending_work().await; return Err(TryStartTurnIfIdleError::new( TryStartTurnIfIdleRejectionReason::PlanMode, @@ -105,6 +112,7 @@ impl Session { .await; if self.input_queue.has_trigger_turn_mailbox_items().await { self.clear_reserved_idle_turn(&turn_state).await; + drop(turn_start); self.maybe_start_turn_for_pending_work().await; return Err(TryStartTurnIfIdleError::new( TryStartTurnIfIdleRejectionReason::PendingTriggerTurn, @@ -131,7 +139,7 @@ impl Session { input.into_iter().map(TurnInput::ResponseItem).collect(), ) .await; - self.start_task(turn_context, Vec::new(), RegularTask::new()) + self.start_task_locked(turn_context, Vec::new(), RegularTask::new()) .await; Ok(()) } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index a7bb9b233d6e..89db411437dc 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -174,6 +174,7 @@ pub(crate) struct Session { pub(crate) submission_lifecycle: SubmissionLifecycle, pub(crate) thread_activity: ThreadActivityHandle, pub(crate) submission_send_lock: Mutex<()>, + pub(crate) turn_start_lock: Mutex<()>, pub(crate) input_queue: InputQueue, pub(crate) guardian_review_session: GuardianReviewSessionManager, pub(crate) services: SessionServices, @@ -1301,6 +1302,7 @@ impl Session { submission_lifecycle: SubmissionLifecycle::default(), thread_activity, submission_send_lock: Mutex::new(()), + turn_start_lock: Mutex::new(()), input_queue: InputQueue::new(), guardian_review_session: GuardianReviewSessionManager::default(), services, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 7560045c409c..e7911e4b46bf 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5526,6 +5526,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { submission_lifecycle: SubmissionLifecycle::default(), thread_activity, submission_send_lock: Mutex::new(()), + turn_start_lock: Mutex::new(()), input_queue: super::input_queue::InputQueue::new(), guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, @@ -7807,6 +7808,7 @@ where submission_lifecycle: SubmissionLifecycle::default(), thread_activity, submission_send_lock: Mutex::new(()), + turn_start_lock: Mutex::new(()), input_queue: super::input_queue::InputQueue::new(), guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, @@ -10045,6 +10047,88 @@ async fn task_finish_emits_thread_idle_lifecycle_after_active_turn_clears() { assert!(session.active_turn.lock().await.is_none()); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn task_finish_retries_trigger_mail_that_arrives_before_active_turn_clears() { + struct BlockingTurnStop { + blocked: std::sync::atomic::AtomicBool, + entered: async_channel::Sender<()>, + release: async_channel::Receiver<()>, + } + + impl codex_extension_api::TurnLifecycleContributor for BlockingTurnStop { + fn on_turn_stop<'a>( + &'a self, + _input: codex_extension_api::TurnStopInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + if self.blocked.swap(true, std::sync::atomic::Ordering::SeqCst) { + return; + } + self.entered.send(()).await.expect("entered receiver open"); + self.release.recv().await.expect("release sender open"); + }) + } + } + + let (mut session, turn_context, events) = make_session_and_context_with_rx().await; + let (entered_tx, entered_rx) = async_channel::bounded(1); + let (release_tx, release_rx) = async_channel::bounded(1); + let mut builder = codex_extension_api::ExtensionRegistryBuilder::::new(); + builder.turn_lifecycle_contributor(Arc::new(BlockingTurnStop { + blocked: std::sync::atomic::AtomicBool::new(false), + entered: entered_tx, + release: release_rx, + })); + Arc::get_mut(&mut session) + .expect("session should be uniquely owned") + .services + .extensions = Arc::new(builder.build()); + session + .spawn_task(Arc::clone(&turn_context), Vec::new(), CompletingTask) + .await; + entered_rx.recv().await.expect("turn stop should block"); + assert!(session.turn_start_lock.try_lock().is_err()); + + session + .input_queue + .enqueue_mailbox_communication(InterAgentCommunication::new( + AgentPath::root(), + AgentPath::root(), + Vec::new(), + "late trigger".to_string(), + /*trigger_turn*/ true, + )) + .await; + let scheduler = tokio::spawn({ + let session = Arc::clone(&session); + async move { session.maybe_start_turn_for_pending_work().await } + }); + release_tx.send(()).await.expect("turn stop should resume"); + scheduler + .await + .expect("pending-work scheduler should not panic"); + + let next_turn_id = timeout(StdDuration::from_secs(2), async { + let mut saw_completed_turn = false; + loop { + match events.recv().await.expect("event channel open").msg { + EventMsg::TurnComplete(event) if event.turn_id == turn_context.sub_id => { + saw_completed_turn = true; + } + EventMsg::TurnStarted(event) => { + assert!(saw_completed_turn); + return event.turn_id; + } + _ => {} + } + } + }) + .await + .expect("late trigger should start a new turn"); + assert_ne!(next_turn_id, turn_context.sub_id); + session.abort_all_tasks(TurnAbortReason::Interrupted).await; +} + #[tokio::test] async fn thread_idle_lifecycle_waits_for_trigger_turn_mailbox_work() { struct ThreadIdleRecorder { @@ -10114,6 +10198,45 @@ async fn try_start_turn_if_idle_rejects_active_turn_without_injecting() { sess.abort_all_tasks(TurnAbortReason::Interrupted).await; } +#[tokio::test] +#[expect( + clippy::await_holding_invalid_type, + reason = "the test holds session state to pause startup at a deterministic boundary" +)] +async fn interrupt_waits_for_idle_turn_start_to_install_task() { + let (sess, _tc, _rx) = make_session_and_context_with_rx().await; + let state = sess.state.lock().await; + let start = tokio::spawn({ + let sess = Arc::clone(&sess); + async move { + sess.try_start_turn_if_idle(vec![user_message("synthetic idle input")]) + .await + } + }); + loop { + match sess.turn_start_lock.try_lock() { + Ok(turn_start) => drop(turn_start), + Err(_) => break, + } + tokio::task::yield_now().await; + } + + let interrupt = tokio::spawn({ + let sess = Arc::clone(&sess); + async move { sess.interrupt_task().await } + }); + tokio::task::yield_now().await; + assert!(!interrupt.is_finished()); + + drop(state); + start + .await + .expect("idle turn start task should not panic") + .expect("idle turn should be accepted"); + interrupt.await.expect("interrupt task should not panic"); + assert!(sess.active_turn.lock().await.is_none()); +} + #[tokio::test] async fn idle_shutdown_claim_waits_for_submissions_and_blocks_automatic_turns() { let (sess, _tc, _rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index a7145f8db852..2dda4b2e29f5 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -311,18 +311,23 @@ where } impl Session { + #[expect( + clippy::await_holding_invalid_type, + reason = "task replacement must remain serialized through startup" + )] pub async fn spawn_task( self: &Arc, turn_context: Arc, input: Vec, task: T, ) { - self.abort_all_tasks(TurnAbortReason::Replaced).await; + let _turn_start = self.turn_start_lock.lock().await; + self.abort_all_tasks_locked(TurnAbortReason::Replaced).await; self.clear_connector_selection().await; - self.start_task(turn_context, input, task).await; + self.start_task_locked(turn_context, input, task).await; } - pub(crate) async fn start_task( + pub(crate) async fn start_task_locked( self: &Arc, turn_context: Arc, input: Vec, @@ -425,8 +430,10 @@ impl Session { } 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; + select! { + _ = task_cancellation_token.cancelled() => {} + _ = sess.on_task_finished(Arc::clone(&ctx_for_finish), task_result) => {} + } } done_clone.notify_waiters(); } @@ -456,9 +463,11 @@ impl Session { /// /// This helper generates a fresh sub-id for the synthetic turn before delegating to the /// explicit-sub-id variant. - pub(crate) async fn maybe_start_turn_for_pending_work(self: &Arc) { - self.maybe_start_turn_for_pending_work_with_sub_id(uuid::Uuid::new_v4().to_string()) - .await; + pub(crate) fn maybe_start_turn_for_pending_work(self: &Arc) -> BoxFuture<'_, ()> { + Box::pin(async move { + self.maybe_start_turn_for_pending_work_with_sub_id(uuid::Uuid::new_v4().to_string()) + .await; + }) } /// Starts a regular turn with the provided sub-id when pending work should wake an idle @@ -466,6 +475,10 @@ impl Session { /// /// The turn is created only when there is mailbox mail marked with `trigger_turn`, and only /// if the session is currently idle. + #[expect( + clippy::await_holding_invalid_type, + reason = "pending-work startup must remain serialized across its async checks" + )] pub(crate) async fn maybe_start_turn_for_pending_work_with_sub_id( self: &Arc, sub_id: String, @@ -473,6 +486,10 @@ impl Session { if !self.input_queue.has_trigger_turn_mailbox_items().await { return; } + let _turn_start = self.turn_start_lock.lock().await; + if !self.input_queue.has_trigger_turn_mailbox_items().await { + return; + } let Some(reservation) = self.try_reserve_activity() else { return; @@ -489,11 +506,24 @@ impl Session { let turn_context = self.new_default_turn_with_sub_id(sub_id).await; self.maybe_emit_model_warnings_for_turn(turn_context.as_ref()) .await; - self.start_task(turn_context, Vec::new(), RegularTask::new()) + self.start_task_locked(turn_context, Vec::new(), RegularTask::new()) .await; } + #[expect( + clippy::await_holding_invalid_type, + reason = "abort and optional restart form one serialized transition" + )] pub async fn abort_all_tasks(self: &Arc, reason: TurnAbortReason) { + let turn_start = self.turn_start_lock.lock().await; + let restart_pending = self.abort_all_tasks_locked(reason).await; + drop(turn_start); + if restart_pending { + self.maybe_start_turn_for_pending_work().await; + } + } + + async fn abort_all_tasks_locked(self: &Arc, reason: TurnAbortReason) -> bool { let mut aborted_turn = false; let mut active_turn_to_clear = None; let mut turn_context = None; @@ -513,21 +543,24 @@ impl Session { self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref()) .await; } - if let Some(active_turn) = active_turn_to_clear { + if let Some(active_turn) = active_turn_to_clear.as_ref() { // Let interrupted tasks observe cancellation before dropping pending approvals, or an // in-flight approval wait can surface as a model-visible rejection before TurnAborted. - self.input_queue.clear_pending(&active_turn).await; - } - if reason == TurnAbortReason::Interrupted && aborted_turn { - self.maybe_start_turn_for_pending_work().await; + self.input_queue.clear_pending(active_turn).await; } + reason == TurnAbortReason::Interrupted && aborted_turn } + #[expect( + clippy::await_holding_invalid_type, + reason = "targeted abort must remain serialized through terminal delivery" + )] pub(crate) async fn abort_turn_if_active( self: &Arc, turn_id: &str, reason: TurnAbortReason, ) -> bool { + let turn_start = self.turn_start_lock.lock().await; let active_turn = { let mut active = self.active_turn.lock().await; if active @@ -557,6 +590,7 @@ impl Session { // in-flight approval wait can surface as a model-visible rejection before TurnAborted. self.input_queue.clear_pending(&active_turn).await; + drop(turn_start); if reason == TurnAbortReason::Interrupted { self.maybe_start_turn_for_pending_work().await; } @@ -564,11 +598,16 @@ impl Session { true } + #[expect( + clippy::await_holding_invalid_type, + reason = "task completion must remain serialized through terminal delivery" + )] pub async fn on_task_finished( self: &Arc, turn_context: Arc, task_result: SessionTaskResult, ) { + let turn_start = self.turn_start_lock.lock().await; let (last_agent_message, abort_reason) = match task_result { Ok(last_agent_message) => (last_agent_message, None), Err(CodexErr::TurnAborted) => (None, Some(TurnAbortReason::Interrupted)), @@ -797,7 +836,9 @@ impl Session { false } }; + drop(turn_start); if cleared_active_turn { + self.maybe_start_turn_for_pending_work().await; self.emit_thread_idle_lifecycle_if_idle().await; } // Regular items were flushed before this terminal event was appended; buffering