codex: address PR review feedback (#18817)

codex-rs/core/src/hook_runtime.rs

@@ -273,8 +273,15 @@ pub(crate) async fn drain_turn_start_transcript_inputs(
     sess: &Arc<Session>,
     turn_context: &Arc<TurnContext>,
 ) -> 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;
     }
 

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;

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.

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;

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<TurnTimingState>,
     pub(crate) turn_start_transcript_inputs: Arc<Mutex<Vec<Vec<UserInput>>>>,
-    pub(crate) transcript_serialization_lock: Arc<Mutex<()>>,
+    pub(crate) transcript_serialization_lock: Arc<Semaphore>,
 }
 impl TurnContext {
     pub(crate) fn model_context_window(&self) -> Option<i64> {
@@ -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)),
         }
     }
 

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;