fix(tui) Defer backtrack trim until rollback confirms (#9401)

Document the backtrack/rollback state machine and invariants between the
transcript overlay, in-flight “live tail”, and core thread state (tui + tui2).

Also adjust behavior for correctness:
- Track a single pending rollback and block additional rollbacks until core responds.
- Defer trimming transcript cells until ThreadRolledBack for the active session.
- Clear the guard on ThreadRollbackFailed so the user can retry.
- After a confirmed trim, schedule a one-shot scrollback refresh on the next draw.
- Clear stale pending rollback state when switching sessions.

---------

Co-authored-by: Josh McKinney <joshka@openai.com>
This commit is contained in:
Ahmed Ibrahim
2026-01-16 22:29:41 -08:00
committed by GitHub
parent 93a5e0fe1c
commit 764f3c7d03
4 changed files with 237 additions and 23 deletions

View File

@@ -3,6 +3,18 @@
//! This file owns backtrack mode (Esc/Enter navigation in the transcript overlay) and also
//! mediates a key rendering boundary for the transcript overlay.
//!
//! Overall goal: keep the main chat view and the transcript overlay in sync while allowing
//! users to "rewind" to an earlier user message. We stage a rollback request, wait for core to
//! confirm it, then trim the local transcript to the matching history boundary. This avoids UI
//! state diverging from the agent if a rollback fails or targets a different thread.
//!
//! Backtrack operates as a small state machine:
//! - The first `Esc` in the main view "primes" the feature and captures a base thread id.
//! - A subsequent `Esc` opens the transcript overlay (`Ctrl+T`) and highlights a user message.
//! - `Enter` requests a rollback from core and records a `pending_rollback` guard.
//! - Only after receiving `EventMsg::ThreadRolledBack` do we trim local transcript state and
//! schedule a one-time scrollback refresh.
//!
//! The transcript overlay (`Ctrl+T`) renders committed transcript cells plus a render-only live
//! tail derived from the current in-flight `ChatWidget.active_cell`.
//!
@@ -20,6 +32,9 @@ use crate::history_cell::UserHistoryCell;
use crate::pager_overlay::Overlay;
use crate::tui;
use crate::tui::TuiEvent;
use codex_core::protocol::CodexErrorInfo;
use codex_core::protocol::ErrorEvent;
use codex_core::protocol::EventMsg;
use codex_core::protocol::Op;
use codex_protocol::ThreadId;
use color_eyre::eyre::Result;
@@ -33,24 +48,53 @@ pub(crate) struct BacktrackState {
/// True when Esc has primed backtrack mode in the main view.
pub(crate) primed: bool,
/// Session id of the base thread to rollback.
///
/// If the current thread changes, backtrack selections become invalid and must be ignored.
pub(crate) base_id: Option<ThreadId>,
/// Index in the transcript of the last user message.
/// Index of the currently highlighted user message.
///
/// This is an index into the filtered "user messages since the last session start" view,
/// not an index into `transcript_cells`. `usize::MAX` indicates "no selection".
pub(crate) nth_user_message: usize,
/// True when the transcript overlay is showing a backtrack preview.
pub(crate) overlay_preview_active: bool,
/// Pending rollback request awaiting confirmation from core.
///
/// This acts as a guardrail: once we request a rollback, we block additional backtrack
/// submissions until core responds with either a success or failure event.
pub(crate) pending_rollback: Option<PendingBacktrackRollback>,
}
/// A user-visible backtrack choice that can be confirmed into a rollback request.
#[derive(Debug, Clone)]
pub(crate) struct BacktrackSelection {
/// The selected user message, counted from the most recent session start.
///
/// This value is used both to compute the rollback depth and to trim the local transcript
/// after core confirms the rollback.
pub(crate) nth_user_message: usize,
/// Composer prefill derived from the selected user message.
///
/// This is applied immediately on selection confirmation; if the rollback fails, the prefill
/// remains as a convenience so the user can retry or edit.
pub(crate) prefill: String,
}
/// An in-flight rollback requested from core.
///
/// We keep enough information to apply the corresponding local trim only if the response targets
/// the same active thread we issued the request for.
#[derive(Debug, Clone)]
pub(crate) struct PendingBacktrackRollback {
pub(crate) selection: BacktrackSelection,
pub(crate) thread_id: Option<ThreadId>,
}
impl App {
/// Route overlay events when transcript overlay is active.
/// - If backtrack preview is active: Esc steps selection; Enter confirms.
/// - Otherwise: Esc begins preview; all other events forward to overlay.
/// interactions (Esc to step target, Enter to confirm) and overlay lifecycle.
/// Route overlay events while the transcript overlay is active.
///
/// If backtrack preview is active, Esc steps the selection and Enter confirms it.
/// Otherwise, Esc begins preview mode and all other events are forwarded to the overlay.
pub(crate) async fn handle_backtrack_overlay_event(
&mut self,
tui: &mut tui::Tui,
@@ -112,22 +156,38 @@ impl App {
}
/// Stage a backtrack and request thread history from the agent.
///
/// We send the rollback request immediately, but we only mutate the transcript after core
/// confirms success so the UI cannot get ahead of the actual thread state.
///
/// The composer prefill is applied immediately as a UX convenience; it does not imply that
/// core has accepted the rollback.
pub(crate) fn apply_backtrack_rollback(&mut self, selection: BacktrackSelection) {
let user_total = user_count(&self.transcript_cells);
if user_total == 0 {
return;
}
if self.backtrack.pending_rollback.is_some() {
self.chat_widget
.add_error_message("Backtrack rollback already in progress.".to_string());
return;
}
let num_turns = user_total.saturating_sub(selection.nth_user_message);
let num_turns = u32::try_from(num_turns).unwrap_or(u32::MAX);
if num_turns == 0 {
return;
}
let prefill = selection.prefill.clone();
self.backtrack.pending_rollback = Some(PendingBacktrackRollback {
selection,
thread_id: self.chat_widget.thread_id(),
});
self.chat_widget.submit_op(Op::ThreadRollback { num_turns });
self.trim_transcript_for_backtrack(selection.nth_user_message);
if !selection.prefill.is_empty() {
self.chat_widget.set_composer_text(selection.prefill);
if !prefill.is_empty() {
self.chat_widget.set_composer_text(prefill);
}
}
@@ -290,7 +350,6 @@ impl App {
self.close_transcript_overlay(tui);
if let Some(selection) = selection {
self.apply_backtrack_rollback(selection);
self.render_transcript_once(tui);
tui.frame_requester().schedule_frame();
}
}
@@ -328,10 +387,39 @@ impl App {
selection: BacktrackSelection,
) {
self.apply_backtrack_rollback(selection);
self.render_transcript_once(tui);
tui.frame_requester().schedule_frame();
}
pub(crate) fn handle_backtrack_event(&mut self, event: &EventMsg) {
match event {
EventMsg::ThreadRolledBack(_) => self.finish_pending_backtrack(),
EventMsg::Error(ErrorEvent {
codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed),
..
}) => {
// Core rejected the rollback; clear the guard so the user can retry.
self.backtrack.pending_rollback = None;
}
_ => {}
}
}
/// Finish a pending rollback by applying the local trim and scheduling a scrollback refresh.
///
/// We ignore events that do not correspond to the currently active thread to avoid applying
/// stale updates after a session switch.
fn finish_pending_backtrack(&mut self) {
let Some(pending) = self.backtrack.pending_rollback.take() else {
return;
};
if pending.thread_id != self.chat_widget.thread_id() {
// Ignore rollbacks targeting a prior thread.
return;
}
self.trim_transcript_for_backtrack(pending.selection.nth_user_message);
self.backtrack_render_pending = true;
}
fn backtrack_selection(&self, nth_user_message: usize) -> Option<BacktrackSelection> {
let base_id = self.backtrack.base_id?;
if self.chat_widget.thread_id() != Some(base_id) {
@@ -350,7 +438,7 @@ impl App {
})
}
/// Trim transcript_cells to preserve only content up to the selected user message.
/// Trim `transcript_cells` to preserve only content before the selected user message.
fn trim_transcript_for_backtrack(&mut self, nth_user_message: usize) {
trim_transcript_cells_to_nth_user(&mut self.transcript_cells, nth_user_message);
}