Simplify backtrack mode sync and generalize pending sync naming

This commit is contained in:
Charles Cunningham
2026-02-02 22:29:40 -08:00
parent 4b7b5cd7cc
commit 0174a3db48
5 changed files with 41 additions and 55 deletions
+4 -2
View File
@@ -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(&current_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(),
+5 -5
View File
@@ -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
+29 -5
View File
@@ -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 {
+3 -24
View File
@@ -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<TextElement>,
/// Local image paths associated with the selected user message.
pub(crate) local_image_paths: Vec<PathBuf>,
/// Collaboration mode active when the selected message was sent.
pub(crate) collaboration_mode: Option<CollaborationModeMask>,
}
/// 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<ThreadId>,
/// Active mask before optimistic rollback preview mode changes.
pub(crate) previous_collaboration_mask: Option<CollaborationModeMask>,
}
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::<UserHistoryCell>())
.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,
})
}
-19
View File
@@ -5372,10 +5372,6 @@ impl ChatWidget {
self.active_mode_kind()
}
pub(crate) fn active_collaboration_mask(&self) -> Option<CollaborationModeMask> {
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<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
}
/// 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