From 24f3af135c70ac48df23711c3634acd8b40cae05 Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Fri, 30 Jan 2026 13:20:44 -0800 Subject: [PATCH] Add backtrack collaboration mode tests --- codex-rs/tui/src/app.rs | 202 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 202 insertions(+) diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 33766c1d90..6d7791f479 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -2529,6 +2529,8 @@ mod tests { use codex_core::protocol::SessionSource; use codex_otel::OtelManager; use codex_protocol::ThreadId; + use codex_protocol::config_types::CollaborationModeMask; + use codex_protocol::config_types::ModeKind; use codex_protocol::user_input::TextElement; use insta::assert_snapshot; use pretty_assertions::assert_eq; @@ -2703,6 +2705,39 @@ mod tests { ) } + fn next_user_turn(op_rx: &mut tokio::sync::mpsc::UnboundedReceiver) -> Op { + loop { + match op_rx.try_recv() { + Ok(op @ Op::UserTurn { .. }) => return op, + Ok(_) => continue, + Err(TryRecvError::Empty) => panic!("expected Op::UserTurn but queue was empty"), + Err(TryRecvError::Disconnected) => { + panic!("expected Op::UserTurn but channel closed") + } + } + } + } + + fn plan_mask() -> CollaborationModeMask { + CollaborationModeMask { + name: "Plan".to_string(), + mode: Some(ModeKind::Plan), + model: None, + reasoning_effort: None, + developer_instructions: None, + } + } + + fn code_mask() -> CollaborationModeMask { + CollaborationModeMask { + name: "Code".to_string(), + mode: Some(ModeKind::Code), + model: None, + reasoning_effort: None, + developer_instructions: None, + } + } + fn test_otel_manager(config: &Config, model: &str) -> OtelManager { let model_info = ModelsManager::construct_model_info_offline(model, config); OtelManager::new( @@ -3014,6 +3049,173 @@ mod tests { assert_eq!(rollback_turns, Some(1)); } + #[tokio::test] + async fn backtrack_to_plan_preserves_plan_mode_on_submit() { + 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 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, + Arc::new(AgentMessageCell::new( + vec![Line::from("plan response")], + true, + )) as Arc, + Arc::new(UserHistoryCell { + message: "code".to_string(), + text_elements: Vec::new(), + local_image_paths: Vec::new(), + collaboration_mode: Some(code_mask), + }) as Arc, + ]; + + 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.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"); + }; + let collaboration_mode = collaboration_mode.expect("expected collaboration mode"); + assert_eq!(collaboration_mode.mode, ModeKind::Plan); + } + + #[tokio::test] + async fn backtrack_plan_then_shift_tab_sends_code_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 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, + Arc::new(AgentMessageCell::new( + vec![Line::from("plan response")], + true, + )) as Arc, + Arc::new(UserHistoryCell { + message: "code".to_string(), + text_elements: Vec::new(), + local_image_paths: Vec::new(), + collaboration_mode: Some(code_mask), + }) as Arc, + ]; + + 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.chat_widget + .handle_key_event(KeyEvent::from(KeyCode::BackTab)); + assert_eq!( + app.chat_widget.active_collaboration_mode_kind(), + ModeKind::Code + ); + + 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"); + }; + let collaboration_mode = collaboration_mode.expect("expected collaboration mode"); + assert_eq!(collaboration_mode.mode, ModeKind::Code); + } + #[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;