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() }