From 0174a3db489af15d60a7100aba3655a79ea05151 Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Sun, 1 Feb 2026 16:11:54 -0800 Subject: [PATCH] Simplify backtrack mode sync and generalize pending sync naming --- codex-rs/core/src/codex.rs | 6 ++++-- codex-rs/core/src/state/session.rs | 10 ++++----- codex-rs/tui/src/app.rs | 34 +++++++++++++++++++++++++----- codex-rs/tui/src/app_backtrack.rs | 27 +++--------------------- codex-rs/tui/src/chatwidget.rs | 19 ----------------- 5 files changed, 41 insertions(+), 55 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 61ae4a9972..bf13e154b6 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2792,9 +2792,10 @@ mod handlers { } let current_context = sess.new_default_turn_with_sub_id(sub_id).await; + let has_model_visible_state_update = next_collaboration_mode.is_some(); let pending_model_visible_state_sync = { let mut state = sess.state.lock().await; - state.take_pending_model_visible_state_sync(next_collaboration_mode.is_some()) + state.take_pending_model_visible_state_sync(has_model_visible_state_update) }; let update_items = sess.build_settings_update_items( Some(&previous_context), @@ -2877,9 +2878,10 @@ mod handlers { // Attempt to inject input into current task if let Err(items) = sess.inject_input(items).await { sess.seed_initial_context_if_needed(¤t_context).await; + let has_model_visible_state_update = next_collaboration_mode.is_some(); let pending_model_visible_state_sync = { let mut state = sess.state.lock().await; - state.take_pending_model_visible_state_sync(next_collaboration_mode.is_some()) + state.take_pending_model_visible_state_sync(has_model_visible_state_update) }; let update_items = sess.build_settings_update_items( previous_context.as_ref(), diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 1d4725d2f3..0492d12999 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -117,15 +117,15 @@ impl SessionState { self.turn_context_history = turn_context_history; } - /// Consume pending model-visible state sync only when a collaboration update is being built. + /// Consume pending model-visible state sync only when a model-visible update is being built. /// - /// If the next operation does not carry a collaboration update (for example, model-only - /// overrides), keep the pending sync for a later turn that does. + /// If the next operation does not carry any model-visible update, keep the pending sync for a + /// later turn that does. pub(crate) fn take_pending_model_visible_state_sync( &mut self, - has_collaboration_update: bool, + has_model_visible_state_update: bool, ) -> PendingModelVisibleStateSync { - if has_collaboration_update { + if has_model_visible_state_update { std::mem::take(&mut self.pending_model_visible_state_sync) } else { PendingModelVisibleStateSync::None diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index a4862d6fa2..0d9deb9587 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -3128,6 +3128,18 @@ mod tests { .expect("backtrack selection"); app.apply_backtrack_rollback(selection); + assert_eq!( + app.chat_widget.active_collaboration_mode_kind(), + ModeKind::Code + ); + + app.handle_backtrack_event(&EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + model_visible_state: Some(ModelVisibleState { + collaboration_mode: Some(collaboration_mode(ModeKind::Plan)), + }), + })); + assert_eq!( app.chat_widget.active_collaboration_mode_kind(), ModeKind::Plan @@ -3208,6 +3220,18 @@ mod tests { .expect("backtrack selection"); app.apply_backtrack_rollback(selection); + assert_eq!( + app.chat_widget.active_collaboration_mode_kind(), + ModeKind::Code + ); + + app.handle_backtrack_event(&EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + model_visible_state: Some(ModelVisibleState { + collaboration_mode: Some(collaboration_mode(ModeKind::Plan)), + }), + })); + assert_eq!( app.chat_widget.active_collaboration_mode_kind(), ModeKind::Plan @@ -3373,25 +3397,25 @@ mod tests { app.apply_backtrack_rollback(selection); assert_eq!( app.chat_widget.active_collaboration_mode_kind(), - ModeKind::Plan + ModeKind::Code ); app.handle_backtrack_event(&EventMsg::ThreadRolledBack(ThreadRolledBackEvent { num_turns: 1, model_visible_state: Some(ModelVisibleState { - collaboration_mode: Some(collaboration_mode(ModeKind::Code)), + collaboration_mode: Some(collaboration_mode(ModeKind::Plan)), }), })); assert_eq!( app.chat_widget.active_collaboration_mode_kind(), - ModeKind::Code + ModeKind::Plan ); assert!(app.backtrack.pending_rollback.is_none()); } #[tokio::test] - async fn backtrack_failed_restores_previous_collaboration_mode() { + async fn backtrack_failed_keeps_current_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); @@ -3451,7 +3475,7 @@ mod tests { assert_eq!( app.chat_widget.active_collaboration_mode_kind(), - ModeKind::Plan + ModeKind::Code ); app.handle_backtrack_event(&EventMsg::Error(ErrorEvent { diff --git a/codex-rs/tui/src/app_backtrack.rs b/codex-rs/tui/src/app_backtrack.rs index c72628d298..c6a27f6159 100644 --- a/codex-rs/tui/src/app_backtrack.rs +++ b/codex-rs/tui/src/app_backtrack.rs @@ -39,7 +39,6 @@ use codex_core::protocol::EventMsg; use codex_core::protocol::Op; use codex_core::protocol::ThreadRolledBackEvent; use codex_protocol::ThreadId; -use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::user_input::TextElement; use color_eyre::eyre::Result; use crossterm::event::KeyCode; @@ -86,8 +85,6 @@ pub(crate) struct BacktrackSelection { pub(crate) text_elements: Vec, /// Local image paths associated with the selected user message. pub(crate) local_image_paths: Vec, - /// Collaboration mode active when the selected message was sent. - pub(crate) collaboration_mode: Option, } /// An in-flight rollback requested from core. @@ -98,8 +95,6 @@ pub(crate) struct BacktrackSelection { pub(crate) struct PendingBacktrackRollback { pub(crate) selection: BacktrackSelection, pub(crate) thread_id: Option, - /// Active mask before optimistic rollback preview mode changes. - pub(crate) previous_collaboration_mask: Option, } impl App { @@ -209,21 +204,12 @@ impl App { return; } - // Optimistically apply the target mode for rollback preview when available. - // If the selected history item has no mode data (legacy/replayed entries), preserve the - // current mode until core returns authoritative rollback state. - let previous_collaboration_mask = self.chat_widget.active_collaboration_mask(); - if let Some(collaboration_mode) = selection.collaboration_mode.clone() { - self.chat_widget - .replace_collaboration_mask(Some(collaboration_mode)); - } 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() { @@ -479,9 +465,6 @@ impl App { if pending.thread_id != self.chat_widget.thread_id() { return; } - // Rollback failed: revert the optimistic mode switch to the pre-rollback mask. - self.chat_widget - .replace_collaboration_mask(pending.previous_collaboration_mask); } _ => {} } @@ -503,9 +486,6 @@ impl App { if let Some(model_visible_state) = rollback.model_visible_state.as_ref() { self.chat_widget .apply_model_visible_state(model_visible_state); - } else { - self.chat_widget - .replace_collaboration_mask(pending.previous_collaboration_mask.clone()); } self.trim_transcript_for_backtrack(pending.selection.nth_user_message); @@ -518,26 +498,25 @@ impl App { return None; } - let (prefill, text_elements, local_image_paths, collaboration_mode) = + let (prefill, text_elements, local_image_paths) = nth_user_position(&self.transcript_cells, nth_user_message) .and_then(|idx| self.transcript_cells.get(idx)) .and_then(|cell| cell.as_any().downcast_ref::()) .map(|cell| { + let _historical_mode = cell.collaboration_mode.as_ref(); ( cell.message.clone(), cell.text_elements.clone(), cell.local_image_paths.clone(), - cell.collaboration_mode.clone(), ) }) - .unwrap_or_else(|| (String::new(), Vec::new(), Vec::new(), None)); + .unwrap_or_else(|| (String::new(), Vec::new(), Vec::new())); Some(BacktrackSelection { nth_user_message, prefill, text_elements, local_image_paths, - collaboration_mode, }) } diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 35fa4d4242..4d54267c34 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -5372,10 +5372,6 @@ impl ChatWidget { self.active_mode_kind() } - pub(crate) fn active_collaboration_mask(&self) -> Option { - self.active_collaboration_mask.clone() - } - fn is_session_configured(&self) -> bool { self.thread_id.is_some() } @@ -5520,21 +5516,6 @@ impl ChatWidget { self.request_redraw(); } - pub(crate) fn replace_collaboration_mask( - &mut self, - mask: Option, - ) -> Option { - 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 - } - /// Synchronize collaboration-mode UI state from core's model-visible snapshot. /// /// This is primarily used after rollback/backtrack so the composer mode and model display