From c0e4efdfb54cfb563ccb88bbb5ee88c78b1604a9 Mon Sep 17 00:00:00 2001 From: Joe Gershenson Date: Tue, 21 Apr 2026 00:58:06 -0700 Subject: [PATCH] codex: serialize interrupt turn start recovery --- codex-rs/core/src/hook_runtime.rs | 94 +++++++------- codex-rs/core/src/session/tests.rs | 145 ++++++++++------------ codex-rs/core/src/session/turn.rs | 23 ++-- codex-rs/core/src/session/turn_context.rs | 8 -- codex-rs/core/src/tasks/mod.rs | 31 +++-- 5 files changed, 133 insertions(+), 168 deletions(-) diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 8d73b38275..3e5403eb04 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -33,9 +33,9 @@ use codex_protocol::user_input::UserInput; use serde_json::Value; use crate::event_mapping::parse_turn_item; +use crate::session::PreviousTurnSettings; use crate::session::session::Session; use crate::session::turn_context::TurnContext; -use crate::session::turn_context::TurnStartUserPromptSubmitOutcome; use crate::tools::sandboxing::PermissionRequestPayload; pub(crate) struct HookRuntimeOutcome { @@ -59,6 +59,11 @@ pub(crate) enum PendingInputRecord { }, } +pub(crate) enum TurnStartTranscriptDrainMode { + RegularTurn, + InterruptRecovery, +} + struct ContextInjectingHookOutcome { hook_events: Vec, outcome: HookRuntimeOutcome, @@ -273,52 +278,45 @@ pub(crate) async fn inspect_pending_input( pub(crate) async fn drain_turn_start_transcript_inputs( sess: &Arc, turn_context: &Arc, + mode: TurnStartTranscriptDrainMode, ) -> bool { let Ok(_permit) = turn_context.transcript_serialization_lock.acquire().await else { return false; }; + + 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; + } + + // Keep the normal turn-start ordering in one serialized region: context first, + // then the user prompt, then hook-provided context and previous-turn settings. + sess.record_context_updates_and_set_reference_context_item(turn_context.as_ref()) + .await; + if run_pending_session_start_hooks(sess, turn_context).await { turn_context.lock_turn_start_transcript_inputs().clear(); return false; } + let mut recorded_start_input = false; loop { - let queued_input = { + let input = { let inputs = turn_context.lock_turn_start_transcript_inputs(); - inputs.first().cloned() + inputs.first().map(|queued| queued.input.clone()) }; - let Some(queued_input) = queued_input else { + let Some(input) = input else { break; }; - let input = queued_input.input; let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input.clone()); let response_item: ResponseItem = initial_input_for_turn.into(); - let user_prompt_submit_outcome = match queued_input.user_prompt_submit_outcome { - Some(outcome) => outcome, - None => { - let outcome = run_user_prompt_submit_hooks( - sess, - turn_context, - UserMessageItem::new(&input).message(), - ) - .await; - let outcome = TurnStartUserPromptSubmitOutcome { - should_stop: outcome.should_stop, - additional_contexts: outcome.additional_contexts, - }; - { - let mut inputs = turn_context.lock_turn_start_transcript_inputs(); - if let Some(queued) = inputs.first_mut() - && queued.input == input - && queued.user_prompt_submit_outcome.is_none() - { - queued.user_prompt_submit_outcome = Some(outcome.clone()); - } - } - outcome - } - }; + let user_prompt_submit_outcome = run_user_prompt_submit_hooks( + sess, + turn_context, + UserMessageItem::new(&input).message(), + ) + .await; if user_prompt_submit_outcome.should_stop { record_additional_contexts( @@ -334,27 +332,12 @@ pub(crate) async fn drain_turn_start_transcript_inputs( return false; } - if !queued_input.user_prompt_recorded { - sess.record_conversation_items( - turn_context.as_ref(), - std::slice::from_ref(&response_item), - ) - .await; - { - let mut inputs = turn_context.lock_turn_start_transcript_inputs(); - if let Some(queued) = inputs.first_mut() - && queued.input == input - { - queued.user_prompt_recorded = true; - } - } - let turn_item = TurnItem::UserMessage(UserMessageItem::new(&input)); - sess.emit_turn_item_started(turn_context.as_ref(), &turn_item) - .await; - sess.emit_turn_item_completed(turn_context.as_ref(), turn_item) - .await; - sess.ensure_rollout_materialized().await; - } + sess.record_user_prompt_and_emit_turn_item( + turn_context.as_ref(), + input.as_slice(), + response_item, + ) + .await; record_additional_contexts( sess, turn_context, @@ -365,6 +348,15 @@ pub(crate) async fn drain_turn_start_transcript_inputs( if inputs.first().is_some_and(|queued| queued.input == input) { inputs.remove(0); } + recorded_start_input = true; + } + + if recorded_start_input { + sess.set_previous_turn_settings(Some(PreviousTurnSettings { + model: turn_context.model_info.slug.clone(), + realtime_active: Some(turn_context.realtime_active), + })) + .await; } true diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 049a7a024d..6c87d1bc91 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -55,8 +55,6 @@ use tracing::Span; use crate::RolloutRecorderParams; use crate::rollout::policy::EventPersistenceMode; use crate::rollout::recorder::RolloutRecorder; -use crate::session::turn_context::TurnStartTranscriptInput; -use crate::session::turn_context::TurnStartUserPromptSubmitOutcome; use crate::state::TaskKind; use crate::tasks::SessionTask; use crate::tasks::SessionTaskContext; @@ -307,12 +305,23 @@ 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; + let input = vec![UserInput::Text { + text: "hello before prewarm".to_string(), + text_elements: Vec::new(), + }]; + let mut expected_history = sess.build_initial_context(tc.as_ref()).await; + expected_history.push(ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "hello before prewarm".to_string(), + }], + end_turn: None, + phase: None, + }); + expected_history.push(crate::tasks::interrupted_turn_history_marker()); + 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 @@ -325,23 +334,38 @@ async fn interrupting_regular_turn_waiting_on_startup_prewarm_emits_turn_aborted sess.abort_all_tasks(TurnAbortReason::Interrupted).await; - let second = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + let (turn_id, reason, completed_at, duration_ms) = + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + let event = rx.recv().await.expect("channel open"); + if let EventMsg::TurnAborted(TurnAbortedEvent { + turn_id, + reason, + completed_at, + duration_ms, + }) = event.msg + { + return (turn_id, reason, completed_at, duration_ms); + } + } + }) .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"); - }; + .expect("expected turn aborted 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()); + + let history = sess.clone_history().await; + assert_eq!(history.raw_items(), expected_history.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), + }) + ); } fn test_model_client_session() -> crate::client::ModelClientSession { @@ -5676,12 +5700,23 @@ impl SessionTask for NeverEndingTask { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[test_log::test] -async fn abort_regular_task_records_prompt_before_interrupt_marker() { +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, @@ -5695,69 +5730,15 @@ async fn abort_regular_task_records_prompt_before_interrupt_marker() { sess.abort_all_tasks(TurnAbortReason::Interrupted).await; let history = sess.clone_history().await; - let expected = vec![ - ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { - text: "hello".to_string(), - }], - end_turn: None, - phase: None, - }, - crate::tasks::interrupted_turn_history_marker(), - ]; - assert_eq!(history.raw_items(), expected.as_slice()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[test_log::test] -async fn abort_regular_task_replays_context_without_replaying_prompt() { - 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 response_item = ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { - text: "hello".to_string(), - }], - end_turn: None, - phase: None, - }; - sess.record_conversation_items(tc.as_ref(), std::slice::from_ref(&response_item)) - .await; - tc.lock_turn_start_transcript_inputs() - .push(TurnStartTranscriptInput { - input, - user_prompt_submit_outcome: Some(TurnStartUserPromptSubmitOutcome { - should_stop: false, - additional_contexts: vec!["hook context".to_string()], - }), - user_prompt_recorded: true, - }); - sess.spawn_task( - Arc::clone(&tc), - Vec::new(), - NeverEndingTask { - kind: TaskKind::Regular, - listen_to_cancellation_token: true, - }, - ) - .await; - - sess.abort_all_tasks(TurnAbortReason::Interrupted).await; - - let history = sess.clone_history().await; - let expected = vec![ - response_item, - DeveloperInstructions::new("hook context".to_string()).into(), - crate::tasks::interrupted_turn_history_marker(), - ]; 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)] diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index a05b0b16d1..5a98e0a406 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::run_inline_remote_auto_compact_task; use crate::connectors; use crate::feedback_tags; use crate::hook_runtime::PendingInputHookDisposition; +use crate::hook_runtime::TurnStartTranscriptDrainMode; use crate::hook_runtime::drain_turn_start_transcript_inputs; use crate::hook_runtime::emit_hook_completed_events; use crate::hook_runtime::inspect_pending_input; @@ -37,7 +38,6 @@ use crate::mentions::collect_tool_mentions_from_messages; use crate::parse_turn_item; use crate::plugins::build_plugin_injections; use crate::resolve_skill_dependencies_for_turn; -use crate::session::PreviousTurnSettings; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::stream_events_utils::HandleOutputCtx; @@ -157,9 +157,6 @@ pub(crate) async fn run_turn( let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref()); - sess.record_context_updates_and_set_reference_context_item(turn_context.as_ref()) - .await; - let loaded_plugins = sess .services .plugins_manager @@ -284,7 +281,13 @@ pub(crate) async fn run_turn( }) .collect::>(); - if !drain_turn_start_transcript_inputs(&sess, &turn_context).await { + if !drain_turn_start_transcript_inputs( + &sess, + &turn_context, + TurnStartTranscriptDrainMode::RegularTurn, + ) + .await + { return None; } sess.services @@ -297,16 +300,6 @@ pub(crate) async fn run_turn( } sess.merge_connector_selection(explicitly_enabled_connectors.clone()) .await; - if !input.is_empty() { - // Track the previous-turn baseline from the regular user-turn path only so - // standalone tasks (compact/shell/review/undo) cannot suppress future - // model/realtime injections. - sess.set_previous_turn_settings(Some(PreviousTurnSettings { - model: turn_context.model_info.slug.clone(), - realtime_active: Some(turn_context.realtime_active), - })) - .await; - } let agent_task = match sess.ensure_agent_task_registered().await { Ok(agent_task) => agent_task, Err(error) => { diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index a6c4c6bc18..867fc0961e 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -27,14 +27,6 @@ impl TurnSkillsContext { #[derive(Clone, Debug, PartialEq)] pub(crate) struct TurnStartTranscriptInput { pub(crate) input: Vec, - pub(crate) user_prompt_submit_outcome: Option, - pub(crate) user_prompt_recorded: bool, -} - -#[derive(Clone, Debug, PartialEq)] -pub(crate) struct TurnStartUserPromptSubmitOutcome { - pub(crate) should_stop: bool, - pub(crate) additional_contexts: Vec, } /// The context needed for a single turn of the thread. diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 957f3521ee..dd7ffc9340 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -22,6 +22,7 @@ use tracing::warn; use crate::contextual_user_message::TURN_ABORTED_CLOSE_TAG; 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::inspect_pending_input; use crate::hook_runtime::record_additional_contexts; @@ -306,8 +307,6 @@ impl Session { .lock_turn_start_transcript_inputs() .push(TurnStartTranscriptInput { input: input.clone(), - user_prompt_submit_outcome: None, - user_prompt_recorded: false, }); } let mut active = self.active_turn.lock().await; @@ -629,6 +628,24 @@ 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. + if reason == TurnAbortReason::Interrupted && task.kind == TaskKind::Regular { + 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; + } + } + select! { _ = task.done.notified() => { }, @@ -639,16 +656,6 @@ impl Session { task.handle.abort(); - if reason == TurnAbortReason::Interrupted && task.kind == TaskKind::Regular { - 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).await; - } - } - let session_ctx = Arc::new(SessionTaskContext::new(Arc::clone(self))); session_task .abort(session_ctx, Arc::clone(&task.turn_context))