diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 3e5403eb04..03c8fbcbf5 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/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, 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, + turn_context: &Arc, + 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; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6c87d1bc91..cea1b3d4b6 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/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() { diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index f1b7bec6bb..94f599eb52 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -70,6 +70,7 @@ pub(crate) struct RunningTask { pub(crate) done: Arc, pub(crate) kind: TaskKind, pub(crate) task: Arc, + pub(crate) records_turn_start_transcript: bool, pub(crate) cancellation_token: CancellationToken, pub(crate) handle: Arc>, pub(crate) turn_context: Arc, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 784972a98e..87e895faf6 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/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, session: Arc, @@ -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, session: Arc, @@ -273,6 +285,7 @@ impl Session { let task: Arc = 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 diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index d64b83f01b..62f6f81624 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/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, session: Arc,