diff --git a/codex-rs/tui/src/app/safety_buffering.rs b/codex-rs/tui/src/app/safety_buffering.rs index 1c10d97bd1..36c54a5258 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 2fcd368687..9f774c94bc 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 3d15e15ed4..c4cca3a8d1 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 8f21fabcb2..128a36a10d 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);