From ba409a3fae7cb43e247e08bf68d6697f0c8b6ef7 Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Thu, 9 Jul 2026 20:02:57 -0300 Subject: [PATCH] fix(core): preserve early interrupted turns --- codex-rs/core/src/hook_runtime.rs | 29 +++-- codex-rs/core/src/session/mod.rs | 43 +++++++- codex-rs/core/src/session/tests.rs | 163 +++++++++++++++++++++++------ codex-rs/core/src/session/turn.rs | 90 +++++++++++++++- codex-rs/core/src/state/turn.rs | 1 + codex-rs/core/src/tasks/mod.rs | 43 ++++++-- codex-rs/core/src/tasks/regular.rs | 163 ++++++++++++++++++++++++++--- codex-rs/core/src/tasks/review.rs | 7 +- codex-rs/core/tests/suite/hooks.rs | 119 +++++++++++++++++++++ 9 files changed, 593 insertions(+), 65 deletions(-) diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 342e4d4b8f..96ae97e743 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -53,6 +53,12 @@ pub(crate) struct HookRuntimeOutcome { pub additional_contexts: Vec, } +#[derive(Clone, Copy)] +pub(crate) enum UserInputEventDisposition { + Emit, + AlreadyEmitted, +} + pub(crate) enum PreToolUseHookResult { Continue { updated_input: Option }, Blocked(String), @@ -541,16 +547,23 @@ pub(crate) async fn record_pending_input( turn_context: &Arc, pending_input: TurnInput, additional_contexts: Vec, + user_input_event_disposition: UserInputEventDisposition, ) { match pending_input { - TurnInput::UserInput { content, client_id } => { - sess.record_user_prompt_and_emit_turn_item( - turn_context.as_ref(), - content.as_slice(), - client_id, - ) - .await; - } + TurnInput::UserInput { content, client_id } => match user_input_event_disposition { + UserInputEventDisposition::Emit => { + sess.record_user_prompt_and_emit_turn_item( + turn_context.as_ref(), + content.as_slice(), + client_id, + ) + .await; + } + UserInputEventDisposition::AlreadyEmitted => { + sess.record_user_prompt(turn_context.as_ref(), content.as_slice()) + .await; + } + }, TurnInput::ResponseItem(item) => { sess.record_conversation_items(turn_context, std::slice::from_ref(&item)) .await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9339413b33..05c5ed3767 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3838,12 +3838,53 @@ impl Session { let response_item = self.response_item_from_user_input(input.to_vec()); self.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) .await; + self.emit_user_input_turn_item(turn_context, input, client_id) + .await; + self.ensure_rollout_materialized().await; + } + + pub(crate) async fn record_user_prompt(&self, turn_context: &TurnContext, input: &[UserInput]) { + let response_item = self.response_item_from_user_input(input.to_vec()); + self.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) + .await; + self.ensure_rollout_materialized().await; + } + + /// Persist transcript lifecycle events for a regular turn's initial user input. + /// + /// Accepted prompts use this after submission hooks. The interrupt fallback uses it with the + /// original input when cancellation wins before those hooks reach the persistence boundary. + /// This method intentionally does not add the prompt to model-visible history. + pub(crate) async fn record_initial_user_input_events( + &self, + turn_context: &TurnContext, + input: &[TurnInput], + ) { + let mut emitted = false; + for input in input { + let TurnInput::UserInput { content, client_id } = input else { + continue; + }; + self.emit_user_input_turn_item(turn_context, content, client_id.clone()) + .await; + emitted = true; + } + if emitted { + self.ensure_rollout_materialized().await; + } + } + + async fn emit_user_input_turn_item( + &self, + turn_context: &TurnContext, + input: &[UserInput], + client_id: Option, + ) { let mut user_message_item = UserMessageItem::new(input); user_message_item.client_id = client_id; let turn_item = TurnItem::UserMessage(user_message_item); self.emit_turn_item_started(turn_context, &turn_item).await; self.emit_turn_item_completed(turn_context, turn_item).await; - self.ensure_rollout_materialized().await; } pub(crate) async fn notify_stream_error( diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 696a8adbd8..dfb421f57b 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -16,6 +16,9 @@ use crate::skills::SkillRenderSideEffects; use crate::skills::render::SkillMetadataBudget; use crate::test_support::models_manager_with_provider; use crate::tools::format_exec_output_str; +use codex_app_server_protocol::ThreadItem as AppServerThreadItem; +use codex_app_server_protocol::UserInput as AppServerUserInput; +use codex_app_server_protocol::build_turns_from_rollout_items; use codex_config::ConfigLayerStack; use codex_config::ConfigLayerStackOrdering; use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; @@ -146,7 +149,6 @@ use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::protocol::TokenCountEvent; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; -use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::protocol::UserMessageEvent; @@ -439,13 +441,25 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled() } #[tokio::test] -async fn interrupting_regular_turn_waiting_on_startup_prewarm_emits_turn_aborted() { - let (sess, tc, rx) = make_session_and_context_with_rx().await; +async fn interrupting_regular_turn_waiting_on_startup_prewarm_persists_user_item() { + let (mut sess, tc, rx) = make_session_and_context_with_rx().await; + let rollout_path = attach_thread_persistence( + Arc::get_mut(&mut sess).expect("session should not have additional references"), + ) + .await; let (_tx, startup_prewarm_rx) = tokio::sync::oneshot::channel::<()>(); let handle = tokio::spawn(async move { let _ = startup_prewarm_rx.await; Ok(test_model_client_session()) }); + let prompt = "persist this interrupted prompt"; + let input = vec![TurnInput::UserInput { + content: vec![UserInput::Text { + text: prompt.to_string(), + text_elements: Vec::new(), + }], + client_id: None, + }]; sess.set_session_startup_prewarm( crate::session_startup_prewarm::SessionStartupPrewarmHandle::new( @@ -455,12 +469,8 @@ async fn interrupting_regular_turn_waiting_on_startup_prewarm_emits_turn_aborted ), ) .await; - sess.spawn_task( - Arc::clone(&tc), - Vec::new(), - crate::tasks::RegularTask::new(), - ) - .await; + sess.spawn_task(Arc::clone(&tc), input, crate::tasks::RegularTask::new()) + .await; let first = tokio::time::timeout(std::time::Duration::from_millis(200), rx.recv()) .await @@ -471,31 +481,120 @@ async fn interrupting_regular_turn_waiting_on_startup_prewarm_emits_turn_aborted EventMsg::TurnStarted(TurnStartedEvent { turn_id, .. }) if turn_id == tc.sub_id )); - sess.abort_all_tasks(TurnAbortReason::Interrupted).await; + tokio::time::timeout( + std::time::Duration::from_secs(1), + sess.abort_all_tasks(TurnAbortReason::Interrupted), + ) + .await + .expect("interrupt should not wait for startup hooks"); - let marker_evt = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) - .await - .expect("expected turn aborted marker event") - .expect("channel open"); - assert!(matches!(marker_evt.msg, EventMsg::RawResponseItem(_))); - - let second = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) - .await - .expect("expected turn aborted event") - .expect("channel open"); - let EventMsg::TurnAborted(TurnAbortedEvent { - turn_id, - reason, - completed_at, - duration_ms, - }) = second.msg - else { - panic!("expected turn aborted event"); + let turn_aborted = loop { + let event = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + .await + .expect("expected turn aborted event") + .expect("channel open"); + if let EventMsg::TurnAborted(event) = event.msg { + break event; + } }; - assert_eq!(turn_id, Some(tc.sub_id.clone())); - assert_eq!(reason, TurnAbortReason::Interrupted); - assert!(completed_at.is_some()); - assert!(duration_ms.is_some()); + assert_eq!(turn_aborted.turn_id, Some(tc.sub_id.clone())); + assert_eq!(turn_aborted.reason, TurnAbortReason::Interrupted); + assert!(turn_aborted.completed_at.is_some()); + assert!(turn_aborted.duration_ms.is_some()); + + let is_turn_aborted_marker = |item: &ResponseItem| { + let ResponseItem::Message { content, .. } = item else { + return false; + }; + content.iter().any(|content_item| { + let ContentItem::InputText { text } = content_item else { + return false; + }; + TurnAborted::matches_text(text) + }) + }; + let live_history = strip_metadata_from_items(sess.clone_history().await.raw_items()); + assert!(!live_history.contains(&user_message(prompt))); + assert_eq!( + live_history + .iter() + .filter(|item| is_turn_aborted_marker(item)) + .count(), + 1 + ); + + let InitialHistory::Resumed(resumed) = RolloutRecorder::get_rollout_history(&rollout_path) + .await + .expect("read rollout history") + else { + panic!("expected resumed rollout history"); + }; + let user_message_positions = resumed + .history + .iter() + .enumerate() + .filter_map(|(index, item)| { + matches!( + item, + RolloutItem::EventMsg(EventMsg::UserMessage(event)) if event.message == prompt + ) + .then_some(index) + }) + .collect::>(); + let turn_aborted_positions = resumed + .history + .iter() + .enumerate() + .filter_map(|(index, item)| { + matches!( + item, + RolloutItem::EventMsg(EventMsg::TurnAborted(event)) + if event.turn_id.as_deref() == Some(tc.sub_id.as_str()) + && event.reason == TurnAbortReason::Interrupted + ) + .then_some(index) + }) + .collect::>(); + assert_eq!(user_message_positions.len(), 1); + assert_eq!(turn_aborted_positions.len(), 1); + assert!(user_message_positions[0] < turn_aborted_positions[0]); + assert!(!resumed.history.iter().any( + |item| matches!(item, RolloutItem::ResponseItem(item) if item == &user_message(prompt)) + )); + + let turns = build_turns_from_rollout_items(&resumed.history); + let interrupted_turn = turns + .iter() + .find(|turn| turn.id == tc.sub_id) + .expect("interrupted turn should be reconstructed"); + let transcript_prompts = interrupted_turn + .items + .iter() + .filter_map(|item| match item { + AppServerThreadItem::UserMessage { content, .. } => Some(content), + _ => None, + }) + .collect::>(); + assert_eq!(transcript_prompts.len(), 1); + assert!(matches!( + transcript_prompts[0].as_slice(), + [AppServerUserInput::Text { text, .. }] if text == prompt + )); + + let (resumed_session, _resumed_turn_context) = make_session_and_context().await; + resumed_session + .record_initial_history(InitialHistory::Resumed(resumed)) + .await; + let resumed_history = + strip_metadata_from_items(resumed_session.clone_history().await.raw_items()); + assert!(!resumed_history.contains(&user_message(prompt))); + assert_eq!( + resumed_history + .iter() + .filter(|item| is_turn_aborted_marker(item)) + .count(), + 1 + ); } fn test_model_client_session() -> crate::client::ModelClientSession { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index bafb368b32..d377dd86ba 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -18,6 +18,7 @@ use crate::compact_remote_v2::run_inline_remote_auto_compact_task as run_inline_ use crate::connectors; use crate::context::ContextualUserFragment; use crate::feedback_tags; +use crate::hook_runtime::UserInputEventDisposition; use crate::hook_runtime::inspect_pending_input; use crate::hook_runtime::record_additional_contexts; use crate::hook_runtime::record_pending_input; @@ -53,6 +54,7 @@ use crate::stream_events_utils::last_assistant_message_from_item; use crate::stream_events_utils::mark_thread_memory_mode_polluted_if_external_context; use crate::stream_events_utils::raw_assistant_output_text_from_item; use crate::stream_events_utils::record_completed_response_item_with_finalized_facts; +use crate::tasks::InitialUserInputEvents; use crate::tasks::emit_compact_metric; use crate::tools::ToolRouter; use crate::tools::context::SharedTurnDiffTracker; @@ -146,6 +148,7 @@ pub(crate) async fn run_turn( input: Vec, prewarmed_client_session: Option, cancellation_token: CancellationToken, + initial_user_input_events: Option<&InitialUserInputEvents>, ) -> CodexResult> { let mut client_session = prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); @@ -184,10 +187,25 @@ pub(crate) async fn run_turn( }; if run_pending_session_start_hooks(&sess, &turn_context).await { + if let Some(initial_user_input_events) = initial_user_input_events { + initial_user_input_events.set_abort_input(Vec::new()); + } return Ok(None); } let mut can_drain_pending_input = input.is_empty(); - if run_hooks_and_record_inputs(&sess, &turn_context, &input).await { + let initial_input_stopped = if let Some(initial_user_input_events) = initial_user_input_events { + run_hooks_and_record_initial_inputs(&sess, &turn_context, &input, initial_user_input_events) + .await + } else { + run_hooks_and_record_inputs( + &sess, + &turn_context, + &input, + UserInputEventDisposition::Emit, + ) + .await + }; + if initial_input_stopped { return Ok(None); } @@ -232,7 +250,14 @@ pub(crate) async fn run_turn( Vec::new() }; - if run_hooks_and_record_inputs(&sess, &turn_context, &pending_input).await { + if run_hooks_and_record_inputs( + &sess, + &turn_context, + &pending_input, + UserInputEventDisposition::Emit, + ) + .await + { break; } @@ -483,6 +508,7 @@ async fn run_hooks_and_record_inputs( sess: &Arc, turn_context: &Arc, input: &[TurnInput], + user_input_event_disposition: UserInputEventDisposition, ) -> bool { let mut blocked_input = false; let mut accepted_user_input = false; @@ -500,6 +526,7 @@ async fn run_hooks_and_record_inputs( turn_context, input_item.clone(), hook_outcome.additional_contexts, + user_input_event_disposition, ) .await; } @@ -507,6 +534,65 @@ async fn run_hooks_and_record_inputs( blocked_input && !accepted_user_input } +#[instrument(level = "trace", skip_all)] +async fn run_hooks_and_record_initial_inputs( + sess: &Arc, + turn_context: &Arc, + input: &[TurnInput], + initial_user_input_events: &InitialUserInputEvents, +) -> bool { + let mut blocked_input = false; + let mut accepted_user_input = false; + let mut accepted: Vec<(TurnInput, Vec)> = Vec::new(); + for (index, input_item) in input.iter().enumerate() { + let hook_outcome = inspect_pending_input(sess, turn_context, input_item).await; + if hook_outcome.should_stop { + blocked_input = true; + let abort_input = accepted + .iter() + .map(|(input, _additional_contexts)| input.clone()) + .chain(input.iter().skip(index + 1).cloned()) + .collect(); + initial_user_input_events.set_abort_input(abort_input); + record_additional_contexts(sess, turn_context, hook_outcome.additional_contexts).await; + } else { + if matches!(input_item, TurnInput::UserInput { content, .. } if !content.is_empty()) { + accepted_user_input = true; + } + accepted.push((input_item.clone(), hook_outcome.additional_contexts)); + let abort_input = accepted + .iter() + .map(|(input, _additional_contexts)| input.clone()) + .chain(input.iter().skip(index + 1).cloned()) + .collect(); + initial_user_input_events.set_abort_input(abort_input); + } + } + + let accepted_input = accepted + .iter() + .map(|(input, _additional_contexts)| input.clone()) + .collect(); + if !initial_user_input_events + .record(Arc::clone(sess), Arc::clone(turn_context), accepted_input) + .await + { + return true; + } + + for (input, additional_contexts) in accepted { + record_pending_input( + sess, + turn_context, + input, + additional_contexts, + UserInputEventDisposition::AlreadyEmitted, + ) + .await; + } + blocked_input && !accepted_user_input +} + #[instrument(level = "trace", skip_all)] async fn build_skills_and_plugins( sess: &Arc, diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index d1e31b7994..12cdbcb13c 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -73,6 +73,7 @@ pub(crate) struct RunningTask { pub(crate) done: Arc, pub(crate) kind: TaskKind, pub(crate) task: Arc, + pub(crate) abort_input: Vec, pub(crate) cancellation_token: CancellationToken, pub(crate) handle: AbortOnDropHandle<()>, pub(crate) turn_context: Arc, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 7a930bb578..05640eb17e 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -25,6 +25,7 @@ use tracing::warn; use crate::codex_thread::BackgroundTerminalInfo; use crate::config::Config; use crate::context::ContextualUserFragment; +use crate::hook_runtime::UserInputEventDisposition; use crate::hook_runtime::inspect_pending_input; use crate::hook_runtime::record_additional_contexts; use crate::hook_runtime::record_pending_input; @@ -58,6 +59,7 @@ use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use codex_protocol::models::ContentItem; pub(crate) use compact::CompactTask; +pub(crate) use regular::InitialUserInputEvents; pub(crate) use regular::RegularTask; pub(crate) use review::ReviewTask; pub(crate) use user_shell::UserShellCommandMode; @@ -241,14 +243,17 @@ pub(crate) trait SessionTask: Send + Sync + 'static { /// /// The default implementation is a no-op; override this if additional /// teardown or notifications are required once - /// [`Session::abort_all_tasks`] cancels the task. - fn abort( - &self, + /// [`Session::abort_all_tasks`] cancels the task. `input` is the same initial input passed to + /// [`SessionTask::run`], retained for cleanup that must handle cancellation before the run + /// future is first polled. + fn abort<'a>( + &'a self, session: Arc, ctx: Arc, - ) -> impl std::future::Future + Send { + input: &'a [TurnInput], + ) -> impl std::future::Future + Send + 'a { async move { - let _ = (session, ctx); + let _ = (session, ctx, input); } } } @@ -270,6 +275,7 @@ pub(crate) trait AnySessionTask: Send + Sync + 'static { &'a self, session: Arc, ctx: Arc, + input: &'a [TurnInput], ) -> BoxFuture<'a, ()>; } @@ -305,8 +311,9 @@ where &'a self, session: Arc, ctx: Arc, + input: &'a [TurnInput], ) -> BoxFuture<'a, ()> { - Box::pin(SessionTask::abort(self, session, ctx)) + Box::pin(SessionTask::abort(self, session, ctx, input)) } } @@ -379,6 +386,7 @@ impl Session { )); let ctx = Arc::clone(&turn_context); let task_for_run = Arc::clone(&task); + let abort_input = input.clone(); let task_input = input; let task_cancellation_token = cancellation_token.child_token(); // Task-owned turn spans keep a core-owned span open for the @@ -441,6 +449,7 @@ impl Session { handle: AbortOnDropHandle::new(handle), kind: task_kind, task, + abort_input, cancellation_token, turn_context: Arc::clone(&turn_context), turn_extension_data, @@ -617,6 +626,7 @@ impl Session { &turn_context, pending_input_item, hook_outcome.additional_contexts, + UserInputEventDisposition::Emit, ) .await; } @@ -847,14 +857,31 @@ impl Session { } } - task.handle.abort(); + let mut handle = task.handle; + handle.abort(); + // Wait for cancellation to finish before task-specific abort cleanup reads state that the + // run future may update between cancellable awaits. Keep this bounded because Tokio cannot + // force a task to stop while it is executing non-yielding code. + if tokio::time::timeout( + Duration::from_millis(GRACEFULL_INTERRUPTION_TIMEOUT_MS), + &mut handle, + ) + .await + .is_err() + { + warn!("task {sub_id} did not stop after forced abort"); + } let session_ctx = Arc::new(SessionTaskContext::new( Arc::clone(self), Arc::clone(&task.turn_extension_data), )); session_task - .abort(session_ctx, Arc::clone(&task.turn_context)) + .abort( + session_ctx, + Arc::clone(&task.turn_context), + &task.abort_input, + ) .await; if reason == TurnAbortReason::Interrupted diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 40837728a2..4723214a1e 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -1,5 +1,9 @@ use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use tokio::sync::watch; use tokio_util::sync::CancellationToken; use crate::session::TurnInput; @@ -16,12 +20,136 @@ use super::SessionTask; use super::SessionTaskContext; use super::SessionTaskResult; +pub(crate) struct InitialUserInputEvents { + abort_input: Mutex>>, + turn_start_started: AtomicBool, + turn_start_recorded: watch::Sender, + started: AtomicBool, + recorded: watch::Sender, +} + +impl Default for InitialUserInputEvents { + fn default() -> Self { + let (recorded, _receiver) = watch::channel(false); + let (turn_start_recorded, _receiver) = watch::channel(false); + Self { + abort_input: Mutex::new(None), + turn_start_started: AtomicBool::new(false), + turn_start_recorded, + started: AtomicBool::new(false), + recorded, + } + } +} + +impl InitialUserInputEvents { + /// Emit the regular-turn boundary exactly once, even when cancellation transfers ownership + /// from the task future to abort cleanup. + async fn record_turn_start( + &self, + session: Arc, + ctx: Arc, + ) { + let mut recorded = self.turn_start_recorded.subscribe(); + let claimed = self + .turn_start_started + .compare_exchange( + /*current*/ false, + /*new*/ true, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok(); + if claimed { + let recorded = self.turn_start_recorded.clone(); + tokio::spawn(async move { + session + .send_event( + ctx.as_ref(), + EventMsg::TurnStarted(TurnStartedEvent { + turn_id: ctx.sub_id.clone(), + trace_id: ctx.trace_id.clone(), + started_at: ctx.turn_timing_state.started_at_unix_secs().await, + model_context_window: ctx.model_context_window(), + collaboration_mode_kind: ctx.collaboration_mode.mode, + }), + ) + .await; + session + .set_server_reasoning_included(/*included*/ false) + .await; + recorded.send_replace(true); + }); + } + + while !*recorded.borrow_and_update() { + // `self.turn_start_recorded` keeps the channel open for this task's lifetime. + let _ = recorded.changed().await; + } + } + + /// Publish the hook-filtered input that an interrupt should preserve if it wins before normal + /// persistence. This is synchronous so a completed hook decision is visible before the next + /// cancellable await. + pub(crate) fn set_abort_input(&self, input: Vec) { + *self + .abort_input + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(input); + } + + fn abort_input(&self, original_input: &[TurnInput]) -> Vec { + self.abort_input + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + .unwrap_or_else(|| original_input.to_vec()) + } + + /// Record one durable copy of the initial user items and report whether this caller won the + /// right to choose which items were recorded. + pub(crate) async fn record( + &self, + session: Arc, + ctx: Arc, + input: Vec, + ) -> bool { + let mut recorded = self.recorded.subscribe(); + let claimed = self + .started + .compare_exchange( + /*current*/ false, + /*new*/ true, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok(); + if claimed { + let recorded = self.recorded.clone(); + tokio::spawn(async move { + session + .record_initial_user_input_events(ctx.as_ref(), &input) + .await; + recorded.send_replace(true); + }); + } + + while !*recorded.borrow_and_update() { + // `self.recorded` keeps the channel open for the duration of this borrow. + let _ = recorded.changed().await; + } + claimed + } +} + #[derive(Default)] -pub(crate) struct RegularTask; +pub(crate) struct RegularTask { + initial_user_input_events: InitialUserInputEvents, +} impl RegularTask { pub(crate) fn new() -> Self { - Self + Self::default() } } @@ -44,18 +172,10 @@ impl SessionTask for RegularTask { let sess = session.clone_session(); let turn_extension_data = session.turn_extension_data(); let run_turn_span = trace_span!("run_turn"); - // Regular turns emit `TurnStarted` inline so first-turn lifecycle does - // not wait on startup prewarm resolution. + self.initial_user_input_events + .record_turn_start(Arc::clone(&sess), Arc::clone(&ctx)) + .await; let prewarmed_client_session = async { - let event = EventMsg::TurnStarted(TurnStartedEvent { - turn_id: ctx.sub_id.clone(), - trace_id: ctx.trace_id.clone(), - started_at: ctx.turn_timing_state.started_at_unix_secs().await, - model_context_window: ctx.model_context_window(), - collaboration_mode_kind: ctx.collaboration_mode.mode, - }); - sess.send_event(ctx.as_ref(), event).await; - sess.set_server_reasoning_included(/*included*/ false).await; sess.consume_startup_prewarm_for_regular_turn(&cancellation_token) .await } @@ -70,6 +190,7 @@ impl SessionTask for RegularTask { }; let mut next_input = input; let mut prewarmed_client_session = prewarmed_client_session; + let mut initial_user_input_events = Some(&self.initial_user_input_events); loop { let last_agent_message = run_turn( Arc::clone(&sess), @@ -78,6 +199,7 @@ impl SessionTask for RegularTask { next_input, prewarmed_client_session.take(), cancellation_token.child_token(), + initial_user_input_events.take(), ) .instrument(run_turn_span.clone()) .await?; @@ -87,4 +209,19 @@ impl SessionTask for RegularTask { next_input = Vec::new(); } } + + async fn abort( + &self, + session: Arc, + ctx: Arc, + input: &[TurnInput], + ) { + self.initial_user_input_events + .record_turn_start(session.clone_session(), Arc::clone(&ctx)) + .await; + let input = self.initial_user_input_events.abort_input(input); + self.initial_user_input_events + .record(session.clone_session(), ctx, input) + .await; + } } diff --git a/codex-rs/core/src/tasks/review.rs b/codex-rs/core/src/tasks/review.rs index 30f8d47d81..085393037b 100644 --- a/codex-rs/core/src/tasks/review.rs +++ b/codex-rs/core/src/tasks/review.rs @@ -88,7 +88,12 @@ impl SessionTask for ReviewTask { Ok(None) } - async fn abort(&self, session: Arc, ctx: Arc) { + async fn abort( + &self, + session: Arc, + ctx: Arc, + _input: &[TurnInput], + ) { exit_review_mode(session.clone_session(), /*review_output*/ None, ctx).await; } } diff --git a/codex-rs/core/tests/suite/hooks.rs b/codex-rs/core/tests/suite/hooks.rs index a5c82d5767..ff540ffc01 100644 --- a/codex-rs/core/tests/suite/hooks.rs +++ b/codex-rs/core/tests/suite/hooks.rs @@ -22,6 +22,7 @@ use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; use codex_protocol::user_input::UserInput; use codex_utils_absolute_path::AbsolutePathBuf; +use core_test_support::fs_wait; use core_test_support::hooks::trust_discovered_hooks; use core_test_support::hooks::trust_hooks; use core_test_support::managed_network_requirements_loader; @@ -259,6 +260,43 @@ if payload.get("prompt") == {blocked_prompt_json}: Ok(()) } +fn write_blocking_user_prompt_submit_hook(home: &Path) -> Result<()> { + let script_path = home.join("blocking_user_prompt_submit_hook.py"); + let started_path = home.join("blocking_user_prompt_submit_hook.started"); + let release_path = home.join("blocking_user_prompt_submit_hook.release"); + let script = format!( + r#"from pathlib import Path +import json +import sys +import time + +json.load(sys.stdin) +Path(r"{started_path}").touch() +release_path = Path(r"{release_path}") +while not release_path.exists(): + time.sleep(0.01) +print(json.dumps({{}})) +"#, + started_path = started_path.display(), + release_path = release_path.display(), + ); + let hooks = serde_json::json!({ + "hooks": { + "UserPromptSubmit": [{ + "hooks": [{ + "type": "command", + "command": format!("python3 {}", script_path.display()), + "statusMessage": "waiting in user prompt submit hook", + }] + }] + } + }); + + fs::write(script_path, script).context("write blocking user prompt submit hook")?; + fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; + Ok(()) +} + fn write_session_start_and_user_prompt_submit_order_hooks(home: &Path) -> Result<()> { let session_start_script_path = home.join("session_start_order_hook.py"); let user_prompt_submit_script_path = home.join("user_prompt_submit_order_hook.py"); @@ -921,6 +959,24 @@ fn rollout_hook_prompt_texts(text: &str) -> Result> { Ok(texts) } +fn rollout_user_message_event_texts(text: &str) -> Result> { + text.lines() + .filter(|line| !line.trim().is_empty()) + .filter_map(|line| { + let rollout: Result = + serde_json::from_str(line).context("parse rollout line"); + match rollout { + Ok(RolloutLine { + item: RolloutItem::EventMsg(EventMsg::UserMessage(event)), + .. + }) => Some(Ok(event.message)), + Ok(_) => None, + Err(err) => Some(Err(err)), + } + }) + .collect() +} + fn request_hook_prompt_texts( request: &core_test_support::responses::ResponsesRequest, ) -> Vec { @@ -1753,6 +1809,14 @@ async fn blocked_user_prompt_submit_persists_additional_context_for_next_turn() "second request should include the accepted prompt", ); + let rollout_path = test.codex.rollout_path().expect("rollout path"); + let rollout_text = fs::read_to_string(rollout_path)?; + assert_eq!( + rollout_user_message_event_texts(&rollout_text)?, + vec!["second prompt".to_string()], + "blocked prompts should not become transcript user items", + ); + let hook_inputs = read_user_prompt_submit_hook_inputs(test.codex_home_path())?; assert_eq!(hook_inputs.len(), 2); assert_eq!( @@ -1780,6 +1844,61 @@ async fn blocked_user_prompt_submit_persists_additional_context_for_next_turn() Ok(()) } +#[tokio::test] +async fn interrupt_during_user_prompt_submit_hook_persists_one_transcript_item() -> Result<()> { + skip_if_host_windows!(Ok(())); + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let mut builder = test_codex() + .with_pre_build_hook(|home| { + write_blocking_user_prompt_submit_hook(home) + .expect("failed to write blocking user prompt submit hook fixture"); + }) + .with_config(trust_discovered_hooks); + let test = builder.build(&server).await?; + let prompt = "persist the prompt interrupted during its hook"; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: prompt.to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + fs_wait::wait_for_path_exists( + test.codex_home_path() + .join("blocking_user_prompt_submit_hook.started"), + Duration::from_secs(5), + ) + .await?; + + test.codex.submit(Op::Interrupt).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnAborted(_)) + }) + .await; + fs::write( + test.codex_home_path() + .join("blocking_user_prompt_submit_hook.release"), + b"release", + )?; + + let rollout_path = test.codex.rollout_path().expect("rollout path"); + let rollout_text = fs::read_to_string(rollout_path)?; + assert_eq!( + rollout_user_message_event_texts(&rollout_text)?, + vec![prompt.to_string()], + ); + + Ok(()) +} + #[tokio::test] async fn blocked_queued_prompt_does_not_strand_earlier_accepted_prompt() -> Result<()> { skip_if_no_network!(Ok(()));