codex: harden turn-start interrupt recovery

codex-rs/core/src/hook_runtime.rs

@@ -31,6 +31,7 @@ use codex_protocol::protocol::HookSource;
 use codex_protocol::protocol::HookStartedEvent;
 use codex_protocol::user_input::UserInput;
 use serde_json::Value;
+use tokio::sync::OwnedSemaphorePermit;
 
 use crate::event_mapping::parse_turn_item;
 use crate::session::PreviousTurnSettings;
@@ -280,10 +281,24 @@ pub(crate) async fn drain_turn_start_transcript_inputs(
     turn_context: &Arc<TurnContext>,
     mode: TurnStartTranscriptDrainMode,
 ) -> bool {
-    let Ok(_permit) = turn_context.transcript_serialization_lock.acquire().await else {
+    let Ok(permit) = turn_context
+        .transcript_serialization_lock
+        .clone()
+        .acquire_owned()
+        .await
+    else {
         return false;
     };
 
+    drain_turn_start_transcript_inputs_with_permit(sess, turn_context, mode, permit).await
+}
+
+pub(crate) async fn drain_turn_start_transcript_inputs_with_permit(
+    sess: &Arc<Session>,
+    turn_context: &Arc<TurnContext>,
+    mode: TurnStartTranscriptDrainMode,
+    _permit: OwnedSemaphorePermit,
+) -> bool {
     let has_queued_start_input = !turn_context.lock_turn_start_transcript_inputs().is_empty();
     if !has_queued_start_input && matches!(mode, TurnStartTranscriptDrainMode::InterruptRecovery) {
         return true;

codex-rs/core/src/session/tests.rs

@@ -5698,49 +5698,6 @@ impl SessionTask for NeverEndingTask {
     }
 }
 
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-#[test_log::test]
-async fn abort_regular_task_records_context_prompt_before_interrupt_marker() {
-    let (sess, tc, _rx) = make_session_and_context_with_rx().await;
-    let input = vec![UserInput::Text {
-        text: "hello".to_string(),
-        text_elements: Vec::new(),
-    }];
-    let mut expected = sess.build_initial_context(tc.as_ref()).await;
-    expected.push(ResponseItem::Message {
-        id: None,
-        role: "user".to_string(),
-        content: vec![ContentItem::InputText {
-            text: "hello".to_string(),
-        }],
-        end_turn: None,
-        phase: None,
-    });
-    expected.push(crate::tasks::interrupted_turn_history_marker());
-    sess.spawn_task(
-        Arc::clone(&tc),
-        input,
-        NeverEndingTask {
-            kind: TaskKind::Regular,
-            listen_to_cancellation_token: true,
-        },
-    )
-    .await;
-
-    sess.abort_all_tasks(TurnAbortReason::Interrupted).await;
-
-    let history = sess.clone_history().await;
-    assert_eq!(history.raw_items(), expected.as_slice());
-    assert!(tc.lock_turn_start_transcript_inputs().is_empty());
-    assert_eq!(
-        sess.previous_turn_settings().await,
-        Some(PreviousTurnSettings {
-            model: tc.model_info.slug.clone(),
-            realtime_active: Some(tc.realtime_active),
-        })
-    );
-}
-
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[test_log::test]
 async fn abort_non_regular_task_keeps_pending_session_start_source() {

codex-rs/core/src/state/turn.rs

@@ -70,6 +70,7 @@ pub(crate) struct RunningTask {
     pub(crate) done: Arc<Notify>,
     pub(crate) kind: TaskKind,
     pub(crate) task: Arc<dyn AnySessionTask>,
+    pub(crate) records_turn_start_transcript: bool,
     pub(crate) cancellation_token: CancellationToken,
     pub(crate) handle: Arc<AbortOnDropHandle<()>>,
     pub(crate) turn_context: Arc<TurnContext>,

codex-rs/core/src/tasks/mod.rs

@@ -24,6 +24,7 @@ use crate::contextual_user_message::TURN_ABORTED_OPEN_TAG;
 use crate::hook_runtime::PendingInputHookDisposition;
 use crate::hook_runtime::TurnStartTranscriptDrainMode;
 use crate::hook_runtime::drain_turn_start_transcript_inputs;
+use crate::hook_runtime::drain_turn_start_transcript_inputs_with_permit;
 use crate::hook_runtime::inspect_pending_input;
 use crate::hook_runtime::record_additional_contexts;
 use crate::hook_runtime::record_pending_input;
@@ -163,6 +164,11 @@ pub(crate) trait SessionTask: Send + Sync + 'static {
     /// Returns the tracing name for a spawned task span.
     fn span_name(&self) -> &'static str;
 
+    /// Whether this task owns the normal user-turn transcript start sequence.
+    fn records_turn_start_transcript(&self) -> bool {
+        false
+    }
+
     /// Executes the task until completion or cancellation.
     ///
     /// Implementations typically stream protocol events using `session` and
@@ -200,6 +206,8 @@ pub(crate) trait AnySessionTask: Send + Sync + 'static {
 
     fn span_name(&self) -> &'static str;
 
+    fn records_turn_start_transcript(&self) -> bool;
+
     fn run(
         self: Arc<Self>,
         session: Arc<SessionTaskContext>,
@@ -227,6 +235,10 @@ where
         SessionTask::span_name(self)
     }
 
+    fn records_turn_start_transcript(&self) -> bool {
+        SessionTask::records_turn_start_transcript(self)
+    }
+
     fn run(
         self: Arc<Self>,
         session: Arc<SessionTaskContext>,
@@ -273,6 +285,7 @@ impl Session {
         let task: Arc<dyn AnySessionTask> = Arc::new(task);
         let task_kind = task.kind();
         let span_name = task.span_name();
+        let records_turn_start_transcript = task.records_turn_start_transcript();
         let started_at = Instant::now();
         turn_context
             .turn_timing_state
@@ -302,7 +315,7 @@ impl Session {
             }
         }
 
-        if task_kind == TaskKind::Regular && !input.is_empty() {
+        if records_turn_start_transcript && !input.is_empty() {
             turn_context
                 .lock_turn_start_transcript_inputs()
                 .push(TurnStartTranscriptInput {
@@ -368,6 +381,7 @@ impl Session {
             handle: Arc::new(AbortOnDropHandle::new(handle)),
             kind: task_kind,
             task,
+            records_turn_start_transcript,
             cancellation_token,
             turn_context: Arc::clone(&turn_context),
             _timer: timer,
@@ -628,35 +642,58 @@ impl Session {
             .cancel_git_enrichment_task();
         let session_task = task.task;
 
-        // The startup prompt queue is serialized by the regular turn itself. If an interrupt
-        // lands while that serialization is in progress, wait for the shared drain to finish
-        // before force-aborting the task so history cannot be left half-written. This can wait
-        // on turn-start hooks, but it keeps hook policy and hook-added context ahead of the
-        // model-visible interrupt marker.
-        if reason == TurnAbortReason::Interrupted && task.kind == TaskKind::Regular {
+        let mut handle_aborted = false;
+        // The startup prompt queue is serialized by the regular turn itself. If the serializer is
+        // idle, grab it and abort the task before recovering the queued prompt; if the serializer
+        // is already active, wait for it to finish before force-aborting so history cannot be left
+        // half-written. This can wait on turn-start hooks, but it keeps hook policy and hook-added
+        // context ahead of the model-visible interrupt marker.
+        if reason == TurnAbortReason::Interrupted && task.records_turn_start_transcript {
             let has_turn_start_transcript_input = {
                 let inputs = task.turn_context.lock_turn_start_transcript_inputs();
                 !inputs.is_empty()
             };
             if has_turn_start_transcript_input {
-                let _ = drain_turn_start_transcript_inputs(
-                    self,
-                    &task.turn_context,
-                    TurnStartTranscriptDrainMode::InterruptRecovery,
-                )
-                .await;
+                match task
+                    .turn_context
+                    .transcript_serialization_lock
+                    .clone()
+                    .try_acquire_owned()
+                {
+                    Ok(permit) => {
+                        task.handle.abort();
+                        handle_aborted = true;
+                        let _ = drain_turn_start_transcript_inputs_with_permit(
+                            self,
+                            &task.turn_context,
+                            TurnStartTranscriptDrainMode::InterruptRecovery,
+                            permit,
+                        )
+                        .await;
+                    }
+                    Err(_) => {
+                        let _ = drain_turn_start_transcript_inputs(
+                            self,
+                            &task.turn_context,
+                            TurnStartTranscriptDrainMode::InterruptRecovery,
+                        )
+                        .await;
+                    }
+                }
             }
         }
 
-        select! {
-            _ = task.done.notified() => {
-            },
-            _ = tokio::time::sleep(Duration::from_millis(GRACEFULL_INTERRUPTION_TIMEOUT_MS)) => {
-                warn!("task {sub_id} didn't complete gracefully after {}ms", GRACEFULL_INTERRUPTION_TIMEOUT_MS);
+        if !handle_aborted {
+            select! {
+                _ = task.done.notified() => {
+                },
+                _ = tokio::time::sleep(Duration::from_millis(GRACEFULL_INTERRUPTION_TIMEOUT_MS)) => {
+                    warn!("task {sub_id} didn't complete gracefully after {}ms", GRACEFULL_INTERRUPTION_TIMEOUT_MS);
+                }
             }
-        }
 
-        task.handle.abort();
+            task.handle.abort();
+        }
 
         let session_ctx = Arc::new(SessionTaskContext::new(Arc::clone(self)));
         session_task

codex-rs/core/src/tasks/regular.rs

@@ -33,6 +33,10 @@ impl SessionTask for RegularTask {
         "session_task.turn"
     }
 
+    fn records_turn_start_transcript(&self) -> bool {
+        true
+    }
+
     async fn run(
         self: Arc<Self>,
         session: Arc<SessionTaskContext>,