Restore TUI turn context after backtracking

Snapshot per-turn TUI context during committed history so rollback can restore the surviving turn's cwd, approval and sandbox settings, collaboration mode, service tier, personality, and related local state.

Add regression coverage for pending and replayed rollback flows.

Co-authored-by: Codex <noreply@openai.com>

codex-rs/tui/src/app.rs

@@ -2577,6 +2577,7 @@ impl App {
                     tui.frame_requester().schedule_frame();
                 }
                 self.transcript_cells.push(cell.clone());
+                self.track_backtrack_transcript_cell(&cell);
                 let mut display = cell.display_lines(tui.terminal.last_known_screen_size.width);
                 if !display.is_empty() {
                     // Only insert a separating blank line for new cells that are not
@@ -4134,7 +4135,16 @@ impl App {
         };
     }
 
-    fn refresh_status_line(&mut self) {
+    pub(crate) fn set_runtime_policy_overrides(
+        &mut self,
+        approval_policy: AskForApproval,
+        sandbox_policy: SandboxPolicy,
+    ) {
+        self.runtime_approval_policy_override = Some(approval_policy);
+        self.runtime_sandbox_policy_override = Some(sandbox_policy);
+    }
+
+    pub(crate) fn refresh_status_line(&mut self) {
         self.chat_widget.refresh_status_line();
     }
 
@@ -4172,6 +4182,7 @@ mod tests {
     use super::*;
     use crate::app_backtrack::BacktrackSelection;
     use crate::app_backtrack::BacktrackState;
+    use crate::app_backtrack::BacktrackTurnContextSnapshot;
     use crate::app_backtrack::user_count;
     use crate::chatwidget::tests::make_chatwidget_manual_with_sender;
     use crate::chatwidget::tests::set_chatgpt_auth;
@@ -4185,14 +4196,18 @@ mod tests {
     use codex_core::CodexAuth;
     use codex_core::config::ConfigBuilder;
     use codex_core::config::ConfigOverrides;
+    use codex_core::config::types::ApprovalsReviewer;
     use codex_core::config::types::ModelAvailabilityNuxConfig;
     use codex_otel::SessionTelemetry;
     use codex_protocol::ThreadId;
     use codex_protocol::config_types::CollaborationMode;
     use codex_protocol::config_types::CollaborationModeMask;
     use codex_protocol::config_types::ModeKind;
+    use codex_protocol::config_types::Personality;
+    use codex_protocol::config_types::ServiceTier;
     use codex_protocol::config_types::Settings;
     use codex_protocol::openai_models::ModelAvailabilityNux;
+    use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
     use codex_protocol::protocol::AgentMessageDeltaEvent;
     use codex_protocol::protocol::AskForApproval;
     use codex_protocol::protocol::Event;
@@ -6442,6 +6457,58 @@ smart_approvals = true
         )
     }
 
+    fn backtrack_test_session_configured(
+        session_id: ThreadId,
+        cwd: PathBuf,
+    ) -> SessionConfiguredEvent {
+        SessionConfiguredEvent {
+            session_id,
+            forked_from_id: None,
+            thread_name: None,
+            model: "gpt-test".to_string(),
+            model_provider_id: "test-provider".to_string(),
+            service_tier: None,
+            approval_policy: AskForApproval::Never,
+            approvals_reviewer: ApprovalsReviewer::User,
+            sandbox_policy: SandboxPolicy::new_read_only_policy(),
+            cwd,
+            reasoning_effort: None,
+            history_log_id: 0,
+            history_entry_count: 0,
+            initial_messages: None,
+            network_proxy: None,
+            rollout_path: Some(PathBuf::new()),
+        }
+    }
+
+    fn backtrack_user_cell(message: &str) -> Arc<dyn HistoryCell> {
+        Arc::new(UserHistoryCell {
+            message: message.to_string(),
+            text_elements: Vec::new(),
+            local_image_paths: Vec::new(),
+            remote_image_urls: Vec::new(),
+        })
+    }
+
+    fn assert_backtrack_turn_context(app: &App, expected: &BacktrackTurnContextSnapshot) {
+        assert_eq!(
+            app.capture_backtrack_turn_context_snapshot(),
+            expected.clone()
+        );
+        assert_eq!(app.config.cwd, expected.cwd.clone());
+        assert_eq!(app.config.approvals_reviewer, expected.approvals_reviewer);
+        assert_eq!(app.config.service_tier, expected.service_tier);
+        assert_eq!(app.config.personality, expected.personality);
+        assert_eq!(
+            app.runtime_approval_policy_override,
+            Some(expected.approval_policy)
+        );
+        assert_eq!(
+            app.runtime_sandbox_policy_override,
+            Some(expected.sandbox_policy.clone())
+        );
+    }
+
     fn next_user_turn_op(op_rx: &mut tokio::sync::mpsc::UnboundedReceiver<Op>) -> Op {
         let mut seen = Vec::new();
         while let Ok(op) = op_rx.try_recv() {
@@ -7356,6 +7423,221 @@ smart_approvals = true
         assert_eq!(user_messages, vec!["first prompt".to_string()]);
     }
 
+    #[tokio::test]
+    async fn pending_backtrack_restores_previous_surviving_turn_context() {
+        let mut app = make_test_app().await;
+        let session_id = ThreadId::new();
+        let session_event =
+            backtrack_test_session_configured(session_id, PathBuf::from("/tmp/backtrack-base"));
+        app.chat_widget.handle_codex_event(Event {
+            id: "session-configured".to_string(),
+            msg: EventMsg::SessionConfigured(session_event.clone()),
+        });
+        app.config = app.chat_widget.config_ref().clone();
+
+        let session_cell = Arc::new(new_session_info(
+            app.chat_widget.config_ref(),
+            "gpt-test",
+            session_event,
+            false,
+            None,
+            None,
+            false,
+        )) as Arc<dyn HistoryCell>;
+        app.track_backtrack_transcript_cell(&session_cell);
+        app.transcript_cells.push(session_cell);
+        let baseline_snapshot = app.capture_backtrack_turn_context_snapshot();
+
+        let turn_one_snapshot = BacktrackTurnContextSnapshot {
+            cwd: PathBuf::from("/tmp/backtrack-turn-one"),
+            approval_policy: AskForApproval::OnRequest,
+            approvals_reviewer: ApprovalsReviewer::GuardianSubagent,
+            sandbox_policy: SandboxPolicy::new_workspace_write_policy(),
+            current_collaboration_mode: CollaborationMode {
+                mode: ModeKind::Default,
+                settings: Settings {
+                    model: "gpt-turn-one".to_string(),
+                    reasoning_effort: Some(ReasoningEffortConfig::Low),
+                    developer_instructions: None,
+                },
+            },
+            active_collaboration_mask: Some(CollaborationModeMask {
+                name: "plan".to_string(),
+                mode: Some(ModeKind::Plan),
+                model: Some("gpt-turn-one".to_string()),
+                reasoning_effort: Some(Some(ReasoningEffortConfig::High)),
+                developer_instructions: None,
+            }),
+            service_tier: Some(ServiceTier::Fast),
+            personality: Some(Personality::Pragmatic),
+        };
+        app.apply_backtrack_turn_context_snapshot(&turn_one_snapshot);
+        let user_one = backtrack_user_cell("first prompt");
+        app.track_backtrack_transcript_cell(&user_one);
+        app.transcript_cells.push(user_one);
+
+        let turn_two_snapshot = BacktrackTurnContextSnapshot {
+            cwd: PathBuf::from("/tmp/backtrack-turn-two"),
+            approval_policy: AskForApproval::Never,
+            approvals_reviewer: ApprovalsReviewer::User,
+            sandbox_policy: SandboxPolicy::new_read_only_policy(),
+            current_collaboration_mode: CollaborationMode {
+                mode: ModeKind::Default,
+                settings: Settings {
+                    model: "gpt-turn-two".to_string(),
+                    reasoning_effort: Some(ReasoningEffortConfig::Medium),
+                    developer_instructions: None,
+                },
+            },
+            active_collaboration_mask: Some(CollaborationModeMask {
+                name: "default".to_string(),
+                mode: Some(ModeKind::Default),
+                model: Some("gpt-turn-two".to_string()),
+                reasoning_effort: Some(Some(ReasoningEffortConfig::Medium)),
+                developer_instructions: None,
+            }),
+            service_tier: None,
+            personality: Some(Personality::Friendly),
+        };
+        app.apply_backtrack_turn_context_snapshot(&turn_two_snapshot);
+        let user_two = backtrack_user_cell("second prompt");
+        app.track_backtrack_transcript_cell(&user_two);
+        app.transcript_cells.push(user_two);
+
+        app.backtrack.pending_rollback = Some(crate::app_backtrack::PendingBacktrackRollback {
+            selection: crate::app_backtrack::BacktrackSelection {
+                nth_user_message: 1,
+                prefill: String::new(),
+                text_elements: Vec::new(),
+                local_image_paths: Vec::new(),
+                remote_image_urls: Vec::new(),
+            },
+            thread_id: app.chat_widget.thread_id(),
+        });
+
+        app.handle_backtrack_event(&EventMsg::ThreadRolledBack(ThreadRolledBackEvent {
+            num_turns: 1,
+        }));
+
+        let user_messages: Vec<String> = app
+            .transcript_cells
+            .iter()
+            .filter_map(|cell| {
+                cell.as_any()
+                    .downcast_ref::<UserHistoryCell>()
+                    .map(|cell| cell.message.clone())
+            })
+            .collect();
+        assert_eq!(user_messages, vec!["first prompt".to_string()]);
+        assert_backtrack_turn_context(&app, &turn_one_snapshot);
+        assert_eq!(
+            app.backtrack.session_baseline_turn_context,
+            Some(baseline_snapshot)
+        );
+        assert_eq!(
+            app.backtrack.current_session_turn_contexts,
+            vec![turn_one_snapshot]
+        );
+    }
+
+    #[tokio::test]
+    async fn non_pending_backtrack_restores_session_baseline_turn_context() {
+        let mut app = make_test_app().await;
+        let session_id = ThreadId::new();
+        let session_event =
+            backtrack_test_session_configured(session_id, PathBuf::from("/tmp/backtrack-base"));
+        app.chat_widget.handle_codex_event(Event {
+            id: "session-configured".to_string(),
+            msg: EventMsg::SessionConfigured(session_event.clone()),
+        });
+        app.config = app.chat_widget.config_ref().clone();
+
+        let session_cell = Arc::new(new_session_info(
+            app.chat_widget.config_ref(),
+            "gpt-test",
+            session_event,
+            false,
+            None,
+            None,
+            false,
+        )) as Arc<dyn HistoryCell>;
+        app.track_backtrack_transcript_cell(&session_cell);
+        app.transcript_cells.push(session_cell);
+        let baseline_snapshot = app.capture_backtrack_turn_context_snapshot();
+
+        let turn_snapshot = BacktrackTurnContextSnapshot {
+            cwd: PathBuf::from("/tmp/backtrack-turn"),
+            approval_policy: AskForApproval::OnRequest,
+            approvals_reviewer: ApprovalsReviewer::GuardianSubagent,
+            sandbox_policy: SandboxPolicy::new_workspace_write_policy(),
+            current_collaboration_mode: CollaborationMode {
+                mode: ModeKind::Default,
+                settings: Settings {
+                    model: "gpt-turn".to_string(),
+                    reasoning_effort: Some(ReasoningEffortConfig::Low),
+                    developer_instructions: None,
+                },
+            },
+            active_collaboration_mask: Some(CollaborationModeMask {
+                name: "plan".to_string(),
+                mode: Some(ModeKind::Plan),
+                model: Some("gpt-turn".to_string()),
+                reasoning_effort: Some(Some(ReasoningEffortConfig::High)),
+                developer_instructions: None,
+            }),
+            service_tier: Some(ServiceTier::Fast),
+            personality: Some(Personality::Pragmatic),
+        };
+        app.apply_backtrack_turn_context_snapshot(&turn_snapshot);
+        let user_cell = backtrack_user_cell("first prompt");
+        app.track_backtrack_transcript_cell(&user_cell);
+        app.transcript_cells.push(user_cell);
+
+        let stale_snapshot = BacktrackTurnContextSnapshot {
+            cwd: PathBuf::from("/tmp/backtrack-stale"),
+            approval_policy: AskForApproval::Never,
+            approvals_reviewer: ApprovalsReviewer::User,
+            sandbox_policy: SandboxPolicy::new_read_only_policy(),
+            current_collaboration_mode: CollaborationMode {
+                mode: ModeKind::Default,
+                settings: Settings {
+                    model: "gpt-stale".to_string(),
+                    reasoning_effort: Some(ReasoningEffortConfig::Medium),
+                    developer_instructions: None,
+                },
+            },
+            active_collaboration_mask: Some(CollaborationModeMask {
+                name: "default".to_string(),
+                mode: Some(ModeKind::Default),
+                model: Some("gpt-stale".to_string()),
+                reasoning_effort: Some(Some(ReasoningEffortConfig::Medium)),
+                developer_instructions: None,
+            }),
+            service_tier: None,
+            personality: Some(Personality::Friendly),
+        };
+        app.apply_backtrack_turn_context_snapshot(&stale_snapshot);
+
+        assert!(app.apply_non_pending_thread_rollback(1));
+
+        let user_messages: Vec<String> = app
+            .transcript_cells
+            .iter()
+            .filter_map(|cell| {
+                cell.as_any()
+                    .downcast_ref::<UserHistoryCell>()
+                    .map(|cell| cell.message.clone())
+            })
+            .collect();
+        assert_eq!(user_messages, Vec::<String>::new());
+        assert_backtrack_turn_context(&app, &baseline_snapshot);
+        assert_eq!(
+            app.backtrack.session_baseline_turn_context,
+            Some(baseline_snapshot)
+        );
+        assert_eq!(app.backtrack.current_session_turn_contexts, Vec::new());
+    }
+
     #[tokio::test]
     async fn queued_rollback_syncs_overlay_and_clears_deferred_history() {
         let mut app = make_test_app().await;

codex-rs/tui/src/app_backtrack.rs

@@ -34,11 +34,18 @@ use crate::history_cell::UserHistoryCell;
 use crate::pager_overlay::Overlay;
 use crate::tui;
 use crate::tui::TuiEvent;
+use codex_core::config::types::ApprovalsReviewer;
 use codex_protocol::ThreadId;
+use codex_protocol::config_types::CollaborationMode;
+use codex_protocol::config_types::CollaborationModeMask;
+use codex_protocol::config_types::Personality;
+use codex_protocol::config_types::ServiceTier;
+use codex_protocol::protocol::AskForApproval;
 use codex_protocol::protocol::CodexErrorInfo;
 use codex_protocol::protocol::ErrorEvent;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::Op;
+use codex_protocol::protocol::SandboxPolicy;
 use codex_protocol::user_input::TextElement;
 use color_eyre::eyre::Result;
 use crossterm::event::KeyCode;
@@ -66,6 +73,22 @@ pub(crate) struct BacktrackState {
     /// This acts as a guardrail: once we request a rollback, we block additional backtrack
     /// submissions until core responds with either a success or failure event.
     pub(crate) pending_rollback: Option<PendingBacktrackRollback>,
+    /// Session baseline turn context from the latest session start.
+    pub(crate) session_baseline_turn_context: Option<BacktrackTurnContextSnapshot>,
+    /// Turn-context snapshots aligned with committed user turns in the current session.
+    pub(crate) current_session_turn_contexts: Vec<BacktrackTurnContextSnapshot>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct BacktrackTurnContextSnapshot {
+    pub(crate) cwd: PathBuf,
+    pub(crate) approval_policy: AskForApproval,
+    pub(crate) approvals_reviewer: ApprovalsReviewer,
+    pub(crate) sandbox_policy: SandboxPolicy,
+    pub(crate) current_collaboration_mode: CollaborationMode,
+    pub(crate) active_collaboration_mask: Option<CollaborationModeMask>,
+    pub(crate) service_tier: Option<ServiceTier>,
+    pub(crate) personality: Option<Personality>,
 }
 
 /// A user-visible backtrack choice that can be confirmed into a rollback request.
@@ -461,6 +484,94 @@ impl App {
         tui.frame_requester().schedule_frame();
     }
 
+    pub(crate) fn capture_backtrack_turn_context_snapshot(&self) -> BacktrackTurnContextSnapshot {
+        BacktrackTurnContextSnapshot {
+            cwd: self.chat_widget.config_ref().cwd.clone(),
+            approval_policy: self
+                .chat_widget
+                .config_ref()
+                .permissions
+                .approval_policy
+                .value(),
+            approvals_reviewer: self.chat_widget.config_ref().approvals_reviewer,
+            sandbox_policy: self
+                .chat_widget
+                .config_ref()
+                .permissions
+                .sandbox_policy
+                .get()
+                .clone(),
+            current_collaboration_mode: self.chat_widget.current_collaboration_mode().clone(),
+            active_collaboration_mask: self.chat_widget.active_collaboration_mask().cloned(),
+            service_tier: self.chat_widget.current_service_tier(),
+            personality: self.chat_widget.config_ref().personality,
+        }
+    }
+
+    pub(crate) fn track_backtrack_transcript_cell(
+        &mut self,
+        cell: &Arc<dyn crate::history_cell::HistoryCell>,
+    ) {
+        if cell.as_any().is::<SessionInfoCell>() {
+            self.backtrack.session_baseline_turn_context =
+                Some(self.capture_backtrack_turn_context_snapshot());
+            self.backtrack.current_session_turn_contexts.clear();
+            return;
+        }
+        if cell.as_any().is::<UserHistoryCell>() {
+            self.backtrack
+                .current_session_turn_contexts
+                .push(self.capture_backtrack_turn_context_snapshot());
+        }
+    }
+
+    pub(crate) fn apply_backtrack_turn_context_snapshot(
+        &mut self,
+        snapshot: &BacktrackTurnContextSnapshot,
+    ) {
+        self.config.cwd = snapshot.cwd.clone();
+        self.config.approvals_reviewer = snapshot.approvals_reviewer;
+        self.config.service_tier = snapshot.service_tier;
+        self.config.personality = snapshot.personality;
+        self.config.model_reasoning_effort = snapshot.current_collaboration_mode.reasoning_effort();
+        if let Err(err) = self
+            .config
+            .permissions
+            .approval_policy
+            .set(snapshot.approval_policy)
+        {
+            tracing::warn!(%err, "failed to restore approval_policy after rollback");
+        }
+        if let Err(err) = self
+            .config
+            .permissions
+            .sandbox_policy
+            .set(snapshot.sandbox_policy.clone())
+        {
+            tracing::warn!(%err, "failed to restore sandbox_policy after rollback");
+        }
+        self.set_runtime_policy_overrides(
+            snapshot.approval_policy,
+            snapshot.sandbox_policy.clone(),
+        );
+        self.file_search.update_search_dir(self.config.cwd.clone());
+        self.chat_widget.restore_backtrack_turn_context(snapshot);
+        self.refresh_status_line();
+    }
+
+    fn restore_backtrack_turn_context_after_trim(&mut self) {
+        let Some(snapshot) = self
+            .backtrack
+            .current_session_turn_contexts
+            .last()
+            .cloned()
+            .or_else(|| self.backtrack.session_baseline_turn_context.clone())
+        else {
+            return;
+        };
+        self.apply_backtrack_turn_context_snapshot(&snapshot);
+    }
+
     pub(crate) fn handle_backtrack_event(&mut self, event: &EventMsg) {
         match event {
             EventMsg::ThreadRolledBack(rollback) => {
@@ -500,6 +611,15 @@ impl App {
         if !trim_transcript_cells_drop_last_n_user_turns(&mut self.transcript_cells, num_turns) {
             return false;
         }
+        let turns_from_end = usize::try_from(num_turns).unwrap_or(usize::MAX);
+        if turns_from_end >= self.backtrack.current_session_turn_contexts.len() {
+            self.backtrack.current_session_turn_contexts.clear();
+        } else {
+            self.backtrack
+                .current_session_turn_contexts
+                .truncate(self.backtrack.current_session_turn_contexts.len() - turns_from_end);
+        }
+        self.restore_backtrack_turn_context_after_trim();
         self.sync_overlay_after_transcript_trim();
         self.backtrack_render_pending = true;
         true
@@ -521,6 +641,10 @@ impl App {
             &mut self.transcript_cells,
             pending.selection.nth_user_message,
         ) {
+            self.backtrack
+                .current_session_turn_contexts
+                .truncate(pending.selection.nth_user_message);
+            self.restore_backtrack_turn_context_after_trim();
             self.sync_overlay_after_transcript_trim();
             self.backtrack_render_pending = true;
         }

codex-rs/tui/src/chatwidget.rs

@@ -38,6 +38,7 @@ use std::time::Duration;
 use std::time::Instant;
 
 use self::realtime::PendingSteerCompareKey;
+use crate::app_backtrack::BacktrackTurnContextSnapshot;
 use crate::app_event::RealtimeAudioDeviceKind;
 #[cfg(all(not(target_os = "linux"), feature = "voice-input"))]
 use crate::audio_device::list_realtime_audio_device_names;
@@ -7869,6 +7870,38 @@ impl ChatWidget {
         self.config.approvals_reviewer = policy;
     }
 
+    pub(crate) fn restore_backtrack_turn_context(
+        &mut self,
+        snapshot: &BacktrackTurnContextSnapshot,
+    ) {
+        self.config.cwd = snapshot.cwd.clone();
+        self.current_cwd = Some(snapshot.cwd.clone());
+        if let Err(err) = self
+            .config
+            .permissions
+            .approval_policy
+            .set(snapshot.approval_policy)
+        {
+            tracing::warn!(%err, "failed to restore approval_policy on chat config");
+        }
+        if let Err(err) = self
+            .config
+            .permissions
+            .sandbox_policy
+            .set(snapshot.sandbox_policy.clone())
+        {
+            tracing::warn!(%err, "failed to restore sandbox_policy on chat config");
+        }
+        self.config.approvals_reviewer = snapshot.approvals_reviewer;
+        self.config.service_tier = snapshot.service_tier;
+        self.config.personality = snapshot.personality;
+        self.current_collaboration_mode = snapshot.current_collaboration_mode.clone();
+        self.active_collaboration_mask = snapshot.active_collaboration_mask.clone();
+        self.update_collaboration_mode_indicator();
+        self.refresh_model_display();
+        self.request_redraw();
+    }
+
     pub(crate) fn set_full_access_warning_acknowledged(&mut self, acknowledged: bool) {
         self.config.notices.hide_full_access_warning = Some(acknowledged);
     }
@@ -8089,6 +8122,10 @@ impl ChatWidget {
         &self.current_collaboration_mode
     }
 
+    pub(crate) fn active_collaboration_mask(&self) -> Option<&CollaborationModeMask> {
+        self.active_collaboration_mask.as_ref()
+    }
+
     pub(crate) fn current_reasoning_effort(&self) -> Option<ReasoningEffortConfig> {
         self.effective_reasoning_effort()
     }