From d83b98516d3653191eaea5a849043e048e568450 Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Fri, 13 Mar 2026 19:55:44 -0700 Subject: [PATCH] Restore TUI turn context after backtracking Snapshot per-turn TUI context during committed history so rollback can restore the surviving turn's cwd, approval and sandbox settings, collaboration mode, service tier, personality, and related local state. Add regression coverage for pending and replayed rollback flows. Co-authored-by: Codex --- codex-rs/tui/src/app.rs | 284 +++++++++++++++++++++++++++++- codex-rs/tui/src/app_backtrack.rs | 124 +++++++++++++ codex-rs/tui/src/chatwidget.rs | 37 ++++ 3 files changed, 444 insertions(+), 1 deletion(-) diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index ca27a2ad1e..6737bdbc22 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -2577,6 +2577,7 @@ impl App { tui.frame_requester().schedule_frame(); } self.transcript_cells.push(cell.clone()); + self.track_backtrack_transcript_cell(&cell); let mut display = cell.display_lines(tui.terminal.last_known_screen_size.width); if !display.is_empty() { // Only insert a separating blank line for new cells that are not @@ -4134,7 +4135,16 @@ impl App { }; } - fn refresh_status_line(&mut self) { + pub(crate) fn set_runtime_policy_overrides( + &mut self, + approval_policy: AskForApproval, + sandbox_policy: SandboxPolicy, + ) { + self.runtime_approval_policy_override = Some(approval_policy); + self.runtime_sandbox_policy_override = Some(sandbox_policy); + } + + pub(crate) fn refresh_status_line(&mut self) { self.chat_widget.refresh_status_line(); } @@ -4172,6 +4182,7 @@ mod tests { use super::*; use crate::app_backtrack::BacktrackSelection; use crate::app_backtrack::BacktrackState; + use crate::app_backtrack::BacktrackTurnContextSnapshot; use crate::app_backtrack::user_count; use crate::chatwidget::tests::make_chatwidget_manual_with_sender; use crate::chatwidget::tests::set_chatgpt_auth; @@ -4185,14 +4196,18 @@ mod tests { use codex_core::CodexAuth; use codex_core::config::ConfigBuilder; use codex_core::config::ConfigOverrides; + use codex_core::config::types::ApprovalsReviewer; use codex_core::config::types::ModelAvailabilityNuxConfig; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::config_types::ModeKind; + use codex_protocol::config_types::Personality; + use codex_protocol::config_types::ServiceTier; use codex_protocol::config_types::Settings; use codex_protocol::openai_models::ModelAvailabilityNux; + use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::AgentMessageDeltaEvent; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::Event; @@ -6442,6 +6457,58 @@ smart_approvals = true ) } + fn backtrack_test_session_configured( + session_id: ThreadId, + cwd: PathBuf, + ) -> SessionConfiguredEvent { + SessionConfiguredEvent { + session_id, + forked_from_id: None, + thread_name: None, + model: "gpt-test".to_string(), + model_provider_id: "test-provider".to_string(), + service_tier: None, + approval_policy: AskForApproval::Never, + approvals_reviewer: ApprovalsReviewer::User, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + cwd, + reasoning_effort: None, + history_log_id: 0, + history_entry_count: 0, + initial_messages: None, + network_proxy: None, + rollout_path: Some(PathBuf::new()), + } + } + + fn backtrack_user_cell(message: &str) -> Arc { + Arc::new(UserHistoryCell { + message: message.to_string(), + text_elements: Vec::new(), + local_image_paths: Vec::new(), + remote_image_urls: Vec::new(), + }) + } + + fn assert_backtrack_turn_context(app: &App, expected: &BacktrackTurnContextSnapshot) { + assert_eq!( + app.capture_backtrack_turn_context_snapshot(), + expected.clone() + ); + assert_eq!(app.config.cwd, expected.cwd.clone()); + assert_eq!(app.config.approvals_reviewer, expected.approvals_reviewer); + assert_eq!(app.config.service_tier, expected.service_tier); + assert_eq!(app.config.personality, expected.personality); + assert_eq!( + app.runtime_approval_policy_override, + Some(expected.approval_policy) + ); + assert_eq!( + app.runtime_sandbox_policy_override, + Some(expected.sandbox_policy.clone()) + ); + } + fn next_user_turn_op(op_rx: &mut tokio::sync::mpsc::UnboundedReceiver) -> Op { let mut seen = Vec::new(); while let Ok(op) = op_rx.try_recv() { @@ -7356,6 +7423,221 @@ smart_approvals = true assert_eq!(user_messages, vec!["first prompt".to_string()]); } + #[tokio::test] + async fn pending_backtrack_restores_previous_surviving_turn_context() { + let mut app = make_test_app().await; + let session_id = ThreadId::new(); + let session_event = + backtrack_test_session_configured(session_id, PathBuf::from("/tmp/backtrack-base")); + app.chat_widget.handle_codex_event(Event { + id: "session-configured".to_string(), + msg: EventMsg::SessionConfigured(session_event.clone()), + }); + app.config = app.chat_widget.config_ref().clone(); + + let session_cell = Arc::new(new_session_info( + app.chat_widget.config_ref(), + "gpt-test", + session_event, + false, + None, + None, + false, + )) as Arc; + app.track_backtrack_transcript_cell(&session_cell); + app.transcript_cells.push(session_cell); + let baseline_snapshot = app.capture_backtrack_turn_context_snapshot(); + + let turn_one_snapshot = BacktrackTurnContextSnapshot { + cwd: PathBuf::from("/tmp/backtrack-turn-one"), + approval_policy: AskForApproval::OnRequest, + approvals_reviewer: ApprovalsReviewer::GuardianSubagent, + sandbox_policy: SandboxPolicy::new_workspace_write_policy(), + current_collaboration_mode: CollaborationMode { + mode: ModeKind::Default, + settings: Settings { + model: "gpt-turn-one".to_string(), + reasoning_effort: Some(ReasoningEffortConfig::Low), + developer_instructions: None, + }, + }, + active_collaboration_mask: Some(CollaborationModeMask { + name: "plan".to_string(), + mode: Some(ModeKind::Plan), + model: Some("gpt-turn-one".to_string()), + reasoning_effort: Some(Some(ReasoningEffortConfig::High)), + developer_instructions: None, + }), + service_tier: Some(ServiceTier::Fast), + personality: Some(Personality::Pragmatic), + }; + app.apply_backtrack_turn_context_snapshot(&turn_one_snapshot); + let user_one = backtrack_user_cell("first prompt"); + app.track_backtrack_transcript_cell(&user_one); + app.transcript_cells.push(user_one); + + let turn_two_snapshot = BacktrackTurnContextSnapshot { + cwd: PathBuf::from("/tmp/backtrack-turn-two"), + approval_policy: AskForApproval::Never, + approvals_reviewer: ApprovalsReviewer::User, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + current_collaboration_mode: CollaborationMode { + mode: ModeKind::Default, + settings: Settings { + model: "gpt-turn-two".to_string(), + reasoning_effort: Some(ReasoningEffortConfig::Medium), + developer_instructions: None, + }, + }, + active_collaboration_mask: Some(CollaborationModeMask { + name: "default".to_string(), + mode: Some(ModeKind::Default), + model: Some("gpt-turn-two".to_string()), + reasoning_effort: Some(Some(ReasoningEffortConfig::Medium)), + developer_instructions: None, + }), + service_tier: None, + personality: Some(Personality::Friendly), + }; + app.apply_backtrack_turn_context_snapshot(&turn_two_snapshot); + let user_two = backtrack_user_cell("second prompt"); + app.track_backtrack_transcript_cell(&user_two); + app.transcript_cells.push(user_two); + + app.backtrack.pending_rollback = Some(crate::app_backtrack::PendingBacktrackRollback { + selection: crate::app_backtrack::BacktrackSelection { + nth_user_message: 1, + prefill: String::new(), + text_elements: Vec::new(), + local_image_paths: Vec::new(), + remote_image_urls: Vec::new(), + }, + thread_id: app.chat_widget.thread_id(), + }); + + app.handle_backtrack_event(&EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + })); + + let user_messages: Vec = app + .transcript_cells + .iter() + .filter_map(|cell| { + cell.as_any() + .downcast_ref::() + .map(|cell| cell.message.clone()) + }) + .collect(); + assert_eq!(user_messages, vec!["first prompt".to_string()]); + assert_backtrack_turn_context(&app, &turn_one_snapshot); + assert_eq!( + app.backtrack.session_baseline_turn_context, + Some(baseline_snapshot) + ); + assert_eq!( + app.backtrack.current_session_turn_contexts, + vec![turn_one_snapshot] + ); + } + + #[tokio::test] + async fn non_pending_backtrack_restores_session_baseline_turn_context() { + let mut app = make_test_app().await; + let session_id = ThreadId::new(); + let session_event = + backtrack_test_session_configured(session_id, PathBuf::from("/tmp/backtrack-base")); + app.chat_widget.handle_codex_event(Event { + id: "session-configured".to_string(), + msg: EventMsg::SessionConfigured(session_event.clone()), + }); + app.config = app.chat_widget.config_ref().clone(); + + let session_cell = Arc::new(new_session_info( + app.chat_widget.config_ref(), + "gpt-test", + session_event, + false, + None, + None, + false, + )) as Arc; + app.track_backtrack_transcript_cell(&session_cell); + app.transcript_cells.push(session_cell); + let baseline_snapshot = app.capture_backtrack_turn_context_snapshot(); + + let turn_snapshot = BacktrackTurnContextSnapshot { + cwd: PathBuf::from("/tmp/backtrack-turn"), + approval_policy: AskForApproval::OnRequest, + approvals_reviewer: ApprovalsReviewer::GuardianSubagent, + sandbox_policy: SandboxPolicy::new_workspace_write_policy(), + current_collaboration_mode: CollaborationMode { + mode: ModeKind::Default, + settings: Settings { + model: "gpt-turn".to_string(), + reasoning_effort: Some(ReasoningEffortConfig::Low), + developer_instructions: None, + }, + }, + active_collaboration_mask: Some(CollaborationModeMask { + name: "plan".to_string(), + mode: Some(ModeKind::Plan), + model: Some("gpt-turn".to_string()), + reasoning_effort: Some(Some(ReasoningEffortConfig::High)), + developer_instructions: None, + }), + service_tier: Some(ServiceTier::Fast), + personality: Some(Personality::Pragmatic), + }; + app.apply_backtrack_turn_context_snapshot(&turn_snapshot); + let user_cell = backtrack_user_cell("first prompt"); + app.track_backtrack_transcript_cell(&user_cell); + app.transcript_cells.push(user_cell); + + let stale_snapshot = BacktrackTurnContextSnapshot { + cwd: PathBuf::from("/tmp/backtrack-stale"), + approval_policy: AskForApproval::Never, + approvals_reviewer: ApprovalsReviewer::User, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + current_collaboration_mode: CollaborationMode { + mode: ModeKind::Default, + settings: Settings { + model: "gpt-stale".to_string(), + reasoning_effort: Some(ReasoningEffortConfig::Medium), + developer_instructions: None, + }, + }, + active_collaboration_mask: Some(CollaborationModeMask { + name: "default".to_string(), + mode: Some(ModeKind::Default), + model: Some("gpt-stale".to_string()), + reasoning_effort: Some(Some(ReasoningEffortConfig::Medium)), + developer_instructions: None, + }), + service_tier: None, + personality: Some(Personality::Friendly), + }; + app.apply_backtrack_turn_context_snapshot(&stale_snapshot); + + assert!(app.apply_non_pending_thread_rollback(1)); + + let user_messages: Vec = app + .transcript_cells + .iter() + .filter_map(|cell| { + cell.as_any() + .downcast_ref::() + .map(|cell| cell.message.clone()) + }) + .collect(); + assert_eq!(user_messages, Vec::::new()); + assert_backtrack_turn_context(&app, &baseline_snapshot); + assert_eq!( + app.backtrack.session_baseline_turn_context, + Some(baseline_snapshot) + ); + assert_eq!(app.backtrack.current_session_turn_contexts, Vec::new()); + } + #[tokio::test] async fn queued_rollback_syncs_overlay_and_clears_deferred_history() { let mut app = make_test_app().await; diff --git a/codex-rs/tui/src/app_backtrack.rs b/codex-rs/tui/src/app_backtrack.rs index 1bb2c29569..8c61407b6c 100644 --- a/codex-rs/tui/src/app_backtrack.rs +++ b/codex-rs/tui/src/app_backtrack.rs @@ -34,11 +34,18 @@ use crate::history_cell::UserHistoryCell; use crate::pager_overlay::Overlay; use crate::tui; use crate::tui::TuiEvent; +use codex_core::config::types::ApprovalsReviewer; use codex_protocol::ThreadId; +use codex_protocol::config_types::CollaborationMode; +use codex_protocol::config_types::CollaborationModeMask; +use codex_protocol::config_types::Personality; +use codex_protocol::config_types::ServiceTier; +use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::SandboxPolicy; use codex_protocol::user_input::TextElement; use color_eyre::eyre::Result; use crossterm::event::KeyCode; @@ -66,6 +73,22 @@ pub(crate) struct BacktrackState { /// 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, + /// Session baseline turn context from the latest session start. + pub(crate) session_baseline_turn_context: Option, + /// Turn-context snapshots aligned with committed user turns in the current session. + pub(crate) current_session_turn_contexts: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct BacktrackTurnContextSnapshot { + pub(crate) cwd: PathBuf, + pub(crate) approval_policy: AskForApproval, + pub(crate) approvals_reviewer: ApprovalsReviewer, + pub(crate) sandbox_policy: SandboxPolicy, + pub(crate) current_collaboration_mode: CollaborationMode, + pub(crate) active_collaboration_mask: Option, + pub(crate) service_tier: Option, + pub(crate) personality: Option, } /// A user-visible backtrack choice that can be confirmed into a rollback request. @@ -461,6 +484,94 @@ impl App { tui.frame_requester().schedule_frame(); } + pub(crate) fn capture_backtrack_turn_context_snapshot(&self) -> BacktrackTurnContextSnapshot { + BacktrackTurnContextSnapshot { + cwd: self.chat_widget.config_ref().cwd.clone(), + approval_policy: self + .chat_widget + .config_ref() + .permissions + .approval_policy + .value(), + approvals_reviewer: self.chat_widget.config_ref().approvals_reviewer, + sandbox_policy: self + .chat_widget + .config_ref() + .permissions + .sandbox_policy + .get() + .clone(), + current_collaboration_mode: self.chat_widget.current_collaboration_mode().clone(), + active_collaboration_mask: self.chat_widget.active_collaboration_mask().cloned(), + service_tier: self.chat_widget.current_service_tier(), + personality: self.chat_widget.config_ref().personality, + } + } + + pub(crate) fn track_backtrack_transcript_cell( + &mut self, + cell: &Arc, + ) { + if cell.as_any().is::() { + self.backtrack.session_baseline_turn_context = + Some(self.capture_backtrack_turn_context_snapshot()); + self.backtrack.current_session_turn_contexts.clear(); + return; + } + if cell.as_any().is::() { + self.backtrack + .current_session_turn_contexts + .push(self.capture_backtrack_turn_context_snapshot()); + } + } + + pub(crate) fn apply_backtrack_turn_context_snapshot( + &mut self, + snapshot: &BacktrackTurnContextSnapshot, + ) { + self.config.cwd = snapshot.cwd.clone(); + self.config.approvals_reviewer = snapshot.approvals_reviewer; + self.config.service_tier = snapshot.service_tier; + self.config.personality = snapshot.personality; + self.config.model_reasoning_effort = snapshot.current_collaboration_mode.reasoning_effort(); + if let Err(err) = self + .config + .permissions + .approval_policy + .set(snapshot.approval_policy) + { + tracing::warn!(%err, "failed to restore approval_policy after rollback"); + } + if let Err(err) = self + .config + .permissions + .sandbox_policy + .set(snapshot.sandbox_policy.clone()) + { + tracing::warn!(%err, "failed to restore sandbox_policy after rollback"); + } + self.set_runtime_policy_overrides( + snapshot.approval_policy, + snapshot.sandbox_policy.clone(), + ); + self.file_search.update_search_dir(self.config.cwd.clone()); + self.chat_widget.restore_backtrack_turn_context(snapshot); + self.refresh_status_line(); + } + + fn restore_backtrack_turn_context_after_trim(&mut self) { + let Some(snapshot) = self + .backtrack + .current_session_turn_contexts + .last() + .cloned() + .or_else(|| self.backtrack.session_baseline_turn_context.clone()) + else { + return; + }; + self.apply_backtrack_turn_context_snapshot(&snapshot); + } + pub(crate) fn handle_backtrack_event(&mut self, event: &EventMsg) { match event { EventMsg::ThreadRolledBack(rollback) => { @@ -500,6 +611,15 @@ impl App { if !trim_transcript_cells_drop_last_n_user_turns(&mut self.transcript_cells, num_turns) { return false; } + let turns_from_end = usize::try_from(num_turns).unwrap_or(usize::MAX); + if turns_from_end >= self.backtrack.current_session_turn_contexts.len() { + self.backtrack.current_session_turn_contexts.clear(); + } else { + self.backtrack + .current_session_turn_contexts + .truncate(self.backtrack.current_session_turn_contexts.len() - turns_from_end); + } + self.restore_backtrack_turn_context_after_trim(); self.sync_overlay_after_transcript_trim(); self.backtrack_render_pending = true; true @@ -521,6 +641,10 @@ impl App { &mut self.transcript_cells, pending.selection.nth_user_message, ) { + self.backtrack + .current_session_turn_contexts + .truncate(pending.selection.nth_user_message); + self.restore_backtrack_turn_context_after_trim(); self.sync_overlay_after_transcript_trim(); self.backtrack_render_pending = true; } diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 4c8e76874c..cc5bb2a507 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -38,6 +38,7 @@ use std::time::Duration; use std::time::Instant; use self::realtime::PendingSteerCompareKey; +use crate::app_backtrack::BacktrackTurnContextSnapshot; use crate::app_event::RealtimeAudioDeviceKind; #[cfg(all(not(target_os = "linux"), feature = "voice-input"))] use crate::audio_device::list_realtime_audio_device_names; @@ -7869,6 +7870,38 @@ impl ChatWidget { self.config.approvals_reviewer = policy; } + pub(crate) fn restore_backtrack_turn_context( + &mut self, + snapshot: &BacktrackTurnContextSnapshot, + ) { + self.config.cwd = snapshot.cwd.clone(); + self.current_cwd = Some(snapshot.cwd.clone()); + if let Err(err) = self + .config + .permissions + .approval_policy + .set(snapshot.approval_policy) + { + tracing::warn!(%err, "failed to restore approval_policy on chat config"); + } + if let Err(err) = self + .config + .permissions + .sandbox_policy + .set(snapshot.sandbox_policy.clone()) + { + tracing::warn!(%err, "failed to restore sandbox_policy on chat config"); + } + self.config.approvals_reviewer = snapshot.approvals_reviewer; + self.config.service_tier = snapshot.service_tier; + self.config.personality = snapshot.personality; + self.current_collaboration_mode = snapshot.current_collaboration_mode.clone(); + self.active_collaboration_mask = snapshot.active_collaboration_mask.clone(); + self.update_collaboration_mode_indicator(); + self.refresh_model_display(); + self.request_redraw(); + } + pub(crate) fn set_full_access_warning_acknowledged(&mut self, acknowledged: bool) { self.config.notices.hide_full_access_warning = Some(acknowledged); } @@ -8089,6 +8122,10 @@ impl ChatWidget { &self.current_collaboration_mode } + pub(crate) fn active_collaboration_mask(&self) -> Option<&CollaborationModeMask> { + self.active_collaboration_mask.as_ref() + } + pub(crate) fn current_reasoning_effort(&self) -> Option { self.effective_reasoning_effort() }