Complete turn abort callbacks before emitting terminal events (#49475)

## Why

`TurnAborted` could be emitted before extension abort callbacks finished, allowing consumers to observe a terminal event while lifecycle cleanup was still pending.

## What changed

Await `on_turn_abort` within task abort handling before emitting `TurnAborted`, retaining callback delivery for already-cancelled tasks. Document that `on_turn_stop` and `on_turn_abort` complete before their corresponding terminal events, if emitted.

## Testing

Add regression coverage for interrupt, conditional interrupt, and shutdown. Hold the abort callback open to verify that `TurnAborted` waits for completion, then verify that the callback runs exactly once with the expected turn ID and abort reason.

GitOrigin-RevId: f34d7e94d04c02544178dc445b85abc65a6ce499

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

@@ -537,11 +537,9 @@ impl Session {
     pub async fn abort_all_tasks(self: &Arc<Self>, reason: TurnAbortReason) {
         let mut aborted_turn = false;
         let mut active_turn_to_clear = None;
-        let mut turn_context = None;
         if let Some(mut active_turn) = self.take_active_turn(&reason).await {
             let task = active_turn.task.take();
             aborted_turn = task.is_some();
-            turn_context = task.as_ref().map(|task| Arc::clone(&task.turn_context));
             if let Some(task) = task {
                 self.handle_task_abort(
                     task,
@@ -556,10 +554,6 @@ impl Session {
             }
         }
 
-        if let Some(turn_context) = turn_context.as_deref() {
-            self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref())
-                .await;
-        }
         if let Some(active_turn) = active_turn_to_clear {
             // Let interrupted tasks observe cancellation before dropping pending approvals, or an
             // in-flight approval wait can surface as a model-visible rejection before TurnAborted.
@@ -609,15 +603,10 @@ impl Session {
         error: Option<ErrorEvent>,
     ) {
         let task = active_turn.task.take();
-        let turn_context = task.as_ref().map(|task| Arc::clone(&task.turn_context));
         if let Some(task) = task {
             self.handle_task_abort(task, reason.clone(), &active_turn.turn_state, error)
                 .await;
         }
-        if let Some(turn_context) = turn_context.as_deref() {
-            self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref())
-                .await;
-        }
         // 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;
@@ -928,6 +917,8 @@ impl Session {
     ) {
         let sub_id = task.turn_context.sub_id.clone();
         if task.cancellation_token.is_cancelled() {
+            self.emit_turn_abort_lifecycle(reason, task.turn_context.extension_data.as_ref())
+                .await;
             return;
         }
 
@@ -1005,6 +996,8 @@ impl Session {
                 turn_id: task.turn_context.sub_id.clone(),
                 profile,
             });
+        self.emit_turn_abort_lifecycle(reason.clone(), task.turn_context.extension_data.as_ref())
+            .await;
         let event = EventMsg::TurnAborted(TurnAbortedEvent {
             turn_id: Some(task.turn_context.sub_id.clone()),
             reason,

codex-rs/core/tests/suite/abort_lifecycle.rs

@@ -0,0 +1,150 @@
+//! Verifies abort callbacks finish exactly once before the terminal event is published.
+
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::time::Duration;
+
+use codex_core::StartThreadOptions;
+use codex_core::TurnInputRequest;
+use codex_extension_api::ExtensionFuture;
+use codex_extension_api::ExtensionRegistryBuilder;
+use codex_extension_api::TurnAbortInput;
+use codex_extension_api::TurnLifecycleContributor;
+use codex_protocol::dynamic_tools::DynamicToolFunctionSpec;
+use codex_protocol::dynamic_tools::DynamicToolSpec;
+use codex_protocol::protocol::EventMsg;
+use codex_protocol::protocol::Op;
+use codex_protocol::protocol::TurnAbortReason;
+use codex_protocol::user_input::UserInput;
+use core_test_support::responses;
+use core_test_support::test_codex::test_codex;
+use core_test_support::wait_for_event;
+use pretty_assertions::assert_eq;
+use test_case::test_case;
+use tokio::sync::Notify;
+
+#[derive(Default)]
+struct AbortGate {
+    entered: Notify,
+    release: Notify,
+    completed: Mutex<Vec<(String, TurnAbortReason)>>,
+}
+
+impl TurnLifecycleContributor for AbortGate {
+    fn on_turn_abort<'a>(&'a self, input: TurnAbortInput<'a>) -> ExtensionFuture<'a, ()> {
+        Box::pin(async move {
+            self.entered.notify_one();
+            self.release.notified().await;
+            self.completed
+                .lock()
+                .expect("abort callback records")
+                .push((input.turn_store.level_id().to_string(), input.reason));
+        })
+    }
+}
+
+#[derive(Clone, Copy)]
+enum AbortOperation {
+    Interrupt,
+    ConditionalInterrupt,
+    Shutdown,
+}
+
+#[test_case(AbortOperation::Interrupt; "interrupt")]
+#[test_case(AbortOperation::ConditionalInterrupt; "conditional_interrupt")]
+#[test_case(AbortOperation::Shutdown; "shutdown")]
+#[tokio::test]
+async fn abort_lifecycle_finishes_before_terminal_event(
+    operation: AbortOperation,
+) -> anyhow::Result<()> {
+    let server = responses::start_mock_server().await;
+    let request = responses::mount_sse_once(
+        &server,
+        responses::sse(vec![
+            responses::ev_response_created("response"),
+            responses::ev_function_call("gate-call", "gate", "{}"),
+            responses::ev_completed("response"),
+        ]),
+    )
+    .await;
+    let gate = Arc::new(AbortGate::default());
+    let mut extensions = ExtensionRegistryBuilder::new();
+    extensions.turn_lifecycle_contributor(gate.clone());
+    let mut builder = test_codex().with_extensions(Arc::new(extensions.build()));
+    let test = builder.build_with_auto_env(&server).await?;
+    let thread = test
+        .thread_manager
+        .start_thread(StartThreadOptions {
+            environments: Some(vec![test.executor_environment().selection().clone()]),
+            dynamic_tools: vec![DynamicToolSpec::Function(DynamicToolFunctionSpec {
+                name: "gate".to_string(),
+                description: "Wait for the host.".to_string(),
+                input_schema: serde_json::json!({"type": "object", "properties": {}}),
+                defer_loading: false,
+            })],
+            ..StartThreadOptions::new(test.config.clone())
+        })
+        .await?
+        .thread;
+    thread
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+            text: "Call gate.".to_string(),
+            text_elements: Vec::new(),
+        }]))
+        .await?;
+    let EventMsg::DynamicToolCallRequest(call) = wait_for_event(&thread, |event| {
+        matches!(event, EventMsg::DynamicToolCallRequest(_))
+    })
+    .await
+    else {
+        unreachable!();
+    };
+    request.single_request();
+
+    match operation {
+        AbortOperation::Interrupt => {
+            thread.submit(Op::Interrupt).await?;
+        }
+        AbortOperation::ConditionalInterrupt => {
+            let (reply, received) = tokio::sync::oneshot::channel();
+            thread
+                .submit(Op::InterruptIfNoPendingInput {
+                    turn_id: call.turn_id.clone(),
+                    reply,
+                })
+                .await?;
+            assert!(received.await?);
+        }
+        AbortOperation::Shutdown => {
+            thread.submit(Op::Shutdown).await?;
+        }
+    }
+    tokio::time::timeout(Duration::from_secs(10), gate.entered.notified()).await?;
+    // Keep the callback suspended: scheduling alone must not hide an early event.
+    let early_event = tokio::time::timeout(
+        Duration::from_millis(100),
+        wait_for_event(&thread, |event| matches!(event, EventMsg::TurnAborted(_))),
+    )
+    .await;
+    gate.release.notify_one();
+    assert!(
+        early_event.is_err(),
+        "TurnAborted arrived before its lifecycle callback finished"
+    );
+    let EventMsg::TurnAborted(aborted) =
+        wait_for_event(&thread, |event| matches!(event, EventMsg::TurnAborted(_))).await
+    else {
+        unreachable!();
+    };
+    assert_eq!(aborted.turn_id, Some(call.turn_id.clone()));
+    assert_eq!(aborted.reason, TurnAbortReason::Interrupted);
+    if !matches!(operation, AbortOperation::Shutdown) {
+        thread.submit(Op::Shutdown).await?;
+    }
+    wait_for_event(&thread, |event| matches!(event, EventMsg::ShutdownComplete)).await;
+    assert_eq!(
+        *gate.completed.lock().expect("abort callback records"),
+        vec![(call.turn_id, TurnAbortReason::Interrupted)],
+    );
+    Ok(())
+}

codex-rs/core/tests/suite/mod.rs

@@ -36,6 +36,7 @@ pub static CODEX_ALIASES_TEMP_DIR: Option<TestBinaryDispatchGuard> = {
     })
 };
 
+mod abort_lifecycle;
 #[cfg(not(target_os = "windows"))]
 mod abort_tasks;
 mod additional_context;

codex-rs/ext/extension-api/src/contributors.rs

@@ -247,6 +247,7 @@ pub trait TurnLifecycleContributor: Send + Sync {
     }
 
     /// Called before the host drops the completed turn runtime and turn store.
+    /// The callback completes before the host emits the turn-complete event.
     fn on_turn_stop<'a>(&'a self, input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> {
         Box::pin(async move {
             let _self = self;
@@ -255,6 +256,7 @@ pub trait TurnLifecycleContributor: Send + Sync {
     }
 
     /// Called after the host aborts a running turn.
+    /// The callback completes before the host emits the turn-aborted event, if any.
     fn on_turn_abort<'a>(&'a self, input: TurnAbortInput<'a>) -> ExtensionFuture<'a, ()> {
         Box::pin(async move {
             let _self = self;