From d46fb69e87541c3a20569e76e2206d45ec5620f4 Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Thu, 9 Jul 2026 19:25:39 -0300 Subject: [PATCH 1/2] fix(tui): retry safety-buffered turns on forks --- codex-rs/tui/src/app.rs | 1 + codex-rs/tui/src/app/event_dispatch.rs | 77 ++------ codex-rs/tui/src/app/safety_buffering.rs | 183 ++++++++++++++++++ .../tui/src/app/safety_buffering_tests.rs | 67 +++++++ codex-rs/tui/src/app/thread_routing.rs | 29 +-- codex-rs/tui/src/app_event.rs | 3 +- .../tui/src/chatwidget/input_submission.rs | 6 +- .../tui/src/chatwidget/safety_buffering.rs | 74 ++++--- .../tui/src/chatwidget/tests/app_server.rs | 75 ++++++- 9 files changed, 394 insertions(+), 121 deletions(-) create mode 100644 codex-rs/tui/src/app/safety_buffering.rs create mode 100644 codex-rs/tui/src/app/safety_buffering_tests.rs diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 554df264e2b1..56a13dd70338 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -214,6 +214,7 @@ mod platform_actions; mod plugin_mentions; mod replay_filter; mod resize_reflow; +mod safety_buffering; mod session_lifecycle; mod side; mod startup_prompts; diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index d994fb1ce7b4..5583d5ac85ab 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -401,74 +401,21 @@ impl App { thread_id, turn_id, model, - mut turn, + turn, + prompt, } => { - if self.active_thread_id != Some(thread_id) - || self.chat_widget.thread_id() != Some(thread_id) - { - return Ok(AppRunControl::Continue); - } - if !self.chat_widget.can_retry_safety_buffered_turn(&turn_id) { - self.app_event_tx.send(AppEvent::UpdateModel(model)); - self.app_event_tx.send(AppEvent::UpdateReasoningEffort(Some( - ReasoningEffortConfig::Low, - ))); - return Ok(AppRunControl::Continue); - } - - let AppCommand::UserTurn { - model: turn_model, - effort, - collaboration_mode, - .. - } = &mut turn - else { - self.chat_widget.add_error_message( - "Failed to retry with a faster model: original turn is unavailable." - .to_string(), - ); - return Ok(AppRunControl::Continue); - }; - *turn_model = model.clone(); - *effort = Some(ReasoningEffortConfig::Low); - *collaboration_mode = collaboration_mode.as_ref().map(|mode| { - mode.with_updates( - Some(model), - Some(Some(ReasoningEffortConfig::Low)), - /*developer_instructions*/ None, - ) - }); - - if let Err(err) = app_server.turn_interrupt(thread_id, turn_id).await { - self.chat_widget - .add_error_message(format!("Failed to retry with a faster model: {err}")); - return Ok(AppRunControl::Continue); - } - let rollback_response = - match app_server.thread_rollback(thread_id, /*num_turns*/ 1).await { - Ok(response) => response, - Err(err) => { - self.chat_widget.add_error_message(format!( - "Failed to retry with a faster model: {err}" - )); - return Ok(AppRunControl::Continue); - } - }; - - self.chat_widget.prepare_safety_buffering_retry(); - self.handle_thread_rollback_response_with_origin( - thread_id, - /*num_turns*/ 1, - &rollback_response, - super::thread_routing::ThreadRollbackOrigin::SafetyBufferingRetry, + self.retry_safety_buffered_turn( + tui, + app_server, + super::safety_buffering::SafetyBufferedRetry { + thread_id, + turn_id, + model, + turn, + prompt, + }, ) .await; - - if let Err(err) = self.submit_thread_op(app_server, thread_id, turn).await { - self.chat_widget.fail_safety_buffering_retry(); - self.chat_widget - .add_error_message(format!("Failed to retry with a faster model: {err}")); - } } AppEvent::RestoreCancelledTurn(prompt) => { self.apply_cancelled_turn_edit(prompt); diff --git a/codex-rs/tui/src/app/safety_buffering.rs b/codex-rs/tui/src/app/safety_buffering.rs new file mode 100644 index 000000000000..1c10d97bd142 --- /dev/null +++ b/codex-rs/tui/src/app/safety_buffering.rs @@ -0,0 +1,183 @@ +//! Safety-buffered turn retries that preserve the source thread. + +use super::session_lifecycle::ThreadAttachPresentation; +use super::*; +use crate::chatwidget::UserMessage; + +pub(super) struct SafetyBufferedRetry { + pub(super) thread_id: ThreadId, + pub(super) turn_id: String, + pub(super) model: String, + pub(super) turn: AppCommand, + pub(super) prompt: UserMessage, +} + +impl App { + pub(super) async fn retry_safety_buffered_turn( + &mut self, + tui: &mut tui::Tui, + app_server: &mut AppServerSession, + retry: SafetyBufferedRetry, + ) { + let SafetyBufferedRetry { + thread_id, + turn_id, + model, + mut turn, + prompt, + } = retry; + if self.active_thread_id != Some(thread_id) + || self.chat_widget.thread_id() != Some(thread_id) + { + return; + } + if !self.chat_widget.can_retry_safety_buffered_turn(&turn_id) { + self.app_event_tx.send(AppEvent::UpdateModel(model)); + self.app_event_tx.send(AppEvent::UpdateReasoningEffort(Some( + ReasoningEffortConfig::Low, + ))); + return; + } + + let AppCommand::UserTurn { + model: turn_model, + effort, + collaboration_mode, + .. + } = &mut turn + else { + self.chat_widget.add_error_message( + "Failed to retry with a faster model: original turn is unavailable.".to_string(), + ); + return; + }; + *turn_model = model.clone(); + *effort = Some(ReasoningEffortConfig::Low); + *collaboration_mode = collaboration_mode.as_ref().map(|mode| { + mode.with_updates( + Some(model), + Some(Some(ReasoningEffortConfig::Low)), + /*developer_instructions*/ None, + ) + }); + + if let Err(err) = app_server.turn_interrupt(thread_id, turn_id.clone()).await { + self.chat_widget + .add_error_message(format!("Failed to retry with a faster model: {err}")); + return; + } + + let thread = match app_server + .thread_read(thread_id, /*include_turns*/ true) + .await + { + Ok(thread) => thread, + Err(err) => { + self.fail_safety_buffered_branch(prompt, err); + return; + } + }; + let fork_point = match safety_retry_fork_point(&thread.turns, &turn_id) { + Ok(fork_point) => fork_point, + Err(err) => { + self.fail_safety_buffered_branch(prompt, err); + return; + } + }; + + let started = if let Some(last_turn_id) = fork_point { + app_server + .fork_thread_after(self.config.clone(), thread_id, last_turn_id) + .await + } else { + app_server + .start_thread_with_session_start_source( + &self.config, + /*session_start_source*/ None, + ) + .await + }; + let started = match started { + Ok(started) => started, + Err(err) => { + self.fail_safety_buffered_branch(prompt, err); + return; + } + }; + let retry_thread_id = started.session.thread_id; + + self.shutdown_current_thread(app_server).await; + if let Err(err) = self + .replace_chat_widget_with_app_server_thread( + tui, + app_server, + started, + ThreadAttachPresentation::SessionLineage, + /*initial_user_message*/ None, + ) + .await + { + self.fail_safety_buffered_branch(prompt, err); + return; + } + + self.chat_widget + .prepare_safety_buffered_retry_submission(prompt.clone()); + if let Err(err) = self + .submit_thread_op(app_server, retry_thread_id, turn) + .await + { + self.fail_safety_buffered_branch(prompt, err); + return; + } + self.chat_widget + .commit_safety_buffered_retry_submission(prompt); + } + + fn fail_safety_buffered_branch(&mut self, prompt: UserMessage, err: impl std::fmt::Display) { + self.chat_widget.cancel_safety_buffered_retry_submission(); + self.chat_widget.restore_user_message_to_composer(prompt); + self.chat_widget + .add_error_message(format!("Failed to retry with a faster model: {err}")); + } +} + +fn safety_retry_fork_point(turns: &[Turn], turn_id: &str) -> Result> { + let Some(turn_index) = turns.iter().position(|turn| turn.id == turn_id) else { + return Err(color_eyre::eyre::eyre!( + "interrupted turn {turn_id} is missing from the source thread" + )); + }; + if turn_index + 1 != turns.len() { + return Err(color_eyre::eyre::eyre!( + "interrupted turn {turn_id} is no longer the latest turn" + )); + } + if turns[turn_index].status == TurnStatus::InProgress { + return Err(color_eyre::eyre::eyre!( + "interrupted turn {turn_id} is still in progress" + )); + } + + let Some(previous_turn) = turns[..turn_index].last() else { + return Ok(None); + }; + if !previous_turn.is_forkable { + return Err(color_eyre::eyre::eyre!( + "previous turn {} has no canonical fork boundary", + previous_turn.id + )); + } + if previous_turn.status == TurnStatus::InProgress { + return Err(color_eyre::eyre::eyre!( + "previous turn {} is still in progress", + previous_turn.id + )); + } + + Ok(Some(previous_turn.id.clone())) +} + +#[cfg(test)] +#[path = "safety_buffering_tests.rs"] +mod tests; diff --git a/codex-rs/tui/src/app/safety_buffering_tests.rs b/codex-rs/tui/src/app/safety_buffering_tests.rs new file mode 100644 index 000000000000..19ac62af0f6f --- /dev/null +++ b/codex-rs/tui/src/app/safety_buffering_tests.rs @@ -0,0 +1,67 @@ +use super::*; +use codex_app_server_protocol::TurnItemsView; + +fn turn(id: &str, status: TurnStatus) -> Turn { + Turn { + id: id.to_string(), + is_forkable: true, + items: Vec::new(), + items_view: TurnItemsView::Full, + status, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + } +} + +#[test] +fn retry_forks_through_the_previous_terminal_turn() { + let turns = vec![ + turn("turn-1", TurnStatus::Completed), + turn("turn-2", TurnStatus::Interrupted), + ]; + + assert_eq!( + safety_retry_fork_point(&turns, "turn-2").expect("fork point"), + Some("turn-1".to_string()) + ); +} + +#[test] +fn retry_starts_fresh_when_the_interrupted_turn_is_first() { + let turns = vec![turn("turn-1", TurnStatus::Interrupted)]; + + assert_eq!( + safety_retry_fork_point(&turns, "turn-1").expect("fork point"), + None + ); +} + +#[test] +fn retry_rejects_a_legacy_synthetic_predecessor() { + let turns = vec![ + Turn { + is_forkable: false, + ..turn("legacy-turn", TurnStatus::Completed) + }, + turn("turn-2", TurnStatus::Interrupted), + ]; + + let err = safety_retry_fork_point(&turns, "turn-2").expect_err("synthetic predecessor"); + + assert_eq!( + err.to_string(), + "previous turn legacy-turn has no canonical fork boundary" + ); +} + +#[test] +fn retry_rejects_a_stale_non_latest_turn() { + let turns = vec![ + turn("turn-1", TurnStatus::Interrupted), + turn("turn-2", TurnStatus::InProgress), + ]; + + assert!(safety_retry_fork_point(&turns, "turn-1").is_err()); +} diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 0d9a6f81e24a..2024ffb3d879 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -8,12 +8,6 @@ use super::session_lifecycle::ThreadAttachPresentation; use super::*; use crate::session_resume::read_session_model; -#[derive(Clone, Copy)] -pub(super) enum ThreadRollbackOrigin { - Backtrack, - SafetyBufferingRetry, -} - impl App { pub(super) async fn shutdown_current_thread(&mut self, app_server: &mut AppServerSession) { if let Some(thread_id) = self.chat_widget.thread_id() { @@ -1435,22 +1429,6 @@ impl App { thread_id: ThreadId, num_turns: u32, response: &ThreadRollbackResponse, - ) { - self.handle_thread_rollback_response_with_origin( - thread_id, - num_turns, - response, - ThreadRollbackOrigin::Backtrack, - ) - .await; - } - - pub(super) async fn handle_thread_rollback_response_with_origin( - &mut self, - thread_id: ThreadId, - num_turns: u32, - response: &ThreadRollbackResponse, - origin: ThreadRollbackOrigin, ) { if let Some(channel) = self.thread_event_channels.get(&thread_id) { let mut store = channel.store.lock().await; @@ -1477,12 +1455,7 @@ impl App { self.clear_active_thread().await; } } - match origin { - ThreadRollbackOrigin::Backtrack => self.handle_backtrack_rollback_succeeded(num_turns), - ThreadRollbackOrigin::SafetyBufferingRetry => { - self.apply_non_pending_thread_rollback(num_turns); - } - } + self.handle_backtrack_rollback_succeeded(num_turns); } pub(super) fn handle_thread_event_now(&mut self, event: ThreadBufferedEvent) { diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index 93d88880ec5b..ca75a624bf0e 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -164,12 +164,13 @@ pub(crate) enum AppEvent { op: AppCommand, }, - /// Interrupt, roll back, and retry a safety-buffered turn with the server-selected model. + /// Interrupt, branch, and retry a safety-buffered turn with the server-selected model. RetrySafetyBufferedTurn { thread_id: ThreadId, turn_id: String, model: String, turn: AppCommand, + prompt: UserMessage, }, /// Deliver a synthetic history lookup response to a specific thread channel. diff --git a/codex-rs/tui/src/chatwidget/input_submission.rs b/codex-rs/tui/src/chatwidget/input_submission.rs index a837fa353126..2fcd368687e7 100644 --- a/codex-rs/tui/src/chatwidget/input_submission.rs +++ b/codex-rs/tui/src/chatwidget/input_submission.rs @@ -390,13 +390,15 @@ impl ChatWidget { } if render_in_history { - self.record_cancel_edit_candidate(UserMessage { + let prompt = UserMessage { text: text.clone(), local_images: local_images.clone(), remote_image_urls: remote_image_urls.clone(), text_elements: text_elements.clone(), mention_bindings: mention_bindings.clone(), - }); + }; + self.record_cancel_edit_candidate(prompt.clone()); + self.record_safety_buffering_prompt(prompt); } // Show replayable user content in conversation history. diff --git a/codex-rs/tui/src/chatwidget/safety_buffering.rs b/codex-rs/tui/src/chatwidget/safety_buffering.rs index 50d0ad75e82a..3d15e15ed4b2 100644 --- a/codex-rs/tui/src/chatwidget/safety_buffering.rs +++ b/codex-rs/tui/src/chatwidget/safety_buffering.rs @@ -17,15 +17,60 @@ struct ActiveSafetyBuffering { agent_message_started: bool, } +#[derive(Debug)] +struct SubmittedSafetyBufferingTurn { + turn_id: String, + turn: AppCommand, + prompt: UserMessage, +} + #[derive(Debug, Default)] pub(super) struct SafetyBufferingState { - submitted_turn: Option<(String, AppCommand)>, + pending_prompt: Option, + submitted_turn: Option, active: Option, } impl ChatWidget { + pub(super) fn record_safety_buffering_prompt(&mut self, prompt: UserMessage) { + self.safety_buffering.pending_prompt = Some(prompt); + } + + /// Prepare a safety-buffered retry that bypasses the composer submission path. + /// + /// The retry command is submitted directly after attaching the forked thread, so it must + /// establish the pending state that a normal composer submission would create. The local row + /// is committed separately, after `turn/start` succeeds, so a rejected submission cannot + /// leave a ghost prompt in the transcript. + pub(crate) fn prepare_safety_buffered_retry_submission(&mut self, prompt: UserMessage) { + self.input_queue.user_turn_pending_start = true; + self.record_safety_buffering_prompt(prompt); + } + + /// Render the local retry row after the app-server has accepted the turn. + /// + /// Rendering locally ensures an identical prompt is not mistaken for the final replayed + /// prompt when the app-server echoes the new user item. + pub(crate) fn commit_safety_buffered_retry_submission(&mut self, prompt: UserMessage) { + let display = + user_message_display_for_history(prompt, &UserMessageHistoryRecord::UserMessageText); + self.on_user_message_display(display, /*turn_id*/ None); + } + + pub(crate) fn cancel_safety_buffered_retry_submission(&mut self) { + self.input_queue.user_turn_pending_start = false; + self.clear_safety_buffering(); + } + pub(crate) fn record_safety_buffering_turn(&mut self, turn_id: String, turn: &AppCommand) { - self.safety_buffering.submitted_turn = Some((turn_id, turn.clone())); + self.safety_buffering.submitted_turn = + self.safety_buffering.pending_prompt.take().map(|prompt| { + SubmittedSafetyBufferingTurn { + turn_id, + turn: turn.clone(), + prompt, + } + }); } pub(super) fn reset_safety_buffering_for_turn_start(&mut self) { @@ -62,24 +107,6 @@ impl ChatWidget { .is_some_and(|active| active.turn_id == turn_id && !active.agent_message_started) } - pub(crate) fn prepare_safety_buffering_retry(&mut self) { - let cancel_edit = std::mem::take(&mut self.cancel_edit); - self.last_rendered_user_message_display = None; - self.finalize_turn(); - self.cancel_edit = cancel_edit; - self.input_queue.user_turn_pending_start = true; - } - - pub(crate) fn fail_safety_buffering_retry(&mut self) { - self.input_queue.user_turn_pending_start = false; - self.clear_safety_buffering(); - let prompt = self.cancel_edit.prompt.take(); - self.clear_cancel_edit(); - if let Some(prompt) = prompt { - self.restore_user_message_to_composer(prompt); - } - } - pub(super) fn on_model_safety_buffering_updated( &mut self, notification: ModelSafetyBufferingUpdatedNotification, @@ -116,8 +143,8 @@ impl ChatWidget { .safety_buffering .submitted_turn .as_ref() - .filter(|(submitted_turn_id, _)| replay_kind.is_none() && submitted_turn_id == &turn_id) - .map(|(_, turn)| turn.clone()); + .filter(|submitted_turn| replay_kind.is_none() && submitted_turn.turn_id == turn_id) + .map(|submitted_turn| (submitted_turn.turn.clone(), submitted_turn.prompt.clone())); let thread_id = self.thread_id; let can_offer_retry = faster_model.is_some() && retry_turn.is_some() && thread_id.is_some(); let previous_active = self @@ -165,7 +192,7 @@ impl ChatWidget { } let header = ColumnRenderable::with(header); let mut items = Vec::new(); - if let (Some(faster_model), Some(turn), Some(thread_id)) = + if let (Some(faster_model), Some((turn, prompt)), Some(thread_id)) = (faster_model, retry_turn, thread_id) { items.push(SelectionItem { @@ -176,6 +203,7 @@ impl ChatWidget { turn_id: turn_id.clone(), model: faster_model.clone(), turn: turn.clone(), + prompt: prompt.clone(), }); })], dismiss_on_select: true, diff --git a/codex-rs/tui/src/chatwidget/tests/app_server.rs b/codex-rs/tui/src/chatwidget/tests/app_server.rs index a522aebd3d43..8f21fabcb2e6 100644 --- a/codex-rs/tui/src/chatwidget/tests/app_server.rs +++ b/codex-rs/tui/src/chatwidget/tests/app_server.rs @@ -145,14 +145,15 @@ async fn safety_buffering_offers_one_retry_with_app_wording() { chat.handle_key_event(KeyEvent::new(KeyCode::Up, KeyModifiers::NONE)); chat.handle_key_event(KeyEvent::new(KeyCode::Up, KeyModifiers::NONE)); chat.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); - let (event_thread_id, event_turn_id, model, turn) = loop { + let (event_thread_id, event_turn_id, model, turn, prompt) = loop { match rx.try_recv() { Ok(AppEvent::RetrySafetyBufferedTurn { thread_id, turn_id, model, turn, - }) => break (thread_id, turn_id, model, turn), + prompt, + }) => break (thread_id, turn_id, model, turn, prompt), Ok(_) => continue, Err(err) => panic!("expected safety-buffering retry event: {err}"), } @@ -161,12 +162,82 @@ async fn safety_buffering_offers_one_retry_with_app_wording() { assert_eq!(event_turn_id, turn_id); assert_eq!(model, "faster-model"); assert_matches!(turn, Op::UserTurn { .. }); + assert_eq!(prompt, UserMessage::from("Explain the request")); assert!( !render_bottom_popup(&chat, /*width*/ 80) .contains("Press enter to confirm or esc to go back") ); } +#[tokio::test] +async fn safety_buffered_retry_preparation_renders_identical_prompt_and_preserves_retry_state() { + let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + let thread_id = ThreadId::new(); + let prompt = UserMessage::from("Repeat this request"); + chat.thread_id = Some(thread_id); + + // Match the state of a newly attached branch whose last replayed prompt has the same display. + chat.submit_user_message(prompt.clone()); + let turn = next_submit_op(&mut op_rx); + assert_eq!(drain_insert_history(&mut rx).len(), 1); + chat.clear_safety_buffering(); + chat.input_queue.user_turn_pending_start = false; + + chat.prepare_safety_buffered_retry_submission(prompt.clone()); + assert!(drain_insert_history(&mut rx).is_empty()); + + let retry_turn_id = "retry-turn"; + chat.record_safety_buffering_turn(retry_turn_id.to_string(), &turn); + chat.commit_safety_buffered_retry_submission(prompt); + + let retry_rows = drain_insert_history(&mut rx); + assert_eq!(retry_rows.len(), 1); + assert!(lines_to_single_string(&retry_rows[0]).contains("Repeat this request")); + + chat.handle_server_notification( + ServerNotification::TurnStarted(TurnStartedNotification { + thread_id: thread_id.to_string(), + turn: AppServerTurn { + id: retry_turn_id.to_string(), + is_forkable: true, + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: AppServerTurnStatus::InProgress, + error: None, + started_at: Some(0), + completed_at: None, + duration_ms: None, + }, + }), + /*replay_kind*/ None, + ); + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(safety_buffering_notification( + thread_id, + retry_turn_id, + Some("another-faster-model"), + )), + /*replay_kind*/ None, + ); + + assert!(chat.can_retry_safety_buffered_turn(retry_turn_id)); + assert!(render_bottom_popup(&chat, /*width*/ 80).contains("Retry with a faster model")); +} + +#[tokio::test] +async fn cancelling_safety_buffered_retry_preparation_unblocks_input() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + + chat.prepare_safety_buffered_retry_submission(UserMessage::from("Retry me")); + assert!(chat.input_queue.user_turn_pending_start); + assert!(drain_insert_history(&mut rx).is_empty()); + + chat.cancel_safety_buffered_retry_submission(); + + assert!(!chat.input_queue.user_turn_pending_start); + assert!(drain_insert_history(&mut rx).is_empty()); +} + #[tokio::test] async fn safety_buffering_remains_visible_until_turn_completes() { let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; From 1d13afad4dcbfbe6fc83f4c401cf773184c10c3d Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Thu, 9 Jul 2026 21:30:33 -0300 Subject: [PATCH 2/2] fix(tui): preserve safety retry state --- codex-rs/tui/src/app/safety_buffering.rs | 9 ++- .../tui/src/chatwidget/input_submission.rs | 17 +++--- .../tui/src/chatwidget/safety_buffering.rs | 1 + .../tui/src/chatwidget/tests/app_server.rs | 56 +++++++++++++++++++ 4 files changed, 74 insertions(+), 9 deletions(-) diff --git a/codex-rs/tui/src/app/safety_buffering.rs b/codex-rs/tui/src/app/safety_buffering.rs index 1c10d97bd142..36c54a5258d1 100644 --- a/codex-rs/tui/src/app/safety_buffering.rs +++ b/codex-rs/tui/src/app/safety_buffering.rs @@ -39,6 +39,9 @@ impl App { return; } + let retry_config = self.chat_widget.config_ref().clone(); + let input_state = self.chat_widget.capture_thread_input_state(); + let AppCommand::UserTurn { model: turn_model, effort, @@ -85,14 +88,15 @@ impl App { } }; + self.config = retry_config.clone(); let started = if let Some(last_turn_id) = fork_point { app_server - .fork_thread_after(self.config.clone(), thread_id, last_turn_id) + .fork_thread_after(retry_config.clone(), thread_id, last_turn_id) .await } else { app_server .start_thread_with_session_start_source( - &self.config, + &retry_config, /*session_start_source*/ None, ) .await @@ -121,6 +125,7 @@ impl App { return; } + self.chat_widget.restore_thread_input_state(input_state); self.chat_widget .prepare_safety_buffered_retry_submission(prompt.clone()); if let Err(err) = self diff --git a/codex-rs/tui/src/chatwidget/input_submission.rs b/codex-rs/tui/src/chatwidget/input_submission.rs index 2fcd368687e7..9f774c94bc70 100644 --- a/codex-rs/tui/src/chatwidget/input_submission.rs +++ b/codex-rs/tui/src/chatwidget/input_submission.rs @@ -390,13 +390,16 @@ impl ChatWidget { } if render_in_history { - let prompt = UserMessage { - text: text.clone(), - local_images: local_images.clone(), - remote_image_urls: remote_image_urls.clone(), - text_elements: text_elements.clone(), - mention_bindings: mention_bindings.clone(), - }; + let prompt = user_message_for_restore( + UserMessage { + text: text.clone(), + local_images: local_images.clone(), + remote_image_urls: remote_image_urls.clone(), + text_elements: text_elements.clone(), + mention_bindings: mention_bindings.clone(), + }, + &history_record, + ); self.record_cancel_edit_candidate(prompt.clone()); self.record_safety_buffering_prompt(prompt); } diff --git a/codex-rs/tui/src/chatwidget/safety_buffering.rs b/codex-rs/tui/src/chatwidget/safety_buffering.rs index 3d15e15ed4b2..c4cca3a8d1b6 100644 --- a/codex-rs/tui/src/chatwidget/safety_buffering.rs +++ b/codex-rs/tui/src/chatwidget/safety_buffering.rs @@ -44,6 +44,7 @@ impl ChatWidget { /// leave a ghost prompt in the transcript. pub(crate) fn prepare_safety_buffered_retry_submission(&mut self, prompt: UserMessage) { self.input_queue.user_turn_pending_start = true; + self.record_cancel_edit_candidate(prompt.clone()); self.record_safety_buffering_prompt(prompt); } diff --git a/codex-rs/tui/src/chatwidget/tests/app_server.rs b/codex-rs/tui/src/chatwidget/tests/app_server.rs index 8f21fabcb2e6..128a36a10d2c 100644 --- a/codex-rs/tui/src/chatwidget/tests/app_server.rs +++ b/codex-rs/tui/src/chatwidget/tests/app_server.rs @@ -169,6 +169,59 @@ async fn safety_buffering_offers_one_retry_with_app_wording() { ); } +#[tokio::test] +async fn safety_buffering_retry_uses_the_visible_history_override() { + let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + let thread_id = ThreadId::new(); + let turn_id = "turn-visible-history"; + chat.thread_id = Some(thread_id); + assert!(chat.submit_user_message_with_history_record( + UserMessage::from("hidden raw prompt"), + UserMessageHistoryRecord::Override(UserMessageHistoryOverride { + text: "visible prompt".to_string(), + text_elements: Vec::new(), + }), + )); + let turn = next_submit_op(&mut op_rx); + chat.record_safety_buffering_turn(turn_id.to_string(), &turn); + chat.handle_server_notification( + ServerNotification::TurnStarted(TurnStartedNotification { + thread_id: thread_id.to_string(), + turn: AppServerTurn { + id: turn_id.to_string(), + is_forkable: true, + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: AppServerTurnStatus::InProgress, + error: None, + started_at: Some(0), + completed_at: None, + duration_ms: None, + }, + }), + /*replay_kind*/ None, + ); + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(safety_buffering_notification( + thread_id, + turn_id, + Some("faster-model"), + )), + /*replay_kind*/ None, + ); + + chat.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + let prompt = loop { + match rx.try_recv() { + Ok(AppEvent::RetrySafetyBufferedTurn { prompt, .. }) => break prompt, + Ok(_) => continue, + Err(err) => panic!("expected safety-buffering retry event: {err}"), + } + }; + + assert_eq!(prompt, UserMessage::from("visible prompt")); +} + #[tokio::test] async fn safety_buffered_retry_preparation_renders_identical_prompt_and_preserves_retry_state() { let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; @@ -185,6 +238,9 @@ async fn safety_buffered_retry_preparation_renders_identical_prompt_and_preserve chat.prepare_safety_buffered_retry_submission(prompt.clone()); assert!(drain_insert_history(&mut rx).is_empty()); + assert_eq!(chat.cancel_edit.prompt, Some(prompt.clone())); + assert!(chat.cancel_edit.eligible); + assert!(!chat.cancel_edit.armed); let retry_turn_id = "retry-turn"; chat.record_safety_buffering_turn(retry_turn_id.to_string(), &turn);