tui: remove draft-only collaboration mode override state

codex-rs/tui/src/chatwidget.rs

@@ -607,6 +607,7 @@ pub(crate) struct UserMessage {
     local_images: Vec<LocalImageAttachment>,
     text_elements: Vec<TextElement>,
     mention_paths: HashMap<String, String>,
+    collaboration_mode_override: Option<CollaborationModeMask>,
 }
 
 impl From<String> for UserMessage {
@@ -617,6 +618,7 @@ impl From<String> for UserMessage {
             // Plain text conversion has no UI element ranges.
             text_elements: Vec::new(),
             mention_paths: HashMap::new(),
+            collaboration_mode_override: None,
         }
     }
 }
@@ -629,6 +631,7 @@ impl From<&str> for UserMessage {
             // Plain text conversion has no UI element ranges.
             text_elements: Vec::new(),
             mention_paths: HashMap::new(),
+            collaboration_mode_override: None,
         }
     }
 }
@@ -655,6 +658,7 @@ pub(crate) fn create_initial_user_message(
             local_images,
             text_elements,
             mention_paths: HashMap::new(),
+            collaboration_mode_override: None,
         })
     }
 }
@@ -669,6 +673,7 @@ fn remap_placeholders_for_message(message: UserMessage, next_label: &mut usize)
         text_elements,
         local_images,
         mention_paths,
+        collaboration_mode_override,
     } = message;
     if local_images.is_empty() {
         return UserMessage {
@@ -676,6 +681,7 @@ fn remap_placeholders_for_message(message: UserMessage, next_label: &mut usize)
             text_elements,
             local_images,
             mention_paths,
+            collaboration_mode_override,
         };
     }
 
@@ -731,6 +737,7 @@ fn remap_placeholders_for_message(message: UserMessage, next_label: &mut usize)
         local_images: remapped_images,
         text_elements: rebuilt_elements,
         mention_paths,
+        collaboration_mode_override,
     }
 }
 
@@ -1404,16 +1411,9 @@ impl ChatWidget {
         }
 
         if let Some(combined) = self.drain_queued_messages_for_restore() {
-            let combined_local_image_paths = combined
-                .local_images
-                .iter()
-                .map(|img| img.path.clone())
-                .collect();
-            self.bottom_pane.set_composer_text(
-                combined.text,
-                combined.text_elements,
-                combined_local_image_paths,
-            );
+            if let Err(message) = self.restore_user_message_to_composer(combined) {
+                self.queued_user_messages.push_front(message);
+            }
             self.refresh_queued_user_messages();
         }
 
@@ -1437,6 +1437,7 @@ impl ChatWidget {
             text_elements: self.bottom_pane.composer_text_elements(),
             local_images: self.bottom_pane.composer_local_images(),
             mention_paths: HashMap::new(),
+            collaboration_mode_override: None,
         };
 
         let mut to_merge: Vec<UserMessage> = self.queued_user_messages.drain(..).collect();
@@ -1449,9 +1450,14 @@ impl ChatWidget {
             text_elements: Vec::new(),
             local_images: Vec::new(),
             mention_paths: HashMap::new(),
+            collaboration_mode_override: None,
         };
         let mut combined_offset = 0usize;
         let mut next_image_label = 1usize;
+        let mut combined_collaboration_mode_override = None;
+        // The restored composer can carry only one override. If merged drafts disagree on mode,
+        // we drop the override instead of arbitrarily picking one.
+        let mut has_conflicting_mode_overrides = false;
 
         for (idx, message) in to_merge.into_iter().enumerate() {
             if idx > 0 {
@@ -1459,6 +1465,16 @@ impl ChatWidget {
                 combined_offset += 1;
             }
             let message = remap_placeholders_for_message(message, &mut next_image_label);
+            // Preserve an override only when every merged draft points to the same mode.
+            if let Some(mask) = &message.collaboration_mode_override {
+                if let Some(existing) = &combined_collaboration_mode_override
+                    && existing != mask
+                {
+                    has_conflicting_mode_overrides = true;
+                } else if combined_collaboration_mode_override.is_none() {
+                    combined_collaboration_mode_override = Some(mask.clone());
+                }
+            }
             let base = combined_offset;
             combined.text.push_str(&message.text);
             combined_offset += message.text.len();
@@ -1472,10 +1488,45 @@ impl ChatWidget {
             combined.local_images.extend(message.local_images);
             combined.mention_paths.extend(message.mention_paths);
         }
+        if !has_conflicting_mode_overrides {
+            combined.collaboration_mode_override = combined_collaboration_mode_override;
+        }
 
         Some(combined)
     }
 
+    fn restore_user_message_to_composer(
+        &mut self,
+        user_message: UserMessage,
+    ) -> Result<(), UserMessage> {
+        let UserMessage {
+            text,
+            local_images,
+            text_elements,
+            mention_paths,
+            collaboration_mode_override,
+        } = user_message;
+        if let Some(mask) = collaboration_mode_override {
+            if self.agent_turn_running && self.active_collaboration_mask.as_ref() != Some(&mask) {
+                self.add_error_message(
+                    "Cannot switch collaboration mode while a turn is running.".to_string(),
+                );
+                return Err(UserMessage {
+                    text,
+                    local_images,
+                    text_elements,
+                    mention_paths,
+                    collaboration_mode_override: Some(mask),
+                });
+            }
+            self.set_collaboration_mask(mask);
+        }
+        let local_image_paths = local_images.into_iter().map(|img| img.path).collect();
+        self.bottom_pane
+            .set_composer_text(text, text_elements, local_image_paths);
+        Ok(())
+    }
+
     fn on_plan_update(&mut self, update: UpdatePlanArgs) {
         self.saw_plan_update_this_turn = true;
         self.add_to_history(history_cell::new_plan_update(update));
@@ -2687,16 +2738,9 @@ impl ChatWidget {
             } if !self.queued_user_messages.is_empty() => {
                 // Prefer the most recently queued item.
                 if let Some(user_message) = self.queued_user_messages.pop_back() {
-                    let local_image_paths = user_message
-                        .local_images
-                        .iter()
-                        .map(|img| img.path.clone())
-                        .collect();
-                    self.bottom_pane.set_composer_text(
-                        user_message.text,
-                        user_message.text_elements,
-                        local_image_paths,
-                    );
+                    if let Err(user_message) = self.restore_user_message_to_composer(user_message) {
+                        self.queued_user_messages.push_back(user_message);
+                    }
                     self.refresh_queued_user_messages();
                     self.request_redraw();
                 }
@@ -2713,9 +2757,10 @@ impl ChatWidget {
                             .take_recent_submission_images_with_placeholders(),
                         text_elements,
                         mention_paths: self.bottom_pane.take_mention_paths(),
+                        collaboration_mode_override: None,
                     };
-                    if self.is_session_configured() {
-                        // Submitted is only emitted when steer is enabled (Enter sends immediately).
+                    if self.is_session_configured() && !self.is_plan_streaming_in_tui() {
+                        // Submitted is only emitted when steer is enabled.
                         // Reset any reasoning header only when we are actually submitting a turn.
                         self.reasoning_buffer.clear();
                         self.full_reasoning_buffer.clear();
@@ -2736,6 +2781,7 @@ impl ChatWidget {
                             .take_recent_submission_images_with_placeholders(),
                         text_elements,
                         mention_paths: self.bottom_pane.take_mention_paths(),
+                        collaboration_mode_override: None,
                     };
                     self.queue_user_message(user_message);
                 }
@@ -3097,6 +3143,7 @@ impl ChatWidget {
                         .take_recent_submission_images_with_placeholders(),
                     text_elements: prepared_elements,
                     mention_paths: self.bottom_pane.take_mention_paths(),
+                    collaboration_mode_override: None,
                 };
                 if self.is_session_configured() {
                     self.reasoning_buffer.clear();
@@ -3234,6 +3281,7 @@ impl ChatWidget {
             local_images,
             text_elements,
             mention_paths,
+            collaboration_mode_override,
         } = user_message;
         if text.is_empty() && local_images.is_empty() {
             return;
@@ -3304,6 +3352,15 @@ impl ChatWidget {
             }
         }
 
+        if let Some(mask) = collaboration_mode_override {
+            if self.agent_turn_running && self.active_collaboration_mask.as_ref() != Some(&mask) {
+                self.add_error_message(
+                    "Cannot switch collaboration mode while a turn is running.".to_string(),
+                );
+                return;
+            }
+            self.set_collaboration_mask(mask);
+        }
         let effective_mode = self.effective_collaboration_mode();
         let collaboration_mode = if self.collaboration_modes_enabled() {
             self.active_collaboration_mask
@@ -5860,6 +5917,10 @@ impl ChatWidget {
         self.bottom_pane.is_task_running() || self.is_review_mode
     }
 
+    fn is_plan_streaming_in_tui(&self) -> bool {
+        self.plan_stream_controller.is_some()
+    }
+
     pub(crate) fn composer_is_empty(&self) -> bool {
         self.bottom_pane.composer_is_empty()
     }
@@ -5869,8 +5930,19 @@ impl ChatWidget {
         text: String,
         collaboration_mode: CollaborationModeMask,
     ) {
-        self.set_collaboration_mask(collaboration_mode);
-        self.submit_user_message(text.into());
+        let should_queue = self.is_plan_streaming_in_tui();
+        let user_message = UserMessage {
+            text,
+            local_images: Vec::new(),
+            text_elements: Vec::new(),
+            mention_paths: HashMap::new(),
+            collaboration_mode_override: Some(collaboration_mode),
+        };
+        if should_queue {
+            self.queue_user_message(user_message);
+        } else {
+            self.submit_user_message(user_message);
+        }
     }
 
     /// True when the UI is in the regular composer state with no running task,

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

@@ -404,6 +404,7 @@ async fn interrupted_turn_restores_queued_messages_with_images_and_elements() {
         }],
         text_elements: first_elements,
         mention_paths: HashMap::new(),
+        collaboration_mode_override: None,
     });
     chat.queued_user_messages.push_back(UserMessage {
         text: second_text,
@@ -413,6 +414,7 @@ async fn interrupted_turn_restores_queued_messages_with_images_and_elements() {
         }],
         text_elements: second_elements,
         mention_paths: HashMap::new(),
+        collaboration_mode_override: None,
     });
     chat.refresh_queued_user_messages();
 
@@ -462,6 +464,57 @@ async fn interrupted_turn_restores_queued_messages_with_images_and_elements() {
     );
 }
 
+#[tokio::test]
+async fn interrupted_turn_restore_preserves_mode_override_for_resubmission() {
+    let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(Some("gpt-5")).await;
+    chat.thread_id = Some(ThreadId::new());
+    chat.set_feature_enabled(Feature::CollaborationModes, true);
+
+    let plan_mask = collaboration_modes::plan_mask(chat.models_manager.as_ref())
+        .expect("expected plan collaboration mode");
+    let code_mask = collaboration_modes::default_mask(chat.models_manager.as_ref())
+        .expect("expected default collaboration mode");
+    let expected_mode = code_mask
+        .mode
+        .expect("expected mode kind on code collaboration mode");
+
+    chat.set_collaboration_mask(plan_mask);
+    chat.on_task_started();
+    chat.queued_user_messages.push_back(UserMessage {
+        text: "Implement the plan.".to_string(),
+        local_images: Vec::new(),
+        text_elements: Vec::new(),
+        mention_paths: HashMap::new(),
+        collaboration_mode_override: Some(code_mask.clone()),
+    });
+    chat.refresh_queued_user_messages();
+
+    chat.handle_codex_event(Event {
+        id: "interrupt".into(),
+        msg: EventMsg::TurnAborted(codex_core::protocol::TurnAbortedEvent {
+            reason: TurnAbortReason::Interrupted,
+        }),
+    });
+
+    assert_eq!(chat.bottom_pane.composer_text(), "Implement the plan.");
+    assert!(chat.queued_user_messages.is_empty());
+    assert_eq!(chat.active_collaboration_mode_kind(), expected_mode);
+
+    chat.handle_key_event(KeyEvent::from(KeyCode::Enter));
+
+    match next_submit_op(&mut op_rx) {
+        Op::UserTurn {
+            collaboration_mode: Some(CollaborationMode { mode, .. }),
+            personality: None,
+            ..
+        } => assert_eq!(mode, expected_mode),
+        other => {
+            panic!("expected Op::UserTurn with restored mode override, got {other:?}")
+        }
+    }
+    assert_eq!(chat.active_collaboration_mode_kind(), expected_mode);
+}
+
 #[tokio::test]
 async fn remap_placeholders_uses_attachment_labels() {
     let placeholder_one = "[Image #1]";
@@ -493,6 +546,7 @@ async fn remap_placeholders_uses_attachment_labels() {
         text_elements: elements,
         local_images: attachments,
         mention_paths: HashMap::new(),
+        collaboration_mode_override: None,
     };
     let mut next_label = 3usize;
     let remapped = remap_placeholders_for_message(message, &mut next_label);
@@ -554,6 +608,7 @@ async fn remap_placeholders_uses_byte_ranges_when_placeholder_missing() {
         text_elements: elements,
         local_images: attachments,
         mention_paths: HashMap::new(),
+        collaboration_mode_override: None,
     };
     let mut next_label = 3usize;
     let remapped = remap_placeholders_for_message(message, &mut next_label);
@@ -1284,6 +1339,97 @@ async fn submit_user_message_with_mode_sets_coding_collaboration_mode() {
     }
 }
 
+#[tokio::test]
+async fn submit_user_message_with_mode_errors_when_mode_changes_during_running_turn() {
+    let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(Some("gpt-5")).await;
+    chat.thread_id = Some(ThreadId::new());
+    chat.set_feature_enabled(Feature::CollaborationModes, true);
+    let plan_mask =
+        collaboration_modes::mask_for_kind(chat.models_manager.as_ref(), ModeKind::Plan)
+            .expect("expected plan collaboration mask");
+    chat.set_collaboration_mask(plan_mask);
+    chat.on_task_started();
+
+    let default_mode = collaboration_modes::default_mask(chat.models_manager.as_ref())
+        .expect("expected default collaboration mode");
+    chat.submit_user_message_with_mode("Implement the plan.".to_string(), default_mode);
+
+    assert_eq!(chat.active_collaboration_mode_kind(), ModeKind::Plan);
+    assert!(chat.queued_user_messages.is_empty());
+    assert_matches!(op_rx.try_recv(), Err(TryRecvError::Empty));
+    let rendered = drain_insert_history(&mut rx)
+        .iter()
+        .map(|lines| lines_to_single_string(lines))
+        .collect::<Vec<_>>()
+        .join("\n");
+    assert!(
+        rendered.contains("Cannot switch collaboration mode while a turn is running."),
+        "expected running-turn error message, got: {rendered:?}"
+    );
+}
+
+#[tokio::test]
+async fn submit_user_message_with_mode_allows_same_mode_during_running_turn() {
+    let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(Some("gpt-5")).await;
+    chat.thread_id = Some(ThreadId::new());
+    chat.set_feature_enabled(Feature::CollaborationModes, true);
+    let plan_mask =
+        collaboration_modes::mask_for_kind(chat.models_manager.as_ref(), ModeKind::Plan)
+            .expect("expected plan collaboration mask");
+    chat.set_collaboration_mask(plan_mask.clone());
+    chat.on_task_started();
+
+    chat.submit_user_message_with_mode("Continue planning.".to_string(), plan_mask);
+
+    assert_eq!(chat.active_collaboration_mode_kind(), ModeKind::Plan);
+    assert!(chat.queued_user_messages.is_empty());
+    match next_submit_op(&mut op_rx) {
+        Op::UserTurn {
+            collaboration_mode:
+                Some(CollaborationMode {
+                    mode: ModeKind::Plan,
+                    ..
+                }),
+            personality: None,
+            ..
+        } => {}
+        other => {
+            panic!("expected Op::UserTurn with plan collab mode, got {other:?}")
+        }
+    }
+}
+
+#[tokio::test]
+async fn submit_user_message_with_mode_submits_when_plan_stream_is_not_active() {
+    let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(Some("gpt-5")).await;
+    chat.thread_id = Some(ThreadId::new());
+    chat.set_feature_enabled(Feature::CollaborationModes, true);
+    let plan_mask =
+        collaboration_modes::mask_for_kind(chat.models_manager.as_ref(), ModeKind::Plan)
+            .expect("expected plan collaboration mask");
+    chat.set_collaboration_mask(plan_mask);
+
+    let default_mode = collaboration_modes::default_mask(chat.models_manager.as_ref())
+        .expect("expected default collaboration mode");
+    let expected_mode = default_mode
+        .mode
+        .expect("expected default collaboration mode kind");
+    chat.submit_user_message_with_mode("Implement the plan.".to_string(), default_mode);
+
+    assert_eq!(chat.active_collaboration_mode_kind(), expected_mode);
+    assert!(chat.queued_user_messages.is_empty());
+    match next_submit_op(&mut op_rx) {
+        Op::UserTurn {
+            collaboration_mode: Some(CollaborationMode { mode, .. }),
+            personality: None,
+            ..
+        } => assert_eq!(mode, expected_mode),
+        other => {
+            panic!("expected Op::UserTurn with default collab mode, got {other:?}")
+        }
+    }
+}
+
 #[tokio::test]
 async fn plan_implementation_popup_skips_replayed_turn_complete() {
     let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5")).await;
@@ -1849,6 +1995,55 @@ async fn exec_begin_restores_status_indicator_after_preamble() {
     assert_eq!(chat.bottom_pane.status_indicator_visible(), true);
 }
 
+#[tokio::test]
+async fn steer_enter_queues_while_plan_stream_is_active() {
+    let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(None).await;
+    chat.thread_id = Some(ThreadId::new());
+    chat.set_feature_enabled(Feature::CollaborationModes, true);
+    let plan_mask =
+        collaboration_modes::mask_for_kind(chat.models_manager.as_ref(), ModeKind::Plan)
+            .expect("expected plan collaboration mask");
+    chat.set_collaboration_mask(plan_mask);
+    chat.on_task_started();
+    chat.on_plan_delta("- Step 1".to_string());
+
+    chat.bottom_pane
+        .set_composer_text("queued submission".to_string(), Vec::new(), Vec::new());
+    chat.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE));
+
+    assert_eq!(chat.active_collaboration_mode_kind(), ModeKind::Plan);
+    assert_eq!(chat.queued_user_messages.len(), 1);
+    assert_eq!(
+        chat.queued_user_messages.front().unwrap().text,
+        "queued submission"
+    );
+    assert_matches!(op_rx.try_recv(), Err(TryRecvError::Empty));
+}
+
+#[tokio::test]
+async fn steer_enter_submits_when_plan_stream_is_not_active() {
+    let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(None).await;
+    chat.thread_id = Some(ThreadId::new());
+    chat.set_feature_enabled(Feature::CollaborationModes, true);
+    let plan_mask =
+        collaboration_modes::mask_for_kind(chat.models_manager.as_ref(), ModeKind::Plan)
+            .expect("expected plan collaboration mask");
+    chat.set_collaboration_mask(plan_mask);
+    chat.on_task_started();
+
+    chat.bottom_pane
+        .set_composer_text("submitted immediately".to_string(), Vec::new(), Vec::new());
+    chat.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE));
+
+    assert!(chat.queued_user_messages.is_empty());
+    match next_submit_op(&mut op_rx) {
+        Op::UserTurn {
+            personality: None, ..
+        } => {}
+        other => panic!("expected Op::UserTurn, got {other:?}"),
+    }
+}
+
 #[tokio::test]
 async fn ctrl_c_shutdown_works_with_caps_lock() {
     let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await;