diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 4666142887..5ace22c811 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -537,11 +537,9 @@ impl Session { pub async fn abort_all_tasks(self: &Arc, 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, ) { 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, diff --git a/codex-rs/core/tests/suite/abort_lifecycle.rs b/codex-rs/core/tests/suite/abort_lifecycle.rs new file mode 100644 index 0000000000..296094d2c1 --- /dev/null +++ b/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>, +} + +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(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 1be4f9710f..2971b1201d 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -36,6 +36,7 @@ pub static CODEX_ALIASES_TEMP_DIR: Option = { }) }; +mod abort_lifecycle; #[cfg(not(target_os = "windows"))] mod abort_tasks; mod additional_context; diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index bb2b162ce0..09892600d7 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/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;