Restore collaboration mode on backtrack failure

codex-rs/tui/src/app.rs

@@ -2524,6 +2524,8 @@ mod tests {
     use codex_core::config::ConfigOverrides;
     use codex_core::models_manager::manager::ModelsManager;
     use codex_core::protocol::AskForApproval;
+    use codex_core::protocol::CodexErrorInfo;
+    use codex_core::protocol::ErrorEvent;
     use codex_core::protocol::Event;
     use codex_core::protocol::EventMsg;
     use codex_core::protocol::SandboxPolicy;
@@ -3219,6 +3221,159 @@ mod tests {
         assert_eq!(collaboration_mode.mode, ModeKind::Code);
     }
 
+    #[tokio::test]
+    async fn backtrack_to_custom_clears_collaboration_mode() {
+        let (mut app, _app_event_rx, mut op_rx) = make_test_app_with_channels().await;
+        app.chat_widget
+            .set_feature_enabled(Feature::CollaborationModes, true);
+
+        let plan_mask = plan_mask();
+        let thread_id = ThreadId::new();
+        app.chat_widget.handle_codex_event(Event {
+            id: String::new(),
+            msg: EventMsg::SessionConfigured(SessionConfiguredEvent {
+                session_id: thread_id,
+                forked_from_id: None,
+                thread_name: None,
+                model: "gpt-test".to_string(),
+                model_provider_id: "test-provider".to_string(),
+                approval_policy: AskForApproval::Never,
+                sandbox_policy: SandboxPolicy::ReadOnly,
+                cwd: PathBuf::from("/home/user/project"),
+                reasoning_effort: None,
+                history_log_id: 0,
+                history_entry_count: 0,
+                initial_messages: None,
+                rollout_path: Some(PathBuf::new()),
+            }),
+        });
+
+        app.chat_widget.set_collaboration_mask(plan_mask.clone());
+
+        app.transcript_cells = vec![
+            Arc::new(UserHistoryCell {
+                message: "custom".to_string(),
+                text_elements: Vec::new(),
+                local_image_paths: Vec::new(),
+                collaboration_mode: None,
+            }) as Arc<dyn HistoryCell>,
+            Arc::new(AgentMessageCell::new(
+                vec![Line::from("custom response")],
+                true,
+            )) as Arc<dyn HistoryCell>,
+            Arc::new(UserHistoryCell {
+                message: "plan".to_string(),
+                text_elements: Vec::new(),
+                local_image_paths: Vec::new(),
+                collaboration_mode: Some(plan_mask),
+            }) as Arc<dyn HistoryCell>,
+        ];
+
+        app.backtrack.base_id = Some(thread_id);
+        app.backtrack.primed = true;
+        app.backtrack.nth_user_message = 0;
+
+        let selection = app
+            .confirm_backtrack_from_main()
+            .expect("backtrack selection");
+        app.apply_backtrack_rollback(selection);
+
+        assert_eq!(
+            app.chat_widget.active_collaboration_mode_kind(),
+            ModeKind::Custom
+        );
+
+        app.chat_widget
+            .set_composer_text("follow-up".to_string(), Vec::new(), Vec::new());
+        app.chat_widget
+            .handle_key_event(KeyEvent::from(KeyCode::Enter));
+
+        let op = next_user_turn(&mut op_rx);
+        let Op::UserTurn {
+            collaboration_mode, ..
+        } = op
+        else {
+            panic!("expected Op::UserTurn");
+        };
+        assert_eq!(collaboration_mode, None);
+    }
+
+    #[tokio::test]
+    async fn backtrack_failed_restores_previous_collaboration_mode() {
+        let (mut app, _app_event_rx, _op_rx) = make_test_app_with_channels().await;
+        app.chat_widget
+            .set_feature_enabled(Feature::CollaborationModes, true);
+
+        let plan_mask = plan_mask();
+        let code_mask = code_mask();
+
+        let thread_id = ThreadId::new();
+        app.chat_widget.handle_codex_event(Event {
+            id: String::new(),
+            msg: EventMsg::SessionConfigured(SessionConfiguredEvent {
+                session_id: thread_id,
+                forked_from_id: None,
+                thread_name: None,
+                model: "gpt-test".to_string(),
+                model_provider_id: "test-provider".to_string(),
+                approval_policy: AskForApproval::Never,
+                sandbox_policy: SandboxPolicy::ReadOnly,
+                cwd: PathBuf::from("/home/user/project"),
+                reasoning_effort: None,
+                history_log_id: 0,
+                history_entry_count: 0,
+                initial_messages: None,
+                rollout_path: Some(PathBuf::new()),
+            }),
+        });
+
+        app.chat_widget.set_collaboration_mask(code_mask.clone());
+
+        app.transcript_cells = vec![
+            Arc::new(UserHistoryCell {
+                message: "plan".to_string(),
+                text_elements: Vec::new(),
+                local_image_paths: Vec::new(),
+                collaboration_mode: Some(plan_mask.clone()),
+            }) as Arc<dyn HistoryCell>,
+            Arc::new(AgentMessageCell::new(
+                vec![Line::from("plan response")],
+                true,
+            )) as Arc<dyn HistoryCell>,
+            Arc::new(UserHistoryCell {
+                message: "code".to_string(),
+                text_elements: Vec::new(),
+                local_image_paths: Vec::new(),
+                collaboration_mode: Some(code_mask),
+            }) as Arc<dyn HistoryCell>,
+        ];
+
+        app.backtrack.base_id = Some(thread_id);
+        app.backtrack.primed = true;
+        app.backtrack.nth_user_message = 0;
+
+        let selection = app
+            .confirm_backtrack_from_main()
+            .expect("backtrack selection");
+        app.apply_backtrack_rollback(selection);
+
+        assert_eq!(
+            app.chat_widget.active_collaboration_mode_kind(),
+            ModeKind::Plan
+        );
+
+        app.handle_backtrack_event(&EventMsg::Error(ErrorEvent {
+            message: "rollback rejected".to_string(),
+            codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed),
+        }));
+
+        assert_eq!(
+            app.chat_widget.active_collaboration_mode_kind(),
+            ModeKind::Code
+        );
+        assert!(app.backtrack.pending_rollback.is_none());
+    }
+
     #[tokio::test]
     async fn new_session_requests_shutdown_for_previous_conversation() {
         let (mut app, mut app_event_rx, mut op_rx) = make_test_app_with_channels().await;

codex-rs/tui/src/app_backtrack.rs

@@ -97,6 +97,7 @@ pub(crate) struct BacktrackSelection {
 pub(crate) struct PendingBacktrackRollback {
     pub(crate) selection: BacktrackSelection,
     pub(crate) thread_id: Option<ThreadId>,
+    pub(crate) previous_collaboration_mask: Option<CollaborationModeMask>,
 }
 
 impl App {
@@ -206,15 +207,16 @@ impl App {
             return;
         }
 
-        if let Some(mask) = selection.collaboration_mode.clone() {
-            self.chat_widget.set_collaboration_mask(mask);
-        }
+        let previous_collaboration_mask = self
+            .chat_widget
+            .replace_collaboration_mask(selection.collaboration_mode.clone());
         let prefill = selection.prefill.clone();
         let text_elements = selection.text_elements.clone();
         let local_image_paths = selection.local_image_paths.clone();
         self.backtrack.pending_rollback = Some(PendingBacktrackRollback {
             selection,
             thread_id: self.chat_widget.thread_id(),
+            previous_collaboration_mask,
         });
         self.chat_widget.submit_op(Op::ThreadRollback { num_turns });
         if !prefill.is_empty() || !text_elements.is_empty() || !local_image_paths.is_empty() {
@@ -464,8 +466,14 @@ impl App {
                 codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed),
                 ..
             }) => {
-                // Core rejected the rollback; clear the guard so the user can retry.
-                self.backtrack.pending_rollback = None;
+                let Some(pending) = self.backtrack.pending_rollback.take() else {
+                    return;
+                };
+                if pending.thread_id != self.chat_widget.thread_id() {
+                    return;
+                }
+                self.chat_widget
+                    .replace_collaboration_mask(pending.previous_collaboration_mask);
             }
             _ => {}
         }

codex-rs/tui/src/chatwidget.rs

@@ -5515,6 +5515,21 @@ impl ChatWidget {
         self.request_redraw();
     }
 
+    pub(crate) fn replace_collaboration_mask(
+        &mut self,
+        mask: Option<CollaborationModeMask>,
+    ) -> Option<CollaborationModeMask> {
+        if !self.collaboration_modes_enabled() {
+            return None;
+        }
+        let previous = self.active_collaboration_mask.clone();
+        self.active_collaboration_mask = mask;
+        self.update_collaboration_mode_indicator();
+        self.refresh_model_display();
+        self.request_redraw();
+        previous
+    }
+
     fn connectors_enabled(&self) -> bool {
         self.config.features.enabled(Feature::Apps)
     }