diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 75c7e692fa..c89e352f78 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -123,7 +123,6 @@ use codex_app_server_protocol::SkillsListResponse; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadLoadedListParams; use codex_app_server_protocol::ThreadMemoryMode; -use codex_app_server_protocol::ThreadRollbackResponse; use codex_app_server_protocol::ThreadStartSource; use codex_app_server_protocol::Turn; use codex_app_server_protocol::TurnError as AppServerTurnError; @@ -214,6 +213,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 ed2478eb02..1134d1f434 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -411,73 +411,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( - thread_id, - /*num_turns*/ 1, - &rollback_response, + 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::AppendMessageHistoryEntry { thread_id, text } => { self.append_message_history_entry(thread_id, text); 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 0000000000..f61649460d --- /dev/null +++ b/codex-rs/tui/src/app/safety_buffering.rs @@ -0,0 +1,231 @@ +//! Safety-buffered turn retries that preserve the source thread. + +use super::session_lifecycle::ThreadAttachPresentation; +use super::*; +use crate::chatwidget::ThreadInputState; +use crate::chatwidget::ThreadInputStateRestoreMode; +use crate::chatwidget::UserMessage; +use codex_app_server_protocol::UserInput; + +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, +} + +enum SafetyRetryForkPoint { + AfterTurn(String), + NewThread, +} + +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) + || self.primary_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 retry_config = self.chat_widget.config_ref().clone(); + let input_state = self.chat_widget.capture_thread_input_state(); + + let AppCommand::UserTurn { + items, + 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(input_state, 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(input_state, prompt, err); + return; + } + }; + items.extend( + thread + .turns + .iter() + .find(|turn| turn.id == turn_id) + .into_iter() + .flat_map(|turn| &turn.items) + .filter_map(|item| match item { + ThreadItem::UserMessage { content, .. } => Some(content), + _ => None, + }) + .skip(/*n*/ 1) + .flat_map(|content| { + std::iter::once(UserInput::Text { + text: "\n".to_string(), + text_elements: Vec::new(), + }) + .chain(content.iter().cloned()) + }), + ); + self.config = retry_config.clone(); + let started = match fork_point { + SafetyRetryForkPoint::AfterTurn(last_turn_id) => { + app_server + .fork_thread_after(retry_config, thread_id, last_turn_id) + .await + } + SafetyRetryForkPoint::NewThread => { + app_server + .start_thread_with_session_start_source( + &retry_config, + /*session_start_source*/ None, + ) + .await + } + }; + let started = match started { + Ok(started) => started, + Err(err) => { + self.fail_safety_buffered_branch(input_state, 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(input_state, prompt, err); + return; + } + + let failure_input_state = input_state.clone(); + self.chat_widget.restore_thread_input_state( + input_state, + ThreadInputStateRestoreMode { + preserve_in_flight_turn: false, + }, + ); + 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(failure_input_state, prompt, err); + return; + } + self.chat_widget + .commit_safety_buffered_retry_submission(prompt); + } + + fn fail_safety_buffered_branch( + &mut self, + input_state: Option, + prompt: UserMessage, + err: impl std::fmt::Display, + ) { + self.chat_widget.restore_thread_input_state( + input_state, + ThreadInputStateRestoreMode { + preserve_in_flight_turn: false, + }, + ); + 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(SafetyRetryForkPoint::NewThread); + }; + if previous_turn.status == TurnStatus::InProgress { + return Err(color_eyre::eyre::eyre!( + "previous turn {} is still in progress", + previous_turn.id + )); + } + + Ok(SafetyRetryForkPoint::AfterTurn(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 0000000000..4d1c2b7196 --- /dev/null +++ b/codex-rs/tui/src/app/safety_buffering_tests.rs @@ -0,0 +1,33 @@ +use super::*; +use codex_app_server_protocol::TurnItemsView; + +fn turn(id: &str, status: TurnStatus) -> Turn { + Turn { + id: id.to_string(), + items: Vec::new(), + items_view: TurnItemsView::Full, + status, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + } +} + +#[test] +fn retry_rejects_a_stale_or_in_progress_turn() { + let stale = vec![ + turn("turn-1", TurnStatus::Interrupted), + turn("turn-2", TurnStatus::InProgress), + ]; + let in_progress = vec![turn("turn-1", TurnStatus::InProgress)]; + let previous_in_progress = vec![ + turn("turn-1", TurnStatus::InProgress), + turn("turn-2", TurnStatus::Interrupted), + ]; + + assert!(safety_retry_fork_point(&stale, "turn-1").is_err()); + assert!(safety_retry_fork_point(&in_progress, "turn-1").is_err()); + assert!(safety_retry_fork_point(&in_progress, "missing").is_err()); + assert!(safety_retry_fork_point(&previous_in_progress, "turn-2").is_err()); +} diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index c448242984..20886160e9 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -5,6 +5,7 @@ mod advanced_reasoning_tests; mod model_catalog; mod plugin_catalog; mod rate_limits; +mod safety_buffering; mod session_summary; mod startup; @@ -45,7 +46,6 @@ use codex_app_server_protocol::AdditionalPermissionProfile; use codex_app_server_protocol::AgentMessageDeltaNotification; use codex_app_server_protocol::AskForApproval; use codex_app_server_protocol::CommandExecutionRequestApprovalParams; -use codex_app_server_protocol::ConfigWarningNotification; use codex_app_server_protocol::FileChangeRequestApprovalParams; use codex_app_server_protocol::FileUpdateChange; use codex_app_server_protocol::ItemStartedNotification; @@ -64,7 +64,6 @@ use codex_app_server_protocol::PermissionsRequestApprovalParams; use codex_app_server_protocol::RequestId as AppServerRequestId; use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ServerRequest; -use codex_app_server_protocol::SessionSource; use codex_app_server_protocol::Thread; use codex_app_server_protocol::ThreadClosedNotification; use codex_app_server_protocol::ThreadItem; @@ -5803,83 +5802,6 @@ async fn refreshed_snapshot_session_persists_resumed_turns() { assert_eq!(store_snapshot.turns, snapshot.turns); } -#[tokio::test] -async fn queued_rollback_syncs_overlay_and_clears_deferred_history() { - let mut app = make_test_app().await; - app.transcript_cells = vec![ - Arc::new(UserHistoryCell { - message: "first".to_string(), - text_elements: Vec::new(), - local_image_paths: Vec::new(), - remote_image_urls: Vec::new(), - }) as Arc, - Arc::new(AgentMessageCell::new( - vec![Line::from("after first")], - /*is_first_line*/ false, - )) as Arc, - Arc::new(UserHistoryCell { - message: "second".to_string(), - text_elements: Vec::new(), - local_image_paths: Vec::new(), - remote_image_urls: Vec::new(), - }) as Arc, - Arc::new(AgentMessageCell::new( - vec![Line::from("after second")], - /*is_first_line*/ false, - )) as Arc, - ]; - app.overlay = Some(Overlay::new_transcript( - app.transcript_cells.clone(), - app.keymap.pager.clone(), - )); - app.deferred_history_lines = vec![Line::from("stale buffered line").into()]; - app.backtrack.overlay_preview_active = true; - app.backtrack.nth_user_message = 1; - app.chat_widget.update_account_state( - /*status_account_display*/ None, /*plan_type*/ None, - /*has_chatgpt_account*/ false, /*has_codex_backend_auth*/ true, - ); - app.chat_widget.insert_str("/usage daily"); - app.chat_widget - .handle_key_event(KeyEvent::new(KeyCode::Esc, KeyModifiers::NONE)); - app.chat_widget - .handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); - app.chat_widget - .handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); - let pending_usage = app - .chat_widget - .active_cell_transcript_lines(/*width*/ 80) - .expect("pending usage transcript"); - assert!(lines_to_single_string(&pending_usage).contains("Token activity\n Loading...")); - - let changed = app.apply_non_pending_thread_rollback(/*num_turns*/ 1); - - assert!(changed); - assert!(app.backtrack_render_pending); - assert!(app.deferred_history_lines.is_empty()); - assert!( - app.chat_widget - .active_cell_transcript_lines(/*width*/ 80) - .is_none_or(|lines| !lines_to_single_string(&lines).contains("Token activity")) - ); - assert_eq!(app.backtrack.nth_user_message, 0); - 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".to_string()]); - let overlay_cell_count = match app.overlay.as_ref() { - Some(Overlay::Transcript(t)) => t.committed_cell_count(), - _ => panic!("expected transcript overlay"), - }; - assert_eq!(overlay_cell_count, app.transcript_cells.len()); -} - #[tokio::test] async fn late_usage_result_can_follow_finalized_plan() { let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; @@ -5912,64 +5834,6 @@ async fn late_usage_result_can_follow_finalized_plan() { ); } -#[tokio::test] -async fn thread_rollback_response_discards_queued_active_thread_events() { - let mut app = make_test_app().await; - let thread_id = ThreadId::new(); - let (tx, rx) = mpsc::channel(8); - app.active_thread_id = Some(thread_id); - app.active_thread_rx = Some(rx); - tx.send(ThreadBufferedEvent::Notification( - ServerNotification::ConfigWarning(ConfigWarningNotification { - summary: "stale warning".to_string(), - details: None, - path: None, - range: None, - }), - )) - .await - .expect("event should queue"); - - app.handle_thread_rollback_response( - thread_id, - /*num_turns*/ 1, - &ThreadRollbackResponse { - thread: Thread { - id: thread_id.to_string(), - extra: None, - session_id: thread_id.to_string(), - forked_from_id: None, - parent_thread_id: None, - preview: String::new(), - ephemeral: false, - history_mode: Default::default(), - model_provider: "openai".to_string(), - created_at: 0, - updated_at: 0, - recency_at: Some(0), - status: codex_app_server_protocol::ThreadStatus::Idle, - path: None, - cwd: test_path_buf("/tmp/project").abs(), - cli_version: "0.0.0".to_string(), - source: SessionSource::Cli, - thread_source: None, - agent_nickname: None, - agent_role: None, - git_info: None, - name: None, - turns: Vec::new(), - }, - }, - ) - .await; - - let rx = app - .active_thread_rx - .as_mut() - .expect("active receiver should remain attached"); - assert!(matches!(rx.try_recv(), Err(TryRecvError::Empty))); -} - #[tokio::test] async fn new_session_requests_shutdown_for_previous_conversation() { Box::pin(async { diff --git a/codex-rs/tui/src/app/tests/safety_buffering.rs b/codex-rs/tui/src/app/tests/safety_buffering.rs new file mode 100644 index 0000000000..f4929027b4 --- /dev/null +++ b/codex-rs/tui/src/app/tests/safety_buffering.rs @@ -0,0 +1,490 @@ +use super::*; +use crate::app::safety_buffering::SafetyBufferedRetry; +use crate::app::session_lifecycle::ThreadAttachPresentation; +use crate::chatwidget::UserMessage; +use codex_app_server_client::AppServerEvent; +use codex_app_server_protocol::ModelSafetyBufferingUpdatedNotification; +use core_test_support::responses; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::streaming_sse::StreamingSseChunk; +use core_test_support::streaming_sse::StreamingSseServer; +use core_test_support::streaming_sse::start_streaming_sse_server; +use pretty_assertions::assert_eq; +use serde_json::Value; +use tokio::sync::oneshot; + +const CURRENT_MODEL: &str = "gpt-5.2"; +const FASTER_MODEL: &str = "gpt-5.4"; +const MODEL_PROVIDER_ID: &str = "safety-retry-test"; +const PREVIOUS_PROMPT: &str = "Establish context"; +const RETRY_PROMPT: &str = "Handle the safety-buffered request"; +const COMMITTED_STEER: &str = "Keep the accepted steer"; +const UNSENT_DRAFT: &str = "Keep this unsent draft"; + +fn response_chunks(response_id: &str) -> Vec { + [ + ev_response_created(response_id), + ev_assistant_message(&format!("message-{response_id}"), "done"), + ev_completed(response_id), + ] + .into_iter() + .map(|event| StreamingSseChunk { + gate: None, + body: responses::sse(vec![event]), + }) + .collect() +} + +fn gated_response_chunks(response_id: &str) -> (Vec, oneshot::Sender<()>) { + let (release_tx, release_rx) = oneshot::channel(); + ( + vec![ + StreamingSseChunk { + gate: None, + body: responses::sse(vec![ev_response_created(response_id)]), + }, + StreamingSseChunk { + gate: Some(release_rx), + body: responses::sse(vec![ev_completed(response_id)]), + }, + ], + release_tx, + ) +} + +fn next_user_turn_event( + app_event_rx: &mut tokio::sync::mpsc::UnboundedReceiver, +) -> AppCommand { + while let Ok(event) = app_event_rx.try_recv() { + if let AppEvent::CodexOp(turn @ AppCommand::UserTurn { .. }) = event { + return turn; + } + } + panic!("expected UserTurn app event"); +} + +fn submit_prompt(app: &mut App, prompt: &str) { + app.chat_widget.apply_external_edit(prompt.to_string()); + app.chat_widget + .handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); +} + +fn drain_active_thread_events(app: &mut App) { + while let Some(event) = app + .active_thread_rx + .as_mut() + .and_then(|receiver| receiver.try_recv().ok()) + { + app.handle_thread_event_now(event); + } +} + +async fn next_turn_started( + app: &mut App, + app_server: &mut AppServerSession, + thread_id: ThreadId, +) -> String { + loop { + let event = tokio::time::timeout( + std::time::Duration::from_secs(/*secs*/ 5), + app_server.next_event(), + ) + .await + .expect("app-server should emit a turn/start event") + .expect("app-server event stream should remain open"); + let started_turn_id = match &event { + AppServerEvent::ServerNotification(ServerNotification::TurnStarted(notification)) + if notification.thread_id == thread_id.to_string() => + { + Some(notification.turn.id.clone()) + } + _ => None, + }; + app.handle_app_server_event(app_server, event).await; + drain_active_thread_events(app); + if let Some(turn_id) = started_turn_id { + return turn_id; + } + } +} + +async fn wait_for_turn_completed( + app: &mut App, + app_server: &mut AppServerSession, + thread_id: ThreadId, +) { + loop { + let event = tokio::time::timeout( + std::time::Duration::from_secs(/*secs*/ 5), + app_server.next_event(), + ) + .await + .expect("app-server should emit a turn/completed event") + .expect("app-server event stream should remain open"); + let completed = matches!( + &event, + AppServerEvent::ServerNotification(ServerNotification::TurnCompleted(notification)) + if notification.thread_id == thread_id.to_string() + ); + app.handle_app_server_event(app_server, event).await; + drain_active_thread_events(app); + if completed { + return; + } + } +} + +async fn drive_until_request_count( + app: &mut App, + app_server: &mut AppServerSession, + server: &StreamingSseServer, + expected_request_count: usize, +) { + let timeout = tokio::time::sleep(std::time::Duration::from_secs(/*secs*/ 5)); + tokio::pin!(timeout); + loop { + tokio::select! { + () = server.wait_for_request_count(expected_request_count) => return, + event = app_server.next_event() => { + let event = event.expect("app-server event stream should remain open"); + app.handle_app_server_event(app_server, event).await; + drain_active_thread_events(app); + } + () = &mut timeout => { + panic!("expected {expected_request_count} Responses API requests"); + } + } + } +} + +fn user_input_texts(body: &Value) -> Vec { + body.get("input") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter(|item| item.get("type").and_then(Value::as_str) == Some("message")) + .filter(|item| item.get("role").and_then(Value::as_str) == Some("user")) + .filter_map(|item| item.get("content").and_then(Value::as_array)) + .flatten() + .filter(|span| span.get("type").and_then(Value::as_str) == Some("input_text")) + .filter_map(|span| span.get("text").and_then(Value::as_str).map(str::to_owned)) + .collect() +} + +fn user_message_count(thread: &Thread, prompt: &str) -> usize { + thread + .turns + .iter() + .flat_map(|turn| &turn.items) + .filter_map(|item| match item { + ThreadItem::UserMessage { content, .. } => Some(content), + _ => None, + }) + .flatten() + .filter(|item| matches!(item, AppServerUserInput::Text { text, .. } if text == prompt)) + .count() +} + +async fn run_safety_retry( + previous_prompt: Option<&str>, + failing_draft: Option<&str>, + committed_steer: Option<&str>, +) -> Result<()> { + let (active_chunks, release_active_response) = gated_response_chunks("active-response"); + let mut release_active_response = Some(release_active_response); + let (steered_chunks, release_steered_response) = gated_response_chunks("steered-response"); + let mut response_sequences = Vec::new(); + if previous_prompt.is_some() { + response_sequences.push(response_chunks("previous-response")); + } + response_sequences.push(active_chunks); + if committed_steer.is_some() { + response_sequences.push(steered_chunks); + } + response_sequences.push(response_chunks("retry-response")); + let expected_request_count = response_sequences.len(); + let (server, _completions) = start_streaming_sse_server(response_sequences).await; + + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; + let codex_home = tempdir()?; + std::fs::write( + codex_home.path().join("config.toml"), + format!( + r#" +model = "{CURRENT_MODEL}" +model_provider = "{MODEL_PROVIDER_ID}" + +[model_providers.{MODEL_PROVIDER_ID}] +name = "Safety retry test" +base_url = "{}/v1" +wire_api = "responses" +request_max_retries = 0 +stream_max_retries = 0 +"#, + server.uri() + ), + )?; + app.config.codex_home = codex_home.path().to_path_buf().abs(); + app.config.sqlite_home = codex_home.path().to_path_buf(); + app.config.model = Some(CURRENT_MODEL.to_string()); + app.config.model_provider_id = MODEL_PROVIDER_ID.to_string(); + app.config.model_provider = ModelProviderInfo { + name: "Safety retry test".to_string(), + base_url: Some(format!("{}/v1", server.uri())), + request_max_retries: Some(0), + stream_max_retries: Some(0), + ..ModelProviderInfo::default() + }; + + let mut tui = crate::tui::test_support::make_test_tui()?; + let mut app_server = Box::pin(crate::start_embedded_app_server_for_picker(&app.config)).await?; + let started = app_server.start_thread(&app.config).await?; + let source_thread_id = started.session.thread_id; + app.replace_chat_widget_with_app_server_thread( + &mut tui, + &mut app_server, + started, + ThreadAttachPresentation::SessionLineage, + /*initial_user_message*/ None, + ) + .await?; + while app_event_rx.try_recv().is_ok() {} + + if let Some(previous_prompt) = previous_prompt { + submit_prompt(&mut app, previous_prompt); + let previous_turn = next_user_turn_event(&mut app_event_rx); + app.submit_thread_op(&mut app_server, source_thread_id, previous_turn) + .await?; + wait_for_turn_completed(&mut app, &mut app_server, source_thread_id).await; + } + + submit_prompt(&mut app, RETRY_PROMPT); + let active_turn = next_user_turn_event(&mut app_event_rx); + app.submit_thread_op(&mut app_server, source_thread_id, active_turn.clone()) + .await?; + let active_turn_id = next_turn_started(&mut app, &mut app_server, source_thread_id).await; + drive_until_request_count( + &mut app, + &mut app_server, + &server, + usize::from(previous_prompt.is_some()) + 1, + ) + .await; + + if let Some(committed_steer) = committed_steer { + submit_prompt(&mut app, committed_steer); + let steer = next_user_turn_event(&mut app_event_rx); + app.submit_thread_op(&mut app_server, source_thread_id, steer) + .await?; + let _ = release_active_response + .take() + .expect("active response should still be gated") + .send(()); + drive_until_request_count( + &mut app, + &mut app_server, + &server, + usize::from(previous_prompt.is_some()) + 2, + ) + .await; + let source = app_server + .thread_read(source_thread_id, /*include_turns*/ true) + .await?; + assert_eq!(user_message_count(&source, committed_steer), 1); + } + + app.handle_app_server_event( + &app_server, + AppServerEvent::ServerNotification(ServerNotification::ModelSafetyBufferingUpdated( + ModelSafetyBufferingUpdatedNotification { + thread_id: source_thread_id.to_string(), + turn_id: active_turn_id.clone(), + model: CURRENT_MODEL.to_string(), + use_cases: Vec::new(), + reasons: Vec::new(), + show_buffering_ui: true, + faster_model: Some(FASTER_MODEL.to_string()), + }, + )), + ) + .await; + drain_active_thread_events(&mut app); + assert!( + app.chat_widget + .can_retry_safety_buffered_turn(&active_turn_id) + ); + if let Some(draft) = failing_draft { + app.chat_widget.apply_external_edit(draft.to_string()); + let source = app_server + .thread_read(source_thread_id, /*include_turns*/ true) + .await?; + std::fs::remove_file( + source + .path + .expect("source thread should have a rollout path"), + )?; + } + + let primary_thread_id = ThreadId::new(); + app.primary_thread_id = Some(primary_thread_id); + Box::pin(app.retry_safety_buffered_turn( + &mut tui, + &mut app_server, + SafetyBufferedRetry { + thread_id: source_thread_id, + turn_id: active_turn_id.clone(), + model: FASTER_MODEL.to_string(), + turn: active_turn.clone(), + prompt: UserMessage::from(RETRY_PROMPT), + }, + )) + .await; + assert_eq!(app.primary_thread_id, Some(primary_thread_id)); + assert_eq!(app.active_thread_id, Some(source_thread_id)); + assert_eq!(app.chat_widget.thread_id(), Some(source_thread_id)); + app.primary_thread_id = Some(source_thread_id); + + Box::pin(app.retry_safety_buffered_turn( + &mut tui, + &mut app_server, + SafetyBufferedRetry { + thread_id: source_thread_id, + turn_id: active_turn_id, + model: FASTER_MODEL.to_string(), + turn: active_turn, + prompt: UserMessage::from(RETRY_PROMPT), + }, + )) + .await; + + if let Some(draft) = failing_draft { + assert_eq!( + app.chat_widget.composer_text_with_pending(), + format!("{RETRY_PROMPT}\n{draft}") + ); + assert_eq!(app.chat_widget.thread_id(), Some(source_thread_id)); + if let Some(release_active_response) = release_active_response.take() { + let _ = release_active_response.send(()); + } + let _ = release_steered_response.send(()); + app_server.shutdown().await?; + server.shutdown().await; + return Ok(()); + } + + drive_until_request_count(&mut app, &mut app_server, &server, expected_request_count).await; + let mut replayed_history = String::new(); + while let Ok(event) = app_event_rx.try_recv() { + if let AppEvent::InsertHistoryCell(cell) = event { + replayed_history.push_str(&lines_to_single_string( + &cell.transcript_lines(/*width*/ 80), + )); + } + } + assert!(!replayed_history.contains("Conversation interrupted")); + + let retry_thread_id = app.chat_widget.thread_id().expect("retry thread id"); + let source = app_server + .thread_read(source_thread_id, /*include_turns*/ true) + .await?; + let retry = app_server + .thread_read(retry_thread_id, /*include_turns*/ true) + .await?; + assert_ne!(retry_thread_id, source_thread_id); + assert_eq!( + source.turns.last().map(|turn| &turn.status), + Some(&TurnStatus::Interrupted) + ); + let expected_forked_from_id = previous_prompt.map(|_| source_thread_id.to_string()); + assert_eq!( + retry.forked_from_id.as_deref(), + expected_forked_from_id.as_deref() + ); + let expected_retry_prompt = match committed_steer { + Some(committed_steer) => format!("{RETRY_PROMPT}\n{committed_steer}"), + None => RETRY_PROMPT.to_string(), + }; + assert_eq!(user_message_count(&source, RETRY_PROMPT), 1); + assert_eq!(user_message_count(&retry, &expected_retry_prompt), 1); + if let Some(committed_steer) = committed_steer { + assert_eq!(user_message_count(&source, committed_steer), 1); + } + if let Some(previous_prompt) = previous_prompt { + assert_eq!(user_message_count(&source, previous_prompt), 1); + assert_eq!(user_message_count(&retry, previous_prompt), 1); + } + + let request_bodies = server + .requests() + .await + .iter() + .map(|request| serde_json::from_slice::(request)) + .collect::>>()?; + let retry_request = request_bodies + .last() + .expect("retry should issue a Responses API request"); + assert_eq!(retry_request["model"].as_str(), Some(FASTER_MODEL)); + assert_eq!(retry_request["reasoning"]["effort"].as_str(), Some("low")); + let relevant_prompts = user_input_texts(retry_request) + .into_iter() + .filter(|text| { + text == PREVIOUS_PROMPT + || text == RETRY_PROMPT + || committed_steer.is_some_and(|steer| text == steer) + }) + .collect::>(); + let mut expected_prompts = match previous_prompt { + Some(previous_prompt) => vec![previous_prompt.to_string(), RETRY_PROMPT.to_string()], + None => vec![RETRY_PROMPT.to_string()], + }; + expected_prompts.extend(committed_steer.map(str::to_string)); + assert_eq!(relevant_prompts, expected_prompts); + + if let Some(release_active_response) = release_active_response.take() { + let _ = release_active_response.send(()); + } + let _ = release_steered_response.send(()); + app_server.shutdown().await?; + server.shutdown().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn safety_retry_forks_after_the_previous_turn_and_uses_faster_settings() -> Result<()> { + run_safety_retry( + Some(PREVIOUS_PROMPT), + /*failing_draft*/ None, + /*committed_steer*/ None, + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn safety_retry_preserves_a_committed_steer_from_the_interrupted_turn() -> Result<()> { + run_safety_retry( + Some(PREVIOUS_PROMPT), + /*failing_draft*/ None, + Some(COMMITTED_STEER), + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn safety_retry_starts_new_thread_for_first_turn_without_duplicating_prompt() -> Result<()> { + run_safety_retry( + /*previous_prompt*/ None, /*failing_draft*/ None, /*committed_steer*/ None, + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn safety_retry_branch_failure_preserves_unsent_draft() -> Result<()> { + run_safety_retry( + Some(PREVIOUS_PROMPT), + Some(UNSENT_DRAFT), + /*committed_steer*/ None, + ) + .await +} diff --git a/codex-rs/tui/src/app/thread_events.rs b/codex-rs/tui/src/app/thread_events.rs index 8d41f1c052..45f1c0321d 100644 --- a/codex-rs/tui/src/app/thread_events.rs +++ b/codex-rs/tui/src/app/thread_events.rs @@ -198,13 +198,6 @@ impl ThreadEventStore { }) } - pub(super) fn apply_thread_rollback(&mut self, response: &ThreadRollbackResponse) { - self.turns = response.thread.turns.clone(); - self.buffer.clear(); - self.pending_interactive_replay = PendingInteractiveReplayState::default(); - self.active_turn_id = None; - } - pub(super) fn snapshot(&self) -> ThreadEventSnapshot { ThreadEventSnapshot { session: self.session.clone(), diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 171f3968ed..4fae845fc1 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -1416,40 +1416,6 @@ impl App { self.chat_widget.handle_skills_list_response(response); } - pub(super) async fn handle_thread_rollback_response( - &mut self, - thread_id: ThreadId, - num_turns: u32, - response: &ThreadRollbackResponse, - ) { - if let Some(channel) = self.thread_event_channels.get(&thread_id) { - let mut store = channel.store.lock().await; - store.apply_thread_rollback(response); - } - if self.active_thread_id == Some(thread_id) - && let Some(mut rx) = self.active_thread_rx.take() - { - let mut disconnected = false; - loop { - match rx.try_recv() { - Ok(_) => {} - Err(TryRecvError::Empty) => break, - Err(TryRecvError::Disconnected) => { - disconnected = true; - break; - } - } - } - - if !disconnected { - self.active_thread_rx = Some(rx); - } else { - self.clear_active_thread().await; - } - } - self.apply_non_pending_thread_rollback(num_turns); - } - pub(super) fn handle_thread_event_now(&mut self, event: ThreadBufferedEvent) { let needs_refresh = matches!( &event, diff --git a/codex-rs/tui/src/app_backtrack.rs b/codex-rs/tui/src/app_backtrack.rs index 65fa7e4c26..49b841186a 100644 --- a/codex-rs/tui/src/app_backtrack.rs +++ b/codex-rs/tui/src/app_backtrack.rs @@ -430,22 +430,6 @@ impl App { self.chat_widget.clear_esc_backtrack_hint(); } - /// Apply rollback semantics for a confirmed safety-buffering retry rollback. - /// - /// Returns `true` when local transcript state changed. - pub(crate) fn apply_non_pending_thread_rollback(&mut self, num_turns: u32) -> bool { - if !trim_transcript_cells_drop_last_n_user_turns(&mut self.transcript_cells, num_turns) { - return false; - } - self.chat_widget.clear_pending_token_activity_refreshes(); - self.chat_widget.clear_pending_rate_limit_reset_hint(); - self.chat_widget - .truncate_agent_copy_history_to_user_turn_count(user_count(&self.transcript_cells)); - self.sync_overlay_after_transcript_trim(); - self.backtrack_render_pending = true; - true - } - fn backtrack_selection(&self, nth_user_message: usize) -> Option { let base_id = self.backtrack.base_id?; if self.chat_widget.thread_id() != Some(base_id) { @@ -477,33 +461,6 @@ impl App { }, }) } - - /// Keep transcript-related UI state aligned after `transcript_cells` was trimmed. - /// - /// This does three things: - /// 1. If transcript overlay is open, replace its committed cells so removed turns disappear. - /// 2. If backtrack preview is active, clamp/recompute the highlighted user selection. - /// 3. Drop deferred transcript lines buffered while overlay was open to avoid flushing lines - /// for cells that were just removed by the trim. - fn sync_overlay_after_transcript_trim(&mut self) { - if let Some(Overlay::Transcript(t)) = &mut self.overlay { - t.replace_cells(self.transcript_cells.clone()); - } - if self.backtrack.overlay_preview_active { - let total_users = user_count(&self.transcript_cells); - let next_selection = if total_users == 0 { - usize::MAX - } else { - self.backtrack - .nth_user_message - .min(total_users.saturating_sub(1)) - }; - self.apply_backtrack_selection_internal(next_selection); - } - // While overlay is open, we buffer rendered history lines and flush them on close. - // If rollback trimmed cells meanwhile, those buffered lines can reference removed turns. - self.deferred_history_lines.clear(); - } } /// Find the persisted turn immediately before a selected transcript prompt. @@ -582,30 +539,6 @@ pub(crate) fn backtrack_fork_last_turn_id( bail!("the selected prompt was not found in the persisted thread") } -pub(crate) fn trim_transcript_cells_drop_last_n_user_turns( - transcript_cells: &mut Vec>, - num_turns: u32, -) -> bool { - if num_turns == 0 { - return false; - } - - let user_positions: Vec = user_positions_iter(transcript_cells).collect(); - let Some(&first_user_idx) = user_positions.first() else { - return false; - }; - - let turns_from_end = usize::try_from(num_turns).unwrap_or(usize::MAX); - let cut_idx = if turns_from_end >= user_positions.len() { - first_user_idx - } else { - user_positions[user_positions.len() - turns_from_end] - }; - let original_len = transcript_cells.len(); - transcript_cells.truncate(cut_idx); - transcript_cells.len() != original_len -} - pub(crate) fn user_count(cells: &[Arc]) -> usize { user_positions_iter(cells).count() } @@ -918,79 +851,6 @@ mod tests { ); } - #[test] - fn trim_drop_last_n_user_turns_applies_rollback_semantics() { - let mut cells: Vec> = vec![ - Arc::new(UserHistoryCell { - message: "first".to_string(), - text_elements: Vec::new(), - local_image_paths: Vec::new(), - remote_image_urls: Vec::new(), - }) as Arc, - Arc::new(AgentMessageCell::new( - vec![Line::from("after first")], - /*is_first_line*/ false, - )) as Arc, - Arc::new(UserHistoryCell { - message: "second".to_string(), - text_elements: Vec::new(), - local_image_paths: Vec::new(), - remote_image_urls: Vec::new(), - }) as Arc, - Arc::new(AgentMessageCell::new( - vec![Line::from("after second")], - /*is_first_line*/ false, - )) as Arc, - ]; - - let changed = - trim_transcript_cells_drop_last_n_user_turns(&mut cells, /*num_turns*/ 1); - - assert!(changed); - assert_eq!(cells.len(), 2); - let first_user = cells[0] - .as_any() - .downcast_ref::() - .expect("first user"); - assert_eq!(first_user.message, "first"); - } - - #[test] - fn trim_drop_last_n_user_turns_allows_overflow() { - let mut cells: Vec> = vec![ - Arc::new(AgentMessageCell::new( - vec![Line::from("intro")], - /*is_first_line*/ true, - )) as Arc, - Arc::new(UserHistoryCell { - message: "first".to_string(), - text_elements: Vec::new(), - local_image_paths: Vec::new(), - remote_image_urls: Vec::new(), - }) as Arc, - Arc::new(AgentMessageCell::new( - vec![Line::from("after")], - /*is_first_line*/ false, - )) as Arc, - ]; - - let changed = trim_transcript_cells_drop_last_n_user_turns(&mut cells, u32::MAX); - - assert!(changed); - assert_eq!(cells.len(), 1); - let intro = cells[0] - .as_any() - .downcast_ref::() - .expect("intro agent"); - let intro_lines = intro.display_lines(u16::MAX); - let intro_text: String = intro_lines[0] - .spans - .iter() - .map(|span| span.content.as_ref()) - .collect(); - assert_eq!(intro_text, "• intro"); - } - #[test] fn agent_group_count_ignores_context_compacted_marker() { let cells: Vec> = vec![ diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index 6e6937b8c0..9ba10f9592 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, fork, 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/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index c54cedc43e..0cb445b3bc 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -85,8 +85,6 @@ use codex_app_server_protocol::ThreadReadParams; use codex_app_server_protocol::ThreadReadResponse; use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadResumeResponse; -use codex_app_server_protocol::ThreadRollbackParams; -use codex_app_server_protocol::ThreadRollbackResponse; use codex_app_server_protocol::ThreadSetNameParams; use codex_app_server_protocol::ThreadSetNameResponse; use codex_app_server_protocol::ThreadSettingsUpdateParams; @@ -1114,24 +1112,6 @@ impl AppServerSession { Ok(()) } - pub(crate) async fn thread_rollback( - &mut self, - thread_id: ThreadId, - num_turns: u32, - ) -> Result { - let request_id = self.next_request_id(); - self.client - .request_typed(ClientRequest::ThreadRollback { - request_id, - params: ThreadRollbackParams { - thread_id: thread_id.to_string(), - num_turns, - }, - }) - .await - .wrap_err("thread/rollback failed in TUI") - } - pub(crate) async fn review_start( &mut self, thread_id: ThreadId, diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index d8a76cf3aa..424e8310f8 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -475,7 +475,6 @@ const APPROVE_FOR_ME_LABEL: &str = "Approve for me"; const AUTO_REVIEW_DESCRIPTION: &str = "Only ask for actions detected as potentially unsafe."; const DEFAULT_OPENAI_BASE_URL: &str = "https://api.openai.com/v1"; const DEFAULT_STATUS_LINE_ITEMS: [&str; 2] = ["model-with-reasoning", "current-dir"]; -const MAX_AGENT_COPY_HISTORY: usize = 32; /// Common initialization parameters shared by all `ChatWidget` constructors. pub(crate) struct ChatWidgetInit { @@ -965,10 +964,6 @@ impl ChatWidget { } } - fn record_visible_user_turn_for_copy(&mut self) { - self.transcript.record_visible_user_turn(); - } - pub(crate) fn open_feedback_note( &mut self, category: crate::app_event::FeedbackCategory, @@ -1308,7 +1303,6 @@ impl ChatWidget { || !display.local_images.is_empty() || !display.remote_image_urls.is_empty() { - self.record_visible_user_turn_for_copy(); self.add_to_history(history_cell::new_user_prompt( display.message, display.text_elements, diff --git a/codex-rs/tui/src/chatwidget/input_restore.rs b/codex-rs/tui/src/chatwidget/input_restore.rs index 7505f626e2..faa1622b0a 100644 --- a/codex-rs/tui/src/chatwidget/input_restore.rs +++ b/codex-rs/tui/src/chatwidget/input_restore.rs @@ -253,9 +253,25 @@ impl ChatWidget { } pub(crate) fn restore_user_message_to_composer(&mut self, user_message: UserMessage) { + let draft = self.bottom_pane.composer_draft_snapshot(); + let pending_pastes = draft.pending_pastes; + let draft_message = UserMessage { + text: draft.text, + text_elements: draft.text_elements, + local_images: draft.local_images, + remote_image_urls: draft.remote_image_urls, + mention_bindings: draft.mention_bindings, + }; + let mut messages = vec![user_message]; + if !draft_message.text.is_empty() + || !draft_message.local_images.is_empty() + || !draft_message.remote_image_urls.is_empty() + { + messages.push(draft_message); + } self.restore_composer_state(Self::composer_state_from_user_message( - user_message, - Vec::new(), + merge_user_messages(messages), + pending_pastes, )); } diff --git a/codex-rs/tui/src/chatwidget/interaction.rs b/codex-rs/tui/src/chatwidget/interaction.rs index 1d38fbdbe4..c95e821789 100644 --- a/codex-rs/tui/src/chatwidget/interaction.rs +++ b/codex-rs/tui/src/chatwidget/interaction.rs @@ -255,14 +255,6 @@ impl ChatWidget { self.copy_last_agent_markdown_with(crate::clipboard_copy::copy_to_clipboard); } - pub(crate) fn truncate_agent_copy_history_to_user_turn_count( - &mut self, - user_turn_count: usize, - ) { - self.transcript - .truncate_copy_history_to_user_turn_count(user_turn_count); - } - /// Inner implementation with an injectable clipboard backend for testing. pub(super) fn copy_last_agent_markdown_with( &mut self, @@ -281,11 +273,6 @@ impl ChatWidget { "Copy failed: {error}" ))), }, - _ if self.transcript.copy_history_evicted_by_rollback => { - self.add_to_history(history_cell::new_error_event(format!( - "Cannot copy that response after rewinding. Only the most recent {MAX_AGENT_COPY_HISTORY} responses are available to /copy." - ))); - } _ => self.add_to_history(history_cell::new_error_event( "No agent response to copy".into(), )), diff --git a/codex-rs/tui/src/chatwidget/safety_buffering.rs b/codex-rs/tui/src/chatwidget/safety_buffering.rs index 964c2ff796..86eb9b9e07 100644 --- a/codex-rs/tui/src/chatwidget/safety_buffering.rs +++ b/codex-rs/tui/src/chatwidget/safety_buffering.rs @@ -63,21 +63,22 @@ 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 safety_buffering_prompt = self.safety_buffering_prompt.take(); + pub(crate) fn prepare_safety_buffered_retry_submission(&mut self, prompt: UserMessage) { self.last_rendered_user_message_display = None; self.finalize_turn(); - self.safety_buffering_prompt = safety_buffering_prompt; + self.safety_buffering_prompt = Some(prompt); self.input_queue.user_turn_pending_start = true; } - pub(crate) fn fail_safety_buffering_retry(&mut self) { + 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); + } + + pub(crate) fn cancel_safety_buffered_retry_submission(&mut self) { self.input_queue.user_turn_pending_start = false; self.clear_safety_buffering(); - let prompt = self.safety_buffering_prompt.take(); - if let Some(prompt) = prompt { - self.restore_user_message_to_composer(prompt); - } } pub(super) fn on_model_safety_buffering_updated( @@ -112,14 +113,23 @@ impl ChatWidget { return; } - let retry_turn = self - .safety_buffering - .submitted_turn - .as_ref() - .filter(|(submitted_turn_id, _)| replay_kind.is_none() && submitted_turn_id == &turn_id) - .map(|(_, turn)| turn.clone()); + let retry_turn = if self.side_conversation_active() { + None + } else { + self.safety_buffering + .submitted_turn + .as_ref() + .filter(|(submitted_turn_id, _)| { + replay_kind.is_none() && submitted_turn_id == &turn_id + }) + .map(|(_, turn)| turn.clone()) + }; let thread_id = self.thread_id; - let can_offer_retry = faster_model.is_some() && retry_turn.is_some() && thread_id.is_some(); + let retry_prompt = self.safety_buffering_prompt.clone(); + let can_offer_retry = faster_model.is_some() + && retry_turn.is_some() + && retry_prompt.is_some() + && thread_id.is_some(); let previous_active = self .safety_buffering .active @@ -165,8 +175,8 @@ impl ChatWidget { } let header = ColumnRenderable::with(header); let mut items = Vec::new(); - if let (Some(faster_model), Some(turn), Some(thread_id)) = - (faster_model, retry_turn, thread_id) + if let (Some(faster_model), Some(turn), Some(prompt), Some(thread_id)) = + (faster_model, retry_turn, retry_prompt, thread_id) { items.push(SelectionItem { name: "Retry with a faster model".to_string(), @@ -176,6 +186,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/snapshots/codex_tui__chatwidget__tests__safety_buffering_side_conversation_without_retry.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__safety_buffering_side_conversation_without_retry.snap new file mode 100644 index 0000000000..56efb00764 --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__safety_buffering_side_conversation_without_retry.snap @@ -0,0 +1,11 @@ +--- +source: tui/src/chatwidget/tests/app_server.rs +expression: popup +--- + Our systems are thinking a bit more about this request before responding. + +› 1. Dismiss and keep waiting + 2. Learn more + + No action is required. Codex will keep waiting, and this menu will close when + the response is ready. diff --git a/codex-rs/tui/src/chatwidget/tests/app_server.rs b/codex-rs/tui/src/chatwidget/tests/app_server.rs index d273b16314..4db7692ff3 100644 --- a/codex-rs/tui/src/chatwidget/tests/app_server.rs +++ b/codex-rs/tui/src/chatwidget/tests/app_server.rs @@ -144,14 +144,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}"), } @@ -160,12 +161,32 @@ 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_buffering_does_not_offer_retry_in_side_conversation() { + let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.set_side_conversation_active(/*active*/ true); + let (thread_id, turn_id, _) = start_safety_buffering_test_turn(&mut chat, &mut op_rx); + + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(safety_buffering_notification( + thread_id, + turn_id, + Some("faster-model"), + )), + /*replay_kind*/ None, + ); + + let popup = render_bottom_popup(&chat, /*width*/ 80); + assert_chatwidget_snapshot!("safety_buffering_side_conversation_without_retry", popup); +} + #[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; @@ -190,20 +211,6 @@ async fn safety_buffering_remains_visible_until_turn_completes() { assert!(!render_bottom_popup(&chat, /*width*/ 80).contains(SAFETY_BUFFERING_HEADER_TEXT)); } -#[tokio::test] -async fn failed_safety_buffering_retry_restores_submitted_prompt() { - let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - start_safety_buffering_test_turn(&mut chat, &mut op_rx); - - chat.prepare_safety_buffering_retry(); - assert!(chat.input_queue.user_turn_pending_start); - - chat.fail_safety_buffering_retry(); - - assert_eq!(chat.composer_text_with_pending(), "Explain the request"); - assert!(!chat.input_queue.user_turn_pending_start); -} - #[tokio::test] async fn safety_buffering_without_retry_shows_short_app_message() { let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/tests/composer_submission.rs b/codex-rs/tui/src/chatwidget/tests/composer_submission.rs index f858457e4f..8ccf34569d 100644 --- a/codex-rs/tui/src/chatwidget/tests/composer_submission.rs +++ b/codex-rs/tui/src/chatwidget/tests/composer_submission.rs @@ -735,6 +735,52 @@ async fn queued_restore_with_remote_images_keeps_local_placeholder_mapping() { assert_eq!(chat.remote_image_urls(), remote_image_urls); } +#[tokio::test] +async fn restored_message_preserves_existing_composer_draft_and_attachments() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + let retry_image = PathBuf::from("/tmp/retry.png"); + let draft_image = PathBuf::from("/tmp/draft.png"); + let paste_placeholder = "[Pasted Content 5 chars]"; + let draft_text = format!("[Image #1] {paste_placeholder} draft"); + chat.bottom_pane.set_composer_text( + draft_text, + vec![TextElement::new( + (0.."[Image #1]".len()).into(), + Some("[Image #1]".to_string()), + )], + vec![draft_image.clone()], + ); + chat.bottom_pane + .set_composer_pending_pastes(vec![(paste_placeholder.to_string(), "hello".to_string())]); + + chat.restore_user_message_to_composer(UserMessage { + text: "[Image #1] retry prompt".to_string(), + local_images: vec![LocalImageAttachment { + placeholder: "[Image #1]".to_string(), + path: retry_image.clone(), + }], + remote_image_urls: Vec::new(), + text_elements: vec![TextElement::new( + (0.."[Image #1]".len()).into(), + Some("[Image #1]".to_string()), + )], + mention_bindings: Vec::new(), + }); + + assert_eq!( + chat.bottom_pane.composer_text(), + format!("[Image #1] retry prompt\n[Image #2] {paste_placeholder} draft") + ); + assert_eq!( + chat.bottom_pane.composer_local_image_paths(), + vec![retry_image, draft_image] + ); + assert_eq!( + chat.bottom_pane.composer_pending_pastes(), + vec![(paste_placeholder.to_string(), "hello".to_string())] + ); +} + #[tokio::test] async fn interrupted_turn_restore_keeps_active_mode_for_resubmission() { let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(Some("gpt-5")).await; diff --git a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs index ab3848eedb..7d807b8001 100644 --- a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs +++ b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs @@ -2071,62 +2071,6 @@ async fn queued_menu_slash_keeps_agent_turn_complete_notification() { assert_matches!(op_rx.try_recv(), Err(TryRecvError::Empty)); } -#[tokio::test] -async fn slash_copy_uses_latest_surviving_response_after_rollback() { - let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - - replay_user_message_text(&mut chat, "user-1", "foo", ReplayKind::ThreadSnapshot); - replay_agent_message( - &mut chat, - "agent-1", - "foo response", - ReplayKind::ThreadSnapshot, - ); - replay_user_message_text(&mut chat, "user-2", "bar", ReplayKind::ThreadSnapshot); - replay_agent_message( - &mut chat, - "agent-2", - "bar response", - ReplayKind::ThreadSnapshot, - ); - let _ = drain_insert_history(&mut rx); - assert_eq!(chat.last_agent_markdown_text(), Some("bar response")); - - chat.truncate_agent_copy_history_to_user_turn_count(/*user_turn_count*/ 1); - - assert_eq!(chat.last_agent_markdown_text(), Some("foo response")); - chat.copy_last_agent_markdown_with(|markdown| { - assert_eq!(markdown, "foo response"); - Ok(None) - }); -} - -#[tokio::test] -async fn slash_copy_reports_when_rewind_exceeds_retained_copy_history() { - let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - - replay_user_message_text(&mut chat, "user-1", "foo", ReplayKind::ThreadSnapshot); - replay_agent_message( - &mut chat, - "agent-1", - "foo response", - ReplayKind::ThreadSnapshot, - ); - let _ = drain_insert_history(&mut rx); - - chat.truncate_agent_copy_history_to_user_turn_count(/*user_turn_count*/ 0); - chat.dispatch_command(SlashCommand::Copy); - - let cells = drain_insert_history(&mut rx); - let rendered = lines_to_single_string(&cells[0]); - assert!( - rendered.contains( - "Cannot copy that response after rewinding. Only the most recent 32 responses are available to /copy." - ), - "expected evicted-history message, got {rendered:?}" - ); -} - #[tokio::test] async fn slash_exit_requests_exit() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/transcript.rs b/codex-rs/tui/src/chatwidget/transcript.rs index 22b046728b..ed66d8a717 100644 --- a/codex-rs/tui/src/chatwidget/transcript.rs +++ b/codex-rs/tui/src/chatwidget/transcript.rs @@ -1,30 +1,14 @@ //! Transcript and active-cell bookkeeping for `ChatWidget`. use super::HistoryCell; -use super::MAX_AGENT_COPY_HISTORY; - -#[derive(Debug)] -pub(super) struct AgentTurnMarkdown { - pub(super) user_turn_count: usize, - pub(super) markdown: String, -} #[derive(Default)] pub(super) struct TranscriptState { pub(super) active_cell: Option>, /// Monotonic-ish counter used to invalidate transcript overlay caching. pub(super) active_cell_revision: u64, - /// Raw markdown of the most recently completed agent response that - /// survived any local thread rollback. + /// Raw markdown of the most recently completed agent response. pub(super) last_agent_markdown: Option, - /// Copyable agent responses keyed by the number of visible user turns at - /// the time the response completed. - pub(super) agent_turn_markdowns: Vec, - /// Number of user turns currently reflected in the visible transcript. - pub(super) visible_user_turn_count: usize, - /// True when rollback discarded the requested copy source because it was - /// older than the retained copy history. - pub(super) copy_history_evicted_by_rollback: bool, /// Raw markdown of the most recently completed proposed plan. pub(super) latest_proposed_plan_markdown: Option, /// Whether this turn already produced a copyable response. @@ -61,48 +45,12 @@ impl TranscriptState { } pub(super) fn record_agent_markdown(&mut self, markdown: String) { - match self.agent_turn_markdowns.last_mut() { - Some(entry) if entry.user_turn_count == self.visible_user_turn_count => { - entry.markdown = markdown.clone(); - } - _ => { - self.agent_turn_markdowns.push(AgentTurnMarkdown { - user_turn_count: self.visible_user_turn_count, - markdown: markdown.clone(), - }); - if self.agent_turn_markdowns.len() > MAX_AGENT_COPY_HISTORY { - self.agent_turn_markdowns.remove(0); - } - } - } self.last_agent_markdown = Some(markdown); - self.copy_history_evicted_by_rollback = false; self.saw_copy_source_this_turn = true; } - pub(super) fn record_visible_user_turn(&mut self) { - self.visible_user_turn_count = self.visible_user_turn_count.saturating_add(1); - } - pub(super) fn reset_copy_history(&mut self) { self.last_agent_markdown = None; - self.agent_turn_markdowns.clear(); - self.visible_user_turn_count = 0; - self.copy_history_evicted_by_rollback = false; - self.saw_copy_source_this_turn = false; - } - - pub(super) fn truncate_copy_history_to_user_turn_count(&mut self, user_turn_count: usize) { - self.visible_user_turn_count = user_turn_count; - let had_copy_history = !self.agent_turn_markdowns.is_empty(); - self.agent_turn_markdowns - .retain(|entry| entry.user_turn_count <= user_turn_count); - self.last_agent_markdown = self - .agent_turn_markdowns - .last() - .map(|entry| entry.markdown.clone()); - self.copy_history_evicted_by_rollback = - had_copy_history && self.last_agent_markdown.is_none(); self.saw_copy_source_this_turn = false; } @@ -134,18 +82,4 @@ mod tests { assert_eq!(state.active_cell_revision, 0); } - - #[test] - fn copy_history_tracks_latest_visible_turn() { - let mut state = TranscriptState::default(); - state.record_visible_user_turn(); - state.record_agent_markdown("first".to_string()); - state.record_visible_user_turn(); - state.record_agent_markdown("second".to_string()); - - state.truncate_copy_history_to_user_turn_count(/*user_turn_count*/ 1); - - assert_eq!(state.last_agent_markdown.as_deref(), Some("first")); - assert!(!state.copy_history_evicted_by_rollback); - } } diff --git a/codex-rs/tui/src/pager_overlay.rs b/codex-rs/tui/src/pager_overlay.rs index 9c5935fb4a..640438557b 100644 --- a/codex-rs/tui/src/pager_overlay.rs +++ b/codex-rs/tui/src/pager_overlay.rs @@ -803,11 +803,6 @@ impl TranscriptOverlay { pub(crate) fn is_done(&self) -> bool { self.is_done } - - #[cfg(test)] - pub(crate) fn committed_cell_count(&self) -> usize { - self.cells.len() - } } pub(crate) struct StaticOverlay {