From 051accdb6a8961017604c3a9d03b81ae20e7ec3b Mon Sep 17 00:00:00 2001 From: Joe Gershenson Date: Tue, 21 Apr 2026 00:20:45 -0700 Subject: [PATCH] codex: address PR review feedback (#18817) --- codex-rs/core/src/hook_runtime.rs | 9 +++- codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/src/session/review.rs | 2 +- codex-rs/core/src/session/tests.rs | 56 +++++++++++++++++++++++ codex-rs/core/src/session/turn_context.rs | 4 +- codex-rs/core/src/tasks/mod.rs | 31 ++++++++----- 6 files changed, 88 insertions(+), 15 deletions(-) diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index a0292d9956..2d3e7087eb 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -273,8 +273,15 @@ pub(crate) async fn drain_turn_start_transcript_inputs( sess: &Arc, turn_context: &Arc, ) -> bool { - let _guard = turn_context.transcript_serialization_lock.lock().await; + let Ok(_permit) = turn_context.transcript_serialization_lock.acquire().await else { + return false; + }; if run_pending_session_start_hooks(sess, turn_context).await { + turn_context + .turn_start_transcript_inputs + .lock() + .await + .clear(); return false; } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index f1fea3517a..07fdda6931 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -128,6 +128,7 @@ use rmcp::model::RequestId; use serde_json::Value; use tokio::sync::Mutex; use tokio::sync::RwLock; +use tokio::sync::Semaphore; use tokio::sync::oneshot; use tokio::sync::watch; use tokio::task::JoinHandle; diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index cfd58e0631..b0b48f4ac9 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -140,7 +140,7 @@ pub(super) async fn spawn_review_thread( turn_skills: TurnSkillsContext::new(parent_turn_context.turn_skills.outcome.clone()), turn_timing_state: Arc::new(TurnTimingState::default()), turn_start_transcript_inputs: Arc::new(Mutex::new(Vec::new())), - transcript_serialization_lock: Arc::new(Mutex::new(())), + transcript_serialization_lock: Arc::new(Semaphore::new(1)), }; // Seed the child task with the review prompt as the initial user message. diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index babf659f6d..bbe21c3bc8 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5708,6 +5708,62 @@ async fn abort_regular_task_records_prompt_before_interrupt_marker() { assert_eq!(history.raw_items(), expected.as_slice()); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[test_log::test] +async fn abort_non_regular_task_keeps_pending_session_start_source() { + let (sess, tc, _rx) = make_session_and_context_with_rx().await; + sess.state + .lock() + .await + .set_pending_session_start_source(Some(codex_hooks::SessionStartSource::Startup)); + let input = vec![UserInput::Text { + text: "review this".to_string(), + text_elements: Vec::new(), + }]; + sess.spawn_task( + Arc::clone(&tc), + input, + NeverEndingTask { + kind: TaskKind::Review, + listen_to_cancellation_token: true, + }, + ) + .await; + + sess.abort_all_tasks(TurnAbortReason::Interrupted).await; + + assert!(matches!( + sess.take_pending_session_start_source().await, + Some(codex_hooks::SessionStartSource::Startup) + )); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[test_log::test] +async fn abort_regular_task_without_start_input_keeps_pending_session_start_source() { + let (sess, tc, _rx) = make_session_and_context_with_rx().await; + sess.state + .lock() + .await + .set_pending_session_start_source(Some(codex_hooks::SessionStartSource::Startup)); + sess.spawn_task( + Arc::clone(&tc), + Vec::new(), + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: true, + }, + ) + .await; + + sess.abort_all_tasks(TurnAbortReason::Interrupted).await; + + assert!(matches!( + sess.take_pending_session_start_source().await, + Some(codex_hooks::SessionStartSource::Startup) + )); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn task_finish_emits_turn_item_lifecycle_for_leftover_pending_user_input() { let (sess, tc, rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 27da2ad5dc..c2cdfa8d43 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -72,7 +72,7 @@ pub(crate) struct TurnContext { pub(crate) turn_skills: TurnSkillsContext, pub(crate) turn_timing_state: Arc, pub(crate) turn_start_transcript_inputs: Arc>>>, - pub(crate) transcript_serialization_lock: Arc>, + pub(crate) transcript_serialization_lock: Arc, } impl TurnContext { pub(crate) fn model_context_window(&self) -> Option { @@ -448,7 +448,7 @@ impl Session { turn_skills: TurnSkillsContext::new(skills_outcome), turn_timing_state: Arc::new(TurnTimingState::default()), turn_start_transcript_inputs: Arc::new(Mutex::new(Vec::new())), - transcript_serialization_lock: Arc::new(Mutex::new(())), + transcript_serialization_lock: Arc::new(Semaphore::new(1)), } } diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 50dc3f9744..b7f0c969ca 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -300,12 +300,6 @@ impl Session { } } - let mut active = self.active_turn.lock().await; - let turn = active.get_or_insert_with(ActiveTurn::default); - debug_assert!(turn.tasks.is_empty()); - let done_clone = Arc::clone(&done); - let session_ctx = Arc::new(SessionTaskContext::new(Arc::clone(self))); - let ctx = Arc::clone(&turn_context); if task_kind == TaskKind::Regular && !input.is_empty() { turn_context .turn_start_transcript_inputs @@ -313,6 +307,12 @@ impl Session { .await .push(input.clone()); } + let mut active = self.active_turn.lock().await; + let turn = active.get_or_insert_with(ActiveTurn::default); + debug_assert!(turn.tasks.is_empty()); + let done_clone = Arc::clone(&done); + let session_ctx = Arc::new(SessionTaskContext::new(Arc::clone(self))); + let ctx = Arc::clone(&turn_context); let task_for_run = Arc::clone(&task); let task_cancellation_token = cancellation_token.child_token(); // Task-owned turn spans keep a core-owned span open for the @@ -626,11 +626,6 @@ impl Session { .cancel_git_enrichment_task(); let session_task = task.task; - if reason == TurnAbortReason::Interrupted { - self.cleanup_after_interrupt(&task.turn_context).await; - let _ = drain_turn_start_transcript_inputs(self, &task.turn_context).await; - } - select! { _ = task.done.notified() => { }, @@ -641,12 +636,26 @@ impl Session { task.handle.abort(); + if reason == TurnAbortReason::Interrupted && task.kind == TaskKind::Regular { + let has_turn_start_transcript_input = !task + .turn_context + .turn_start_transcript_inputs + .lock() + .await + .is_empty(); + if has_turn_start_transcript_input { + let _ = drain_turn_start_transcript_inputs(self, &task.turn_context).await; + } + } + let session_ctx = Arc::new(SessionTaskContext::new(Arc::clone(self))); session_task .abort(session_ctx, Arc::clone(&task.turn_context)) .await; if reason == TurnAbortReason::Interrupted { + self.cleanup_after_interrupt(&task.turn_context).await; + let marker = interrupted_turn_history_marker(); self.record_into_history(std::slice::from_ref(&marker), task.turn_context.as_ref()) .await;