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/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<ThreadInputState>) {
-        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<ThreadInputState>,
+        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::<VecDeque<_>>();
+                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,

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());

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(),

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();

codex-rs/tui/src/chatwidget/user_messages.rs

@@ -132,12 +132,18 @@ pub(crate) struct ThreadInputState {
     pub(super) queued_user_messages: VecDeque<QueuedUserMessage>,
     pub(super) queued_user_message_history_records: VecDeque<UserMessageHistoryRecord>,
     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<CollaborationModeMask>,
     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<String> for UserMessage {
     fn from(text: String) -> Self {
         Self {