From 77a3f4e80f4bda601e3a204089bd13479001ea7d Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Wed, 15 Jul 2026 01:42:26 +0000 Subject: [PATCH] Preserve in-flight state when restoring thread input (#33203) ## What changed - Add an explicit restore mode for thread input snapshots and use it when replaying thread snapshots. - Preserve the pending-steer interrupt flag alongside running-turn and queued-input state when restoring an in-flight turn. - When in-flight state is not preserved, clear stale running flags and move pending steers ahead of already queued follow-ups while retaining their history records. ## Testing - Expand TUI tests to cover both restoration modes, including sleep-inhibitor state, pending-steer ordering, history records, and queued submission. GitOrigin-RevId: f4e8b2bdf065ef6c6644b7f0ea9cf32c8a9af619 --- codex-rs/tui/src/app/thread_routing.rs | 9 +- codex-rs/tui/src/chatwidget.rs | 1 + codex-rs/tui/src/chatwidget/input_restore.rs | 76 +++++++++++----- .../chatwidget/tests/composer_submission.rs | 91 ++++++++++++++++--- .../tui/src/chatwidget/tests/review_mode.rs | 38 ++++---- .../src/chatwidget/tests/slash_commands.rs | 21 ++++- codex-rs/tui/src/chatwidget/user_messages.rs | 6 ++ 7 files changed, 184 insertions(+), 58 deletions(-) diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 17dc36d9b6..171f3968ed 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -6,6 +6,7 @@ use super::session_lifecycle::ThreadAttachPresentation; use super::*; +use crate::chatwidget::ThreadInputStateRestoreMode; use crate::session_resume::read_session_model; impl App { @@ -1344,8 +1345,12 @@ impl App { } self.chat_widget .set_queue_autosend_suppressed(/*suppressed*/ true); - self.chat_widget - .restore_thread_input_state(snapshot.input_state); + self.chat_widget.restore_thread_input_state( + snapshot.input_state, + ThreadInputStateRestoreMode { + preserve_in_flight_turn: true, + }, + ); if !snapshot.turns.is_empty() { self.chat_widget .replay_thread_turns(snapshot.turns, ReplayKind::ThreadSnapshot); diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 662f1d5618..d8a76cf3aa 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -427,6 +427,7 @@ use self::user_messages::QueuedUserMessage; use self::user_messages::ShellEscapePolicy; use self::user_messages::ThreadComposerState; pub(crate) use self::user_messages::ThreadInputState; +pub(crate) use self::user_messages::ThreadInputStateRestoreMode; pub(crate) use self::user_messages::UserMessage; use self::user_messages::UserMessageDisplay; #[cfg(test)] diff --git a/codex-rs/tui/src/chatwidget/input_restore.rs b/codex-rs/tui/src/chatwidget/input_restore.rs index 3aaa276247..7505f626e2 100644 --- a/codex-rs/tui/src/chatwidget/input_restore.rs +++ b/codex-rs/tui/src/chatwidget/input_restore.rs @@ -1,6 +1,7 @@ //! Input queue restore and thread-input snapshot behavior for `ChatWidget`. use std::collections::HashSet; +use std::collections::VecDeque; use super::user_messages::remap_colliding_paste_placeholders; use super::*; @@ -338,6 +339,9 @@ impl ChatWidget { .queued_user_message_history_records .clone(), user_turn_pending_start: self.input_queue.user_turn_pending_start, + submit_pending_steers_after_interrupt: self + .input_queue + .submit_pending_steers_after_interrupt, current_collaboration_mode: self.current_collaboration_mode.clone(), active_collaboration_mask: self.active_collaboration_mask.clone(), task_running: self.bottom_pane.is_task_running(), @@ -345,15 +349,26 @@ impl ChatWidget { }) } - pub(crate) fn restore_thread_input_state(&mut self, input_state: Option) { - let restored_task_running = input_state.as_ref().is_some_and(|state| state.task_running); + pub(crate) fn restore_thread_input_state( + &mut self, + input_state: Option, + restore_mode: ThreadInputStateRestoreMode, + ) { + let preserve_in_flight_turn = restore_mode.preserve_in_flight_turn; + let restored_task_running = + preserve_in_flight_turn && input_state.as_ref().is_some_and(|state| state.task_running); if let Some(input_state) = input_state { self.current_collaboration_mode = input_state.current_collaboration_mode; self.active_collaboration_mask = input_state.active_collaboration_mask; self.safety_buffering_prompt = input_state.safety_buffering_prompt; - self.turn_lifecycle - .restore_running(input_state.agent_turn_running, Instant::now()); - self.input_queue.user_turn_pending_start = input_state.user_turn_pending_start; + self.turn_lifecycle.restore_running( + preserve_in_flight_turn && input_state.agent_turn_running, + Instant::now(), + ); + self.input_queue.user_turn_pending_start = + preserve_in_flight_turn && input_state.user_turn_pending_start; + self.input_queue.submit_pending_steers_after_interrupt = + preserve_in_flight_turn && input_state.submit_pending_steers_after_interrupt; self.update_collaboration_mode_indicator(); self.refresh_model_dependent_surfaces(); self.restore_composer_state(input_state.composer.unwrap_or_default()); @@ -363,22 +378,37 @@ impl ChatWidget { UserMessageHistoryRecord::UserMessageText, ); let mut pending_steer_compare_keys = input_state.pending_steer_compare_keys; - self.input_queue.pending_steers = input_state - .pending_steers - .into_iter() - .zip(pending_steer_history_records) - .map(|(user_message, history_record)| PendingSteer { - compare_key: pending_steer_compare_keys.pop_front().unwrap_or_else(|| { - PendingSteerCompareKey { - message: user_message.text.clone(), - image_count: user_message.local_images.len() - + user_message.remote_image_urls.len(), - } - }), - history_record, - user_message, - }) - .collect(); + let pending_steers = input_state.pending_steers; + let mut queued_user_messages = input_state.queued_user_messages; + let mut queued_user_message_history_records = + input_state.queued_user_message_history_records; + if preserve_in_flight_turn { + self.input_queue.pending_steers = pending_steers + .into_iter() + .zip(pending_steer_history_records) + .map(|(user_message, history_record)| PendingSteer { + compare_key: pending_steer_compare_keys.pop_front().unwrap_or_else(|| { + PendingSteerCompareKey { + message: user_message.text.clone(), + image_count: user_message.local_images.len() + + user_message.remote_image_urls.len(), + } + }), + history_record, + user_message, + }) + .collect(); + } else { + self.input_queue.pending_steers.clear(); + let mut safety_retry_follow_ups = pending_steers + .into_iter() + .map(QueuedUserMessage::from) + .collect::>(); + safety_retry_follow_ups.append(&mut queued_user_messages); + queued_user_messages = safety_retry_follow_ups; + pending_steer_history_records.append(&mut queued_user_message_history_records); + queued_user_message_history_records = pending_steer_history_records; + } self.input_queue.rejected_steers_queue = input_state.rejected_steers_queue; self.input_queue.rejected_steer_history_records = input_state.rejected_steer_history_records; @@ -386,9 +416,9 @@ impl ChatWidget { self.input_queue.rejected_steers_queue.len(), UserMessageHistoryRecord::UserMessageText, ); - self.input_queue.queued_user_messages = input_state.queued_user_messages; + self.input_queue.queued_user_messages = queued_user_messages; self.input_queue.queued_user_message_history_records = - input_state.queued_user_message_history_records; + queued_user_message_history_records; self.input_queue.queued_user_message_history_records.resize( self.input_queue.queued_user_messages.len(), UserMessageHistoryRecord::UserMessageText, diff --git a/codex-rs/tui/src/chatwidget/tests/composer_submission.rs b/codex-rs/tui/src/chatwidget/tests/composer_submission.rs index 920b224265..f858457e4f 100644 --- a/codex-rs/tui/src/chatwidget/tests/composer_submission.rs +++ b/codex-rs/tui/src/chatwidget/tests/composer_submission.rs @@ -1068,38 +1068,101 @@ async fn pending_steer_interrupt_uses_remapped_binding() { } #[tokio::test] -async fn restore_thread_input_state_syncs_sleep_inhibitor_state() { - let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; +async fn restore_thread_input_state_applies_running_state_policy() { + let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.thread_id = Some(ThreadId::new()); chat.set_feature_enabled(Feature::PreventIdleSleep, /*enabled*/ true); - chat.restore_thread_input_state(Some(ThreadInputState { - composer: None, + let pending_history = UserMessageHistoryRecord::Override(UserMessageHistoryOverride { + text: "submitted history".to_string(), + text_elements: Vec::new(), + }); + let queued_history = UserMessageHistoryRecord::Override(UserMessageHistoryOverride { + text: "queued history".to_string(), + text_elements: Vec::new(), + }); + let input_state = ThreadInputState { + composer: Some(ThreadComposerState { + text: "composer draft".to_string(), + ..Default::default() + }), safety_buffering_prompt: Some(UserMessage::from("buffered prompt")), - pending_steers: VecDeque::new(), - pending_steer_history_records: VecDeque::new(), + pending_steers: VecDeque::from([UserMessage::from("submitted to the interrupted turn")]), + pending_steer_history_records: VecDeque::from([pending_history.clone()]), pending_steer_compare_keys: VecDeque::new(), rejected_steers_queue: VecDeque::new(), rejected_steer_history_records: VecDeque::new(), - queued_user_messages: VecDeque::new(), - queued_user_message_history_records: VecDeque::new(), - user_turn_pending_start: false, + queued_user_messages: VecDeque::from([UserMessage::from("already queued").into()]), + queued_user_message_history_records: VecDeque::from([queued_history.clone()]), + user_turn_pending_start: true, + submit_pending_steers_after_interrupt: true, current_collaboration_mode: chat.current_collaboration_mode.clone(), active_collaboration_mask: chat.active_collaboration_mask.clone(), task_running: true, agent_turn_running: true, - })); + }; + chat.restore_thread_input_state( + Some(input_state.clone()), + ThreadInputStateRestoreMode { + preserve_in_flight_turn: true, + }, + ); assert!(chat.turn_lifecycle.agent_turn_running); assert!(chat.turn_lifecycle.sleep_inhibitor.is_turn_running()); assert!(chat.bottom_pane.is_task_running()); + assert!(chat.input_queue.user_turn_pending_start); + assert!(chat.input_queue.submit_pending_steers_after_interrupt); + let captured_input_state = chat + .capture_thread_input_state() + .expect("thread input state"); + assert!(captured_input_state.submit_pending_steers_after_interrupt); assert_eq!( - chat.capture_thread_input_state() - .expect("thread input state") - .safety_buffering_prompt, + captured_input_state.safety_buffering_prompt, + Some(UserMessage::from("buffered prompt")) + ); + assert_eq!(chat.input_queue.pending_steers.len(), 1); + assert_eq!( + chat.safety_buffering_prompt, Some(UserMessage::from("buffered prompt")) ); - chat.restore_thread_input_state(/*input_state*/ None); + chat.restore_thread_input_state( + Some(input_state), + ThreadInputStateRestoreMode { + preserve_in_flight_turn: false, + }, + ); + + assert!(!chat.turn_lifecycle.agent_turn_running); + assert!(!chat.turn_lifecycle.sleep_inhibitor.is_turn_running()); + assert!(!chat.bottom_pane.is_task_running()); + assert!(!chat.input_queue.user_turn_pending_start); + assert!(!chat.input_queue.submit_pending_steers_after_interrupt); + assert!(chat.input_queue.pending_steers.is_empty()); + assert_eq!(chat.bottom_pane.composer_text(), "composer draft"); + assert_eq!( + chat.safety_buffering_prompt, + Some(UserMessage::from("buffered prompt")) + ); + assert_eq!( + chat.queued_user_message_texts(), + vec!["submitted to the interrupted turn", "already queued"] + ); + assert_eq!( + chat.input_queue.queued_user_message_history_records, + VecDeque::from([pending_history, queued_history]) + ); + assert!(chat.maybe_send_next_queued_input()); + assert_matches!(next_submit_op(&mut op_rx), Op::UserTurn { .. }); + assert_eq!(chat.queued_user_message_texts(), vec!["already queued"]); + + chat.restore_thread_input_state( + /*input_state*/ None, + ThreadInputStateRestoreMode { + preserve_in_flight_turn: true, + }, + ); assert!(!chat.turn_lifecycle.agent_turn_running); assert!(!chat.turn_lifecycle.sleep_inhibitor.is_turn_running()); diff --git a/codex-rs/tui/src/chatwidget/tests/review_mode.rs b/codex-rs/tui/src/chatwidget/tests/review_mode.rs index 9ef9c46a53..b1f3360c0d 100644 --- a/codex-rs/tui/src/chatwidget/tests/review_mode.rs +++ b/codex-rs/tui/src/chatwidget/tests/review_mode.rs @@ -373,22 +373,28 @@ async fn restore_thread_input_state_restores_pending_steers_without_downgrading_ let mut queued_user_messages = VecDeque::new(); queued_user_messages.push_back(UserMessage::from("queued draft").into()); - chat.restore_thread_input_state(Some(ThreadInputState { - composer: None, - safety_buffering_prompt: None, - pending_steers, - pending_steer_history_records: VecDeque::new(), - pending_steer_compare_keys, - rejected_steers_queue, - rejected_steer_history_records: VecDeque::new(), - queued_user_messages, - queued_user_message_history_records: VecDeque::new(), - user_turn_pending_start: false, - current_collaboration_mode: chat.current_collaboration_mode.clone(), - active_collaboration_mask: chat.active_collaboration_mask.clone(), - task_running: false, - agent_turn_running: false, - })); + chat.restore_thread_input_state( + Some(ThreadInputState { + composer: None, + safety_buffering_prompt: None, + pending_steers, + pending_steer_history_records: VecDeque::new(), + pending_steer_compare_keys, + rejected_steers_queue, + rejected_steer_history_records: VecDeque::new(), + queued_user_messages, + queued_user_message_history_records: VecDeque::new(), + user_turn_pending_start: false, + submit_pending_steers_after_interrupt: false, + current_collaboration_mode: chat.current_collaboration_mode.clone(), + active_collaboration_mask: chat.active_collaboration_mask.clone(), + task_running: false, + agent_turn_running: false, + }), + ThreadInputStateRestoreMode { + preserve_in_flight_turn: true, + }, + ); assert_eq!( chat.queued_user_message_texts(), diff --git a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs index c3a43e0e9e..ab3848eedb 100644 --- a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs +++ b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs @@ -539,9 +539,19 @@ async fn queued_inline_rename_does_not_drain_again_before_turn_started() { ); let input_state = chat.capture_thread_input_state().unwrap(); assert!(input_state.user_turn_pending_start); - chat.restore_thread_input_state(/*input_state*/ None); + chat.restore_thread_input_state( + /*input_state*/ None, + ThreadInputStateRestoreMode { + preserve_in_flight_turn: true, + }, + ); assert!(!chat.input_queue.user_turn_pending_start); - chat.restore_thread_input_state(Some(input_state)); + chat.restore_thread_input_state( + Some(input_state), + ThreadInputStateRestoreMode { + preserve_in_flight_turn: true, + }, + ); assert!(chat.input_queue.user_turn_pending_start); assert_eq!( chat.queued_user_message_texts(), @@ -987,7 +997,12 @@ async fn restored_queued_goal_slash_command_emits_set_goal_event() { let (mut restored_chat, mut restored_rx, mut restored_op_rx) = make_chatwidget_manual(/*model_override*/ None).await; restored_chat.set_feature_enabled(Feature::Goals, /*enabled*/ true); - restored_chat.restore_thread_input_state(Some(input_state)); + restored_chat.restore_thread_input_state( + Some(input_state), + ThreadInputStateRestoreMode { + preserve_in_flight_turn: true, + }, + ); let thread_id = ThreadId::new(); restored_chat.thread_id = Some(thread_id); restored_chat.maybe_send_next_queued_input(); diff --git a/codex-rs/tui/src/chatwidget/user_messages.rs b/codex-rs/tui/src/chatwidget/user_messages.rs index d166692a65..255ebb6e11 100644 --- a/codex-rs/tui/src/chatwidget/user_messages.rs +++ b/codex-rs/tui/src/chatwidget/user_messages.rs @@ -132,12 +132,18 @@ pub(crate) struct ThreadInputState { pub(super) queued_user_messages: VecDeque, pub(super) queued_user_message_history_records: VecDeque, pub(super) user_turn_pending_start: bool, + pub(super) submit_pending_steers_after_interrupt: bool, pub(super) current_collaboration_mode: CollaborationMode, pub(super) active_collaboration_mask: Option, pub(super) task_running: bool, pub(super) agent_turn_running: bool, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct ThreadInputStateRestoreMode { + pub(crate) preserve_in_flight_turn: bool, +} + impl From for UserMessage { fn from(text: String) -> Self { Self {