Sync collaboration mode on rollback

codex-rs/core/src/codex.rs

@@ -2578,11 +2578,17 @@ mod handlers {
     use codex_protocol::request_user_input::RequestUserInputResponse;
 
     use crate::context_manager::is_user_turn_boundary;
+    use crate::models_manager::collaboration_mode_presets::mask_from_instructions;
     use codex_protocol::config_types::CollaborationMode;
+    use codex_protocol::config_types::CollaborationModeMask;
     use codex_protocol::config_types::ModeKind;
     use codex_protocol::config_types::Settings;
     use codex_protocol::dynamic_tools::DynamicToolResponse;
     use codex_protocol::mcp::RequestId as ProtocolRequestId;
+    use codex_protocol::models::ContentItem;
+    use codex_protocol::models::ResponseItem;
+    use codex_protocol::protocol::COLLABORATION_MODE_CLOSE_TAG;
+    use codex_protocol::protocol::COLLABORATION_MODE_OPEN_TAG;
     use codex_protocol::user_input::UserInput;
     use codex_rmcp_client::ElicitationAction;
     use codex_rmcp_client::ElicitationResponse;
@@ -2976,6 +2982,13 @@ mod handlers {
 
         let mut history = sess.clone_history().await;
         history.drop_last_n_user_turns(num_turns);
+        if let Some(mask) = last_collaboration_mask(history.raw_items()) {
+            let mut state = sess.state.lock().await;
+            state.session_configuration.collaboration_mode = state
+                .session_configuration
+                .collaboration_mode
+                .apply_mask(&mask);
+        }
 
         // Replace with the raw items. We don't want to replace with a normalized
         // version of the history.
@@ -3102,6 +3115,33 @@ mod handlers {
         true
     }
 
+    fn last_collaboration_mask(items: &[ResponseItem]) -> Option<CollaborationModeMask> {
+        items.iter().rev().find_map(|item| {
+            let ResponseItem::Message { role, content, .. } = item else {
+                return None;
+            };
+            if role != "developer" {
+                return None;
+            }
+            let text = content.iter().find_map(|item| {
+                if let ContentItem::InputText { text } = item {
+                    Some(text.as_str())
+                } else {
+                    None
+                }
+            })?;
+            let instructions = extract_collaboration_instructions(text)?;
+            mask_from_instructions(instructions)
+        })
+    }
+
+    fn extract_collaboration_instructions(text: &str) -> Option<&str> {
+        let start = text.find(COLLABORATION_MODE_OPEN_TAG)? + COLLABORATION_MODE_OPEN_TAG.len();
+        let rest = text.get(start..)?;
+        let end = rest.find(COLLABORATION_MODE_CLOSE_TAG)?;
+        rest.get(..end)
+    }
+
     pub async fn review(
         sess: &Arc<Session>,
         config: &Arc<Config>,

codex-rs/core/src/models_manager/collaboration_mode_presets.rs

@@ -18,6 +18,18 @@ pub(super) fn builtin_collaboration_mode_presets() -> Vec<CollaborationModeMask>
     ]
 }
 
+pub(crate) fn mask_from_instructions(instructions: &str) -> Option<CollaborationModeMask> {
+    let normalized = instructions.trim();
+    builtin_collaboration_mode_presets()
+        .into_iter()
+        .find(|mask| {
+            mask.developer_instructions
+                .as_ref()
+                .and_then(|value| value.as_ref())
+                .is_some_and(|text| text.trim() == normalized)
+        })
+}
+
 #[cfg(any(test, feature = "test-support"))]
 pub fn test_builtin_collaboration_mode_presets() -> Vec<CollaborationModeMask> {
     builtin_collaboration_mode_presets()

codex-rs/tui/src/app_backtrack.rs

@@ -206,10 +206,9 @@ impl App {
             return;
         }
 
-        let collaboration_override = selection
-            .collaboration_mode
-            .clone()
-            .and_then(|mask| self.chat_widget.set_collaboration_mask_from_history(mask));
+        if let Some(mask) = selection.collaboration_mode.clone() {
+            self.chat_widget.set_collaboration_mask(mask);
+        }
         let prefill = selection.prefill.clone();
         let text_elements = selection.text_elements.clone();
         let local_image_paths = selection.local_image_paths.clone();
@@ -218,21 +217,6 @@ impl App {
             thread_id: self.chat_widget.thread_id(),
         });
         self.chat_widget.submit_op(Op::ThreadRollback { num_turns });
-        if let Some(collaboration_mode) = collaboration_override
-            && self.chat_widget.thread_id().is_some()
-        {
-            self.chat_widget.submit_op(Op::OverrideTurnContext {
-                cwd: None,
-                approval_policy: None,
-                sandbox_policy: None,
-                windows_sandbox_level: None,
-                model: None,
-                effort: None,
-                summary: None,
-                collaboration_mode: Some(collaboration_mode),
-                personality: None,
-            });
-        }
         if !prefill.is_empty() || !text_elements.is_empty() || !local_image_paths.is_empty() {
             self.chat_widget
                 .set_composer_text(prefill, text_elements, local_image_paths);

codex-rs/tui/src/chatwidget.rs

@@ -5515,18 +5515,6 @@ impl ChatWidget {
         self.request_redraw();
     }
 
-    pub(crate) fn set_collaboration_mask_from_history(
-        &mut self,
-        mask: CollaborationModeMask,
-    ) -> Option<CollaborationMode> {
-        if !self.collaboration_modes_enabled() {
-            return None;
-        }
-        let next_mode = self.current_collaboration_mode.apply_mask(&mask);
-        self.set_collaboration_mask(mask);
-        Some(next_mode)
-    }
-
     fn connectors_enabled(&self) -> bool {
         self.config.features.enabled(Feature::Apps)
     }